From 59e7f7402e0c706a7416e4fab0430a708b8d2d87 Mon Sep 17 00:00:00 2001
From: Maksim Zdobnikau <maksim.zdobnikau@swmansion.com>
Date: Thu, 19 Dec 2024 20:27:42 +0100
Subject: [PATCH] add size limit checks for Sierra and CASM programs

---
 .../compilers/starknet_contract/compiler.rs   |  6 ++-
 .../starknet_contract/validations.rs          | 51 +++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/scarb/src/compiler/compilers/starknet_contract/compiler.rs b/scarb/src/compiler/compilers/starknet_contract/compiler.rs
index 4091226c6..83eed5862 100644
--- a/scarb/src/compiler/compilers/starknet_contract/compiler.rs
+++ b/scarb/src/compiler/compilers/starknet_contract/compiler.rs
@@ -17,7 +17,7 @@ use tracing::{debug, trace, trace_span};
 
 use super::contract_selector::ContractSelector;
 use crate::compiler::compilers::starknet_contract::contract_selector::GLOB_PATH_SELECTOR;
-use crate::compiler::compilers::starknet_contract::validations::check_allowed_libfuncs;
+use crate::compiler::compilers::starknet_contract::validations::{check_allowed_libfuncs, check_sierra_size_limits, check_casm_size_limits};
 use crate::compiler::compilers::{ensure_gas_enabled, ArtifactsWriter};
 use crate::compiler::helpers::{build_compiler_config, collect_main_crate_ids};
 use crate::compiler::{CairoCompilationUnit, CompilationUnitAttributes, Compiler};
@@ -113,6 +113,8 @@ impl Compiler for StarknetContractCompiler {
 
         check_allowed_libfuncs(&props, &contracts, &classes, db, &unit, ws)?;
 
+        check_sierra_size_limits(&classes, ws);
+
         let casm_classes: Vec<Option<CasmContractClass>> = if props.casm {
             let span = trace_span!("compile_sierra");
             let _guard = span.enter();
@@ -135,6 +137,8 @@ impl Compiler for StarknetContractCompiler {
             classes.iter().map(|_| None).collect()
         };
 
+        check_casm_size_limits(&casm_classes, ws);
+
         let target_name = &unit.main_component().target_name();
 
         let writer = ArtifactsWriter::new(target_name.clone(), target_dir, props);
diff --git a/scarb/src/compiler/compilers/starknet_contract/validations.rs b/scarb/src/compiler/compilers/starknet_contract/validations.rs
index 7daa93d03..cc23fa8a6 100644
--- a/scarb/src/compiler/compilers/starknet_contract/validations.rs
+++ b/scarb/src/compiler/compilers/starknet_contract/validations.rs
@@ -12,6 +12,7 @@ use cairo_lang_starknet_classes::allowed_libfuncs::{
     AllowedLibfuncsError, ListSelector, BUILTIN_EXPERIMENTAL_LIBFUNCS_LIST,
 };
 use cairo_lang_starknet_classes::contract_class::ContractClass;
+use cairo_lang_starknet_classes::casm_contract_class::CasmContractClass;
 use cairo_lang_utils::Upcast;
 use indoc::{formatdoc, writedoc};
 use std::fmt::Write;
@@ -19,6 +20,10 @@ use std::iter::zip;
 use tracing::debug;
 
 const AUTO_WITHDRAW_GAS_FLAG: &str = "add_withdraw_gas";
+const MAX_SIERRA_PROGRAM_FELTS: usize = 81290;
+const MAX_CASM_PROGRAM_FELTS: usize = 81290;
+const MAX_CONTRACT_CLASS_BYTES: usize = 4089446;
+const MAX_COMPILED_CONTRACT_CLASS_BYTES: usize = 4089446;
 
 pub fn ensure_gas_enabled(db: &mut RootDatabase) -> anyhow::Result<()> {
     let flag = FlagId::new(db.as_files_group_mut(), AUTO_WITHDRAW_GAS_FLAG);
@@ -117,3 +122,49 @@ pub fn check_allowed_libfuncs(
 
     Ok(())
 }
+
+pub fn check_sierra_size_limits(
+    classes: &[ContractClass],
+    ws: &Workspace<'_>,
+) {
+    for class in classes {
+        let sierra_felts = class.sierra_program.to_vec().len();
+        if sierra_felts > MAX_SIERRA_PROGRAM_FELTS {
+            ws.config().ui().warn(formatdoc! {r#"
+                Sierra program exceeds maximum byte-code size:
+                {MAX_SIERRA_PROGRAM_FELTS} felts allowed on Starknet. Actual size: {sierra_felts} felts.
+            "#});
+        }
+
+        let class_size = serde_json::to_vec(class).unwrap().len();
+        if class_size > MAX_CONTRACT_CLASS_BYTES {
+            ws.config().ui().warn(formatdoc! {r#"
+                Contract class size exceeds maximum allowed size:
+                {MAX_CONTRACT_CLASS_BYTES} bytes allowed on Starknet. Actual size: {class_size} bytes.
+            "#});
+        }
+    }
+}
+
+pub fn check_casm_size_limits(
+    casm_classes: &[Option<CasmContractClass>],
+    ws: &Workspace<'_>,
+) {
+    for casm_class in casm_classes.iter().flatten() {
+        let casm_felts = casm_class.bytecode.len();
+        if casm_felts > MAX_CASM_PROGRAM_FELTS {
+            ws.config().ui().warn(formatdoc! {r#"
+                CASM program exceeds maximum byte-code size:
+                {MAX_CASM_PROGRAM_FELTS} felts allowed on Starknet. Actual size: {casm_felts} felts.
+            "#});
+        }
+
+        let compiled_class_size = serde_json::to_vec(casm_class).unwrap().len();
+        if compiled_class_size > MAX_COMPILED_CONTRACT_CLASS_BYTES {
+            ws.config().ui().warn(formatdoc! {r#"
+                Compiled contract class size exceeds maximum allowed size:
+                {MAX_COMPILED_CONTRACT_CLASS_BYTES} bytes allowed on Starknet. Actual size: {compiled_class_size} bytes.
+            "#});
+        }
+    }
+}