Cheatcode function

corelib/src/starknet/testing.cairo

@@ -1,9 +1,11 @@
 use starknet::ContractAddress;
+use array::ArrayTrait;
+use array::SpanTrait;
+use traits::Into;
 
 extern fn set_caller_address(address: ContractAddress) implicits() nopanic;
 extern fn set_contract_address(address: ContractAddress) implicits() nopanic;
 extern fn set_sequencer_address(address: ContractAddress) implicits() nopanic;
-extern fn set_block_number(block_number: u64) implicits() nopanic;
 extern fn set_block_timestamp(block_timestamp: u64) implicits() nopanic;
 extern fn set_version(version: felt252) implicits() nopanic;
 extern fn set_account_contract_address(address: ContractAddress) implicits() nopanic;
@@ -15,3 +17,18 @@ extern fn set_signature(signature: Span<felt252>) implicits() nopanic;
 extern fn pop_log(
     address: ContractAddress
 ) -> Option<(Span<felt252>, Span<felt252>)> implicits() nopanic;
+
+// A general cheatcode function used to simplify implementation of Starknet testing functions.
+// External users of the cairo crates can also implement their own cheatcodes
+// by injecting custom `CairoHintProcessor`.
+extern fn cheatcode<const selector: felt252>(
+    input: Span<felt252>
+) -> Span<felt252> implicits() nopanic;
+
+// Set the block number to the provided value.
+fn set_block_number(block_number: u64) {
+    let mut data = ArrayTrait::new();
+    data.append(block_number.into());
+
+    cheatcode::<'set_block_number'>(data.span());
+}

crates/cairo-lang-casm/src/hints/mod.rs

@@ -1,5 +1,6 @@
 use std::fmt::{Display, Formatter};
 
+use cairo_lang_utils::bigint::BigIntAsHex;
 use indoc::writedoc;
 use parity_scale_codec_derive::{Decode, Encode};
 use schemars::JsonSchema;
@@ -48,30 +49,28 @@ pub enum StarknetHint {
     #[codec(index = 0)]
     SystemCall { system: ResOperand },
     #[codec(index = 1)]
-    SetBlockNumber { value: ResOperand },
-    #[codec(index = 2)]
     SetBlockTimestamp { value: ResOperand },
-    #[codec(index = 3)]
+    #[codec(index = 2)]
     SetCallerAddress { value: ResOperand },
-    #[codec(index = 4)]
+    #[codec(index = 3)]
     SetContractAddress { value: ResOperand },
-    #[codec(index = 5)]
+    #[codec(index = 4)]
     SetSequencerAddress { value: ResOperand },
-    #[codec(index = 6)]
+    #[codec(index = 5)]
     SetVersion { value: ResOperand },
-    #[codec(index = 7)]
+    #[codec(index = 6)]
     SetAccountContractAddress { value: ResOperand },
-    #[codec(index = 8)]
+    #[codec(index = 7)]
     SetMaxFee { value: ResOperand },
-    #[codec(index = 9)]
+    #[codec(index = 8)]
     SetTransactionHash { value: ResOperand },
-    #[codec(index = 10)]
+    #[codec(index = 9)]
     SetChainId { value: ResOperand },
-    #[codec(index = 11)]
+    #[codec(index = 10)]
     SetNonce { value: ResOperand },
-    #[codec(index = 12)]
+    #[codec(index = 11)]
     SetSignature { start: ResOperand, end: ResOperand },
-    #[codec(index = 13)]
+    #[codec(index = 12)]
     PopLog {
         value: ResOperand,
         opt_variant: CellRef,
@@ -80,6 +79,14 @@ pub enum StarknetHint {
         data_start: CellRef,
         data_end: CellRef,
     },
+    #[codec(index = 13)]
+    Cheatcode {
+        selector: BigIntAsHex,
+        input_start: ResOperand,
+        input_end: ResOperand,
+        output_start: CellRef,
+        output_end: CellRef,
+    },
 }
 
 // Represents a cairo core hint.
@@ -699,9 +706,6 @@ impl Display for StarknetHint {
             StarknetHint::SystemCall { system } => {
                 write!(f, "syscall_handler.syscall(syscall_ptr={})", ResOperandFormatter(system))
             }
-            StarknetHint::SetBlockNumber { value } => {
-                write!(f, "syscall_handler.block_number = {}", ResOperandFormatter(value))
-            }
             StarknetHint::SetBlockTimestamp { value } => {
                 write!(f, "syscall_handler.block_timestamp = {}", ResOperandFormatter(value))
             }
@@ -748,16 +752,12 @@ impl Display for StarknetHint {
                     ResOperandFormatter(end)
                 )
             }
-            StarknetHint::PopLog {
-                value: _,
-                opt_variant: _,
-                keys_start: _,
-                keys_end: _,
-                data_start: _,
-                data_end: _,
-            } => {
+            StarknetHint::PopLog { .. } => {
                 write!(f, "raise NotImplemented")
             }
+            StarknetHint::Cheatcode { .. } => {
+                write!(f, "raise NotImplementedError")
+            }
         }
     }
 }

