Modify program

src/bin/cairo-native-stress.rs

@@ -10,7 +10,7 @@ use std::path::Path;
 use std::{collections::HashMap, fs, rc::Rc, time::Instant};
 
 use cairo_lang_sierra::ids::FunctionId;
-use cairo_lang_sierra::program::Program;
+use cairo_lang_sierra::program::{GenericArg, Program};
 use cairo_lang_sierra::program_registry::ProgramRegistry;
 use cairo_lang_starknet::compile::compile_path;
 use cairo_native::metadata::gas::GasMetadata;
@@ -22,12 +22,19 @@ use cairo_native::{
 use cairo_native::{module_to_object, object_to_shared_lib, OptLevel};
 use clap::Parser;
 use libloading::Library;
+use num_bigint::BigInt;
 use tracing::{debug, debug_span, info, info_span, warn};
 use tracing_subscriber::{EnvFilter, FmtSubscriber};
 
 /// The directory used to store compiled native programs
 const AOT_CACHE_DIR: &str = ".aot-cache";
 
+/// An unique value hardcoded into the first contract that it's
+/// used as an anchor point to safely modify the return value on
+/// the following contracts.
+/// It can be any value as long as it's unique in the initial contract.
+const RETURN_ANCHOR: u32 = 835;
+
 /// A stress tester for Cairo Native
 ///
 /// It Sierra programs compiles with Cairo Native, caches, and executes them with AOT runner.
@@ -75,9 +82,11 @@ fn main() {
         for program_number in 0..cli_args.programs {
             let _enter_program_span = debug_span!("program");
 
-            // The round and program count is used as a key. After making sure each iteration uses
-            // a different unique program, the program hash should be used.
-            let key = round * cli_args.programs + program_number;
+            let global_program_number = round * cli_args.programs + program_number;
+            let program = modify_starknet_contract(program.clone(), global_program_number);
+
+            // The round and program count is used as a key. Later on the program hash should be used.
+            let key = global_program_number;
             if cache.get(&key).is_some() {
                 panic!("all program keys should be different")
             }
@@ -104,7 +113,8 @@ fn main() {
                     )
                     .expect("failed to execute contract");
                 let elapsed = now.elapsed().as_millis();
-                debug!(time = elapsed, "executed test program");
+                let result = execution_result.return_values[0];
+                debug!(time = elapsed, result = %result, "executed test program");
                 execution_result
             };
 
@@ -138,18 +148,20 @@ fn main() {
 /// We should modify the program returned from this to obtain
 /// different unique programs without recompiling each time
 fn generate_starknet_contract() -> (FunctionId, cairo_lang_sierra::program::Program) {
-    let program_str = "\
+    let program_str = format!(
+        "\
 #[starknet::contract]
-mod Contract {
+mod Contract {{
     #[storage]
-    struct Storage {}
+    struct Storage {{}}
 
     #[external(v0)]
-    fn main(self: @ContractState) -> felt252 {
-        return 252;
-    }
-}
-";
+    fn main(self: @ContractState) -> felt252 {{
+        return {RETURN_ANCHOR};
+    }}
+}}
+"
+    );
 
     let mut program_file = tempfile::Builder::new()
         .prefix("test_")
@@ -180,6 +192,34 @@ mod Contract {
     (entry_point, program)
 }
 
+/// Modifies the given contract by replacing the `RETURN_ANCHOR` wtith `new_return_value`
+///
+/// The contract must only contain the value `RETURN_ANCHOR` exactly once
+fn modify_starknet_contract(mut program: Program, new_return_value: u32) -> Program {
+    let mut anchor_counter = 0;
+
+    for type_declaration in &mut program.type_declarations {
+        for generic_arg in &mut type_declaration.long_id.generic_args {
+            let anchor = BigInt::from(RETURN_ANCHOR);
+
+            match generic_arg {
+                GenericArg::Value(return_value) if *return_value == anchor => {
+                    *return_value = BigInt::from(new_return_value);
+                    anchor_counter += 1;
+                }
+                _ => {}
+            };
+        }
+    }
+
+    assert!(
+        anchor_counter == 1,
+        "RETURN_ANCHOR was not found exactly once"
+    );
+
+    program
+}
+
 /// A naive implementation of an AOT Program Cache.
 ///
 /// Stores `AotNativeExecutor`s by a given key. Each executors has it's corresponding