Add Cheatcode Syscalls (#588)

* add scaffolding

* add placeholder for callback functions

* add testing_state to syscall handler example

* add add parameters to callbacks

* add set_version

* add wrap_set_account_contract_address

* add other callbacks

* add testing methods to syscall handler

* add set_sequencer_address cheatcode

* implement cheatcodes for testing syscallhandler

* make tests runnable

* fix callbacks

* fix warnings

* add missing constants

* add some debug prints and cheatcode method

* change cheatcode return value to arrayAbi

* remove testing syscalls in favour of having only cheatcode syscall

* remove testing syscalls from example syscall handler

* small fixes

* Remove stale impls

* Remove old code approach

* allow dead_code

* Remove debug

* Rename parameter

* Add dummy implementation

* Communicate with rust runtime

* Add selector to cheatcode signature

* Add remaining arguments to cheatcode libfunc

* Add comments

* Fix selector ptr bug

* Remove unused

* Remove dbg!

* Set vtable pointer

* Improve legibility

* Remove dbg!()

* Call cheatcode from runtime

* Add default return if no system handler is set

* Make tests pass

* Use references to never drop values

* Remove unused

* Alloc syscall handler always

testing syscalls don't needed it acording to sierra, so we must allocate
memory for it always

* Fix clippy warning

* Simplify pointer logic

* Remvoe dbg!

* Add dummy implementation of cheatcode syscall

* Fix memory leak

syscall_handler was a &mut &mut, and needed to be a &mut

* Add documentation

* Delete cheatcode file

* Add tests for sequencer address

* Add comment

* Remove comment (already implemented)

* Add all cheatcode syscalls

* Implement all other syscalls in testing

* Add testing cheatcode implementation to starknet example

* Add todo! implementation of cheatcode

* Remove warnings

* Make self_ptr private

* Add pop_l2_to_l1_messages

* Remove extra space

* Implement cheatcode with unimplemented by default

* User assert instead of return

* Restore syscall handler and don't use arena

* Add comment

* Add another comment

* Added cheatcode feature to Cargo.toml

* Add feature flags

* Add documentation

* Fix clippy

---------

Co-authored-by: Julián González Calderón <[email protected]>

Cargo.toml

@@ -51,6 +51,7 @@ scarb = ["build-cli", "dep:scarb-ui", "dep:scarb-metadata", "dep:serde_json"]
 with-debug-utils = []
 with-runtime = ["dep:cairo-native-runtime"]
 with-serde = ["dep:serde"]
+with-cheatcode = []
 
 [dependencies]
 bumpalo = "3.16.0"

examples/starknet.rs

@@ -10,11 +10,50 @@ use cairo_native::{
     utils::find_entry_point_by_idx,
 };
 use starknet_types_core::felt::Felt;
-use std::path::Path;
+use std::{
+    collections::{HashMap, VecDeque},
+    path::Path,
+};
 use tracing_subscriber::{EnvFilter, FmtSubscriber};
 
-#[derive(Debug)]
-struct SyscallHandler;
+type Log = (Vec<Felt>, Vec<Felt>);
+type L2ToL1Message = (Felt, Vec<Felt>);
+
+#[derive(Debug, Default)]
+struct ContractLogs {
+    events: VecDeque<Log>,
+    l2_to_l1_messages: VecDeque<L2ToL1Message>,
+}
+
+#[derive(Debug, Default)]
+struct TestingState {
+    sequencer_address: Felt,
+    caller_address: Felt,
+    contract_address: Felt,
+    account_contract_address: Felt,
+    transaction_hash: Felt,
+    nonce: Felt,
+    chain_id: Felt,
+    version: Felt,
+    max_fee: u64,
+    block_number: u64,
+    block_timestamp: u64,
+    signature: Vec<Felt>,
+    logs: HashMap<Felt, ContractLogs>,
+}
+
+#[derive(Debug, Default)]
+struct SyscallHandler {
+    testing_state: TestingState,
+}
+
+impl SyscallHandler {
+    pub fn new() -> Self {
+        Self {
+            testing_state: TestingState::default(),
+        }
+    }
+}
 
 impl StarknetSyscallHandler for SyscallHandler {
     fn get_block_hash(&mut self, block_number: u64, _gas: &mut u128) -> SyscallResult<Felt> {
@@ -265,6 +304,98 @@ impl StarknetSyscallHandler for SyscallHandler {
     ) -> SyscallResult<(U256, U256)> {
         unimplemented!()
     }
+
+    #[cfg(feature = "with-cheatcode")]
+    fn cheatcode(&mut self, selector: Felt, input: &[Felt]) -> Vec<Felt> {
+        let selector_bytes = selector.to_bytes_be();
+
+        let selector = match std::str::from_utf8(&selector_bytes) {
+            Ok(selector) => selector.trim_start_matches('\0'),
+            Err(_) => return Vec::new(),
+        };
+
+        match selector {
+            "set_sequencer_address" => {
+                self.testing_state.sequencer_address = input[0];
+                vec![]
+            }
+            "set_caller_address" => {
+                self.testing_state.caller_address = input[0];
+                vec![]
+            }
+            "set_contract_address" => {
+                self.testing_state.contract_address = input[0];
+                vec![]
+            }
+            "set_account_contract_address" => {
+                self.testing_state.account_contract_address = input[0];
+                vec![]
+            }
+            "set_transaction_hash" => {
+                self.testing_state.transaction_hash = input[0];
+                vec![]
+            }
+            "set_nonce" => {
+                self.testing_state.nonce = input[0];
+                vec![]
+            }
+            "set_version" => {
+                self.testing_state.version = input[0];
+                vec![]
+            }
+            "set_chain_id" => {
+                self.testing_state.chain_id = input[0];
+                vec![]
+            }
+            "set_max_fee" => {
+                let max_fee = input[0].to_biguint().try_into().unwrap();
+                self.testing_state.max_fee = max_fee;
+                vec![]
+            }
+            "set_block_number" => {
+                let block_number = input[0].to_biguint().try_into().unwrap();
+                self.testing_state.block_number = block_number;
+                vec![]
+            }
+            "set_block_timestamp" => {
+                let block_timestamp = input[0].to_biguint().try_into().unwrap();
+                self.testing_state.block_timestamp = block_timestamp;
+                vec![]
+            }
+            "set_signature" => {
+                self.testing_state.signature = input.to_vec();
+                vec![]
+            }
+            "pop_log" => self
+                .testing_state
+                .logs
+                .get_mut(&input[0])
+                .and_then(|logs| logs.events.pop_front())
+                .map(|mut log| {
+                    let mut serialized_log = Vec::new();
+                    serialized_log.push(log.0.len().into());
+                    serialized_log.append(&mut log.0);
+                    serialized_log.push(log.1.len().into());
+                    serialized_log.append(&mut log.1);
+                    serialized_log
+                })
+                .unwrap_or_default(),
+            "pop_l2_to_l1_message" => self
+                .testing_state
+                .logs
+                .get_mut(&input[0])
+                .and_then(|logs| logs.l2_to_l1_messages.pop_front())
+                .map(|mut log| {
+                    let mut serialized_log = Vec::new();
+                    serialized_log.push(log.0);
+                    serialized_log.push(log.1.len().into());
+                    serialized_log.append(&mut log.1);
+                    serialized_log
+                })
+                .unwrap_or_default(),
+            _ => vec![],
+        }
+    }
 }
 
 fn main() {
@@ -305,7 +436,7 @@ fn main() {
     let native_executor = JitNativeExecutor::from_native_module(native_program, Default::default());
 
     let result = native_executor
-        .invoke_contract_dynamic(fn_id, &[Felt::ONE], Some(u128::MAX), SyscallHandler)
+        .invoke_contract_dynamic(fn_id, &[Felt::ONE], Some(u128::MAX), SyscallHandler::new())
         .expect("failed to execute the given contract");
 
     println!();

src/executor.rs

@@ -151,7 +151,6 @@ fn invoke_dynamic(
     mut syscall_handler: Option<impl StarknetSyscallHandler>,
 ) -> Result<ExecutionResult, Error> {
     tracing::info!("Invoking function with signature: {function_signature:?}.");
-
     let arena = Bump::new();
     let mut invoke_data = ArgumentMapper::new(&arena, registry);
 
@@ -196,6 +195,22 @@ fn invoke_dynamic(
         None
     };
 
+    // The Cairo compiler doesn't specify that the cheatcode syscall needs the syscall handler,
+    // so we must always allocate it in case it needs it, regardless of whether it's passed
+    // as an argument to the entry point or not.
+    let mut syscall_handler = syscall_handler
+        .as_mut()
+        .map(|syscall_handler| StarknetSyscallHandlerCallbacks::new(syscall_handler));
+    // We only care for the previous syscall handler if we actually modify it
+    #[cfg(feature = "with-cheatcode")]
+    let previous_syscall_handler = syscall_handler.as_mut().map(|syscall_handler| {
+        let previous_syscall_handler = crate::starknet::SYSCALL_HANDLER_VTABLE.get();
+        let syscall_handler_ptr = std::ptr::addr_of!(*syscall_handler) as *mut ();
+        crate::starknet::SYSCALL_HANDLER_VTABLE.set(syscall_handler_ptr);
+
+        previous_syscall_handler
+    });
+
     // Generate argument list.
     let mut iter = args.iter();
     for type_id in function_signature.param_types.iter().filter(|id| {
@@ -209,19 +224,14 @@ fn invoke_dynamic(
                 &[gas as u64, (gas >> 64) as u64],
             ),
             CoreTypeConcrete::StarkNet(StarkNetTypeConcrete::System(_)) => {
-                match syscall_handler.as_mut() {
-                    Some(syscall_handler) => {
-                        let syscall_handler =
-                            arena.alloc(StarknetSyscallHandlerCallbacks::new(syscall_handler));
-                        invoke_data.push_aligned(
-                            get_integer_layout(64).align(),
-                            &[syscall_handler as *mut _ as u64],
-                        )
-                    }
-                    None => {
-                        panic!("Syscall handler is required");
-                    }
-                }
+                let syscall_handler = syscall_handler
+                    .as_mut()
+                    .expect("syscall handler is required");
+
+                invoke_data.push_aligned(
+                    get_integer_layout(64).align(),
+                    &[syscall_handler as *mut _ as u64],
+                );
             }
             type_info => invoke_data
                 .push(
@@ -252,6 +262,13 @@ fn invoke_dynamic(
         );
     }
 
+    // If the syscall handler was changed, then reset the previous one.
+    // It's only necessary to restore the pointer if it's been modified i.e. if previous_syscall_handler is Some(...)
+    #[cfg(feature = "with-cheatcode")]
+    if let Some(previous_syscall_handler) = previous_syscall_handler {
+        crate::starknet::SYSCALL_HANDLER_VTABLE.set(previous_syscall_handler);
+    }
+
     // Parse final gas.
     unsafe fn read_value<T>(ptr: &mut NonNull<()>) -> &T {
         let align_offset = ptr

src/libfuncs/starknet.rs

@@ -15,7 +15,7 @@ use cairo_lang_sierra::{
         consts::SignatureAndConstConcreteLibfunc,
         core::{CoreLibfunc, CoreType},
         lib_func::SignatureOnlyConcreteLibfunc,
-        starknet::StarkNetConcreteLibfunc,
+        starknet::{testing::TestingConcreteLibfunc, StarkNetConcreteLibfunc},
         ConcreteLibfunc,
     },
     program_registry::ProgramRegistry,
@@ -40,6 +40,9 @@ use std::alloc::Layout;
 
 mod secp256;
 
+#[cfg(feature = "with-cheatcode")]
+mod testing;
+
 /// Select and call the correct libfunc builder function from the selector.
 pub fn build<'ctx, 'this>(
     context: &'ctx Context,
@@ -138,7 +141,14 @@ pub fn build<'ctx, 'this>(
         StarkNetConcreteLibfunc::Secp256(selector) => self::secp256::build(
             context, registry, entry, location, helper, metadata, selector,
         ),
-        StarkNetConcreteLibfunc::Testing(_) => todo!("implement starknet testing libfunc"),
+        #[cfg(feature = "with-cheatcode")]
+        StarkNetConcreteLibfunc::Testing(TestingConcreteLibfunc::Cheatcode(info)) => {
+            self::testing::build(context, registry, entry, location, helper, metadata, info)
+        }
+        #[cfg(not(feature = "with-cheatcode"))]
+        StarkNetConcreteLibfunc::Testing(TestingConcreteLibfunc::Cheatcode(_)) => {
+            unimplemented!("feature 'with-cheatcode' is required to compile with cheatcode syscall")
+        }
     }
 }