crates/cairo-lang-runner/src/casm_run/mod.rs

@@ -324,9 +324,6 @@ impl HintProcessor for CairoHintProcessor<'_> {
             StarknetHint::SystemCall { system } => {
                 self.execute_syscall(system, vm, exec_scopes)?;
             }
-            StarknetHint::SetBlockNumber { value } => {
-                self.starknet_state.exec_info.block_info.block_number = get_val(vm, value)?;
-            }
             StarknetHint::SetSequencerAddress { value } => {
                 self.starknet_state.exec_info.block_info.sequencer_address = get_val(vm, value)?;
             }
@@ -395,6 +392,33 @@ impl HintProcessor for CairoHintProcessor<'_> {
                     insert_value_to_cellref!(vm, opt_variant, 1)?;
                 }
             }
+            StarknetHint::Cheatcode { selector, input_start, input_end, .. } => {
+                let selector = &selector.value.to_bytes_be().1;
+                let selector = std::str::from_utf8(selector).map_err(|_| {
+                    HintError::CustomHint(Box::from("failed to parse selector".to_string()))
+                })?;
+
+                let input_start = extract_relocatable(vm, input_start)?;
+                let input_end = extract_relocatable(vm, input_end)?;
+                let inputs = vm_get_range(vm, input_start, input_end)?;
+
+                match selector {
+                    "set_block_number" => match &inputs[..] {
+                        [input] => {
+                            self.starknet_state.exec_info.block_info.block_number = input.clone();
+                        }
+                        _ => {
+                            return Err(HintError::CustomHint(Box::from(
+                                "set_block_number cheatcode invalid args: pass span of an array \
+                                 with exactly one element",
+                            )));
+                        }
+                    },
+                    _ => Err(HintError::CustomHint(Box::from(format!(
+                        "Unknown cheatcode selector: {selector}"
+                    ))))?,
+                }
+            }
         };
         Ok(())
     }

crates/cairo-lang-sierra-ap-change/src/core_libfunc_ap_change.rs

@@ -258,6 +258,7 @@ pub fn core_libfunc_ap_change<InfoProvider: InvocationApChangeInfoProvider>(
             }
             StarkNetConcreteLibfunc::Testing(libfunc) => match libfunc {
                 TestingConcreteLibfunc::PopLog(_) => vec![ApChange::Known(5), ApChange::Known(5)],
+                TestingConcreteLibfunc::Cheatcode(_) => vec![ApChange::Known(2)],
                 _ => vec![ApChange::Known(0)],
             },
         },

crates/cairo-lang-sierra-to-casm/src/invocations/starknet/testing.rs

@@ -2,6 +2,7 @@ use cairo_lang_casm::builder::{CasmBuilder, Var};
 use cairo_lang_casm::casm_build_extend;
 use cairo_lang_casm::hints::StarknetHint;
 use cairo_lang_sierra::extensions::starknet::testing::TestingConcreteLibfunc;
+use cairo_lang_utils::bigint::BigIntAsHex;
 
 use crate::invocations::{
     add_input_variables, get_non_fallthrough_statement_id, CompiledInvocation,
@@ -20,10 +21,6 @@ pub fn build(
         Ok(value)
     };
     match libfunc {
-        TestingConcreteLibfunc::SetBlockNumber(_) => {
-            let value = declare_single_value()?;
-            casm_build_extend! {casm_builder, hint StarknetHint::SetBlockNumber {value: value}; };
-        }
         TestingConcreteLibfunc::SetBlockTimestamp(_) => {
             let value = declare_single_value()?;
             casm_build_extend! {casm_builder, hint StarknetHint::SetBlockTimestamp {value: value}; };
@@ -102,6 +99,40 @@ pub fn build(
                 CostValidationInfo::default(),
             ));
         }
+        TestingConcreteLibfunc::Cheatcode(c) => {
+            let [input] = builder.try_get_refs()?;
+            let [input_start, input_end] = input.try_unpack()?;
+
+            add_input_variables! {casm_builder,
+                deref input_start;
+                deref input_end;
+            }
+
+            casm_build_extend! {casm_builder,
+                tempvar output_start;
+                tempvar output_end;
+            }
+
+            casm_builder.add_hint(
+                |[input_start, input_end], [output_start, output_end]| StarknetHint::Cheatcode {
+                    selector: BigIntAsHex { value: c.selector.clone() },
+                    input_start,
+                    input_end,
+                    output_start,
+                    output_end,
+                },
+                [input_start, input_end],
+                [output_start, output_end],
+            );
+
+            casm_build_extend! {casm_builder, ap += 2; }
+
+            return Ok(builder.build_from_casm_builder(
+                casm_builder,
+                [("Fallthrough", &[&[output_start, output_end]], None)],
+                CostValidationInfo::default(),
+            ));
+        }
     }
     casm_build_extend! {casm_builder, ap += 0; };
 

crates/cairo-lang-sierra/src/extensions/modules/starknet/testing.rs

@@ -1,5 +1,7 @@
 use std::marker::PhantomData;
 
+use num_bigint::BigInt;
+
 use super::felt252_span_ty;
 use super::interoperability::ContractAddressType;
 use crate::define_libfunc_hierarchy;
@@ -8,13 +10,15 @@ use crate::extensions::int::unsigned::Uint64Type;
 use crate::extensions::int::unsigned128::Uint128Type;
 use crate::extensions::lib_func::{
     BranchSignature, DeferredOutputKind, LibfuncSignature, OutputVarInfo, ParamSignature,
-    SierraApChange, SignatureSpecializationContext,
+    SierraApChange, SignatureSpecializationContext, SpecializationContext,
 };
 use crate::extensions::{
-    NamedType, NoGenericArgsGenericLibfunc, NoGenericArgsGenericType, OutputVarReferenceInfo,
-    SpecializationError,
+    NamedLibfunc, NamedType, NoGenericArgsGenericLibfunc, NoGenericArgsGenericType,
+    OutputVarReferenceInfo, SignatureBasedConcreteLibfunc, SpecializationError,
 };
 use crate::ids::ConcreteTypeId;
+use crate::program::GenericArg;
+
 /// Trait for implementing test setters.
 pub trait TestSetterTraits: Default {
     /// The generic libfunc id for the setter libfunc.
@@ -64,13 +68,6 @@ impl<TTestSetterTraits: TestSetterTraits> NoGenericArgsGenericLibfunc
     }
 }
 
-#[derive(Default)]
-pub struct SetBlockNumberTrait {}
-impl BasicTypeTestSetterTraits for SetBlockNumberTrait {
-    const STR_ID: &'static str = "set_block_number";
-    type ValueType = Uint64Type;
-}
-
 #[derive(Default)]
 pub struct SetBlockTimestampTrait {}
 impl BasicTypeTestSetterTraits for SetBlockTimestampTrait {
@@ -196,9 +193,74 @@ impl NoGenericArgsGenericLibfunc for PopLogLibfunc {
     }
 }
 
+/// Libfunc for creating a general cheatcode.
+#[derive(Default)]
+pub struct CheatcodeLibfunc {}
+impl NamedLibfunc for CheatcodeLibfunc {
+    type Concrete = CheatcodeConcreteLibfunc;
+    const STR_ID: &'static str = "cheatcode";
+
+    fn specialize_signature(
+        &self,
+        context: &dyn SignatureSpecializationContext,
+        args: &[GenericArg],
+    ) -> Result<LibfuncSignature, SpecializationError> {
+        if args.len() != 1 {
+            return Err(SpecializationError::WrongNumberOfGenericArgs);
+        }
+
+        let span_ty = felt252_span_ty(context)?;
+        Ok(LibfuncSignature {
+            param_signatures: vec![
+                // input
+                ParamSignature::new(span_ty.clone()),
+            ],
+            branch_signatures: vec![BranchSignature {
+                vars: vec![
+                    // output
+                    OutputVarInfo {
+                        ty: span_ty,
+                        ref_info: OutputVarReferenceInfo::NewTempVar { idx: 0 },
+                    },
+                ],
+                ap_change: SierraApChange::Known { new_vars_only: true },
+            }],
+            fallthrough: Some(0),
+        })
+    }
+
+    fn specialize(
+        &self,
+        context: &dyn SpecializationContext,
+        args: &[GenericArg],
+    ) -> Result<Self::Concrete, SpecializationError> {
+        match args {
+            [GenericArg::Value(selector)] => Ok(CheatcodeConcreteLibfunc {
+                selector: selector.clone(),
+                signature: <Self as NamedLibfunc>::specialize_signature(
+                    self,
+                    context.upcast(),
+                    args,
+                )?,
+            }),
+            [_] => Err(SpecializationError::UnsupportedGenericArg),
+            _ => Err(SpecializationError::WrongNumberOfGenericArgs),
+        }
+    }
+}
+
+pub struct CheatcodeConcreteLibfunc {
+    pub selector: BigInt,
+    pub signature: LibfuncSignature,
+}
+impl SignatureBasedConcreteLibfunc for CheatcodeConcreteLibfunc {
+    fn signature(&self) -> &LibfuncSignature {
+        &self.signature
+    }
+}
+
 define_libfunc_hierarchy! {
     pub enum TestingLibfunc {
-         SetBlockNumber(TestSetterLibfunc<SetBlockNumberTrait>),
          SetBlockTimestamp(TestSetterLibfunc<SetBlockTimestampTrait>),
          SetCallerAddress(TestSetterLibfunc<SetCallerAddressTrait>),
          SetContractAddress(TestSetterLibfunc<SetContractAddressTrait>),
@@ -211,5 +273,6 @@ define_libfunc_hierarchy! {
          SetNonce(TestSetterLibfunc<SetNonceTrait>),
          SetSignature(TestSetterLibfunc<SetSignatureTrait>),
          PopLog(PopLogLibfunc),
+         Cheatcode(CheatcodeLibfunc),
     }, TestingConcreteLibfunc
 }

crates/cairo-lang-starknet/src/allowed_libfuncs_test.rs

@@ -12,7 +12,7 @@ use super::{
 fn experimental_list_includes_all() {
     let blocked_libfuncs = [
         "print",
-        "set_block_number",
+        "cheatcode",
         "set_block_timestamp",
         "set_caller_address",
         "set_contract_address",