Add generate_arg cheatcode

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/fuzzer.rs

@@ -0,0 +1,38 @@
+use anyhow::ensure;
+use num_bigint::RandBigInt;
+use rand::prelude::StdRng;
+use starknet_types_core::felt::Felt;
+use std::sync::{Arc, Mutex};
+
+pub(crate) fn generate_arg(
+    fuzzer_rng: Option<Arc<Mutex<StdRng>>>,
+    min_value: Felt,
+    max_value: Felt,
+) -> anyhow::Result<Felt> {
+    let min_big_int = if min_value > (Felt::MAX + Felt::from(i128::MIN)) && min_value > max_value {
+        // Negative value x is serialized as P + x, where P is the STARK prime number
+        // hence to deserialize and get the actual x we need to subtract P (== Felt::MAX + 1)
+        min_value.to_bigint() - Felt::MAX.to_bigint() - 1
+    } else {
+        min_value.to_bigint()
+    };
+
+    let max_big_int = max_value.to_bigint();
+
+    ensure!(
+        min_big_int <= max_big_int,
+        format!("`generate_arg` cheatcode: `min_value` must be <= `max_value`, provided values after deserialization: {min_big_int} and {max_big_int}")
+    );
+
+    let value = if let Some(fuzzer_rng) = fuzzer_rng {
+        fuzzer_rng
+            .lock()
+            .unwrap()
+            .gen_bigint_range(&min_big_int, &(max_big_int + 1))
+    } else {
+        // `generate_arg` cheatcode can be also used outside the fuzzer context
+        rand::thread_rng().gen_bigint_range(&min_big_int, &(max_big_int + 1))
+    };
+
+    Ok(Felt::from(value))
+}

crates/cheatnet/src/runtime_extensions/forge_runtime_extension/mod.rs

@@ -40,6 +40,7 @@ use conversions::felt::TryInferFormat;
 use conversions::serde::deserialize::BufferReader;
 use conversions::serde::serialize::CairoSerialize;
 use data_transformer::cairo_types::CairoU256;
+use rand::prelude::StdRng;
 use runtime::{
     CheatcodeHandlingResult, EnhancedHintError, ExtendedRuntime, ExtensionLogic,
     SyscallHandlingResult,
@@ -48,16 +49,19 @@ use starknet::signers::SigningKey;
 use starknet_api::{core::ClassHash, deprecated_contract_class::EntryPointType::L1Handler};
 use starknet_types_core::felt::Felt;
 use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
 
 pub mod cheatcodes;
 pub mod contracts_data;
 mod file_operations;
+mod fuzzer;
 
 pub type ForgeRuntime<'a> = ExtendedRuntime<ForgeExtension<'a>>;
 
 pub struct ForgeExtension<'a> {
     pub environment_variables: &'a HashMap<String, String>,
     pub contracts_data: &'a ContractsData,
+    pub fuzzer_rng: Option<Arc<Mutex<StdRng>>>,
 }
 
 // This runtime extension provides an implementation logic for functions from snforge_std library.
@@ -476,6 +480,14 @@ impl<'a> ExtensionLogic for ForgeExtension<'a> {
             "generate_random_felt" => Ok(CheatcodeHandlingResult::from_serializable(
                 generate_random_felt(),
             )),
+            "generate_arg" => {
+                let min_value: Felt = input_reader.read()?;
+                let max_value: Felt = input_reader.read()?;
+
+                Ok(CheatcodeHandlingResult::from_serializable(
+                    fuzzer::generate_arg(self.fuzzer_rng.clone(), min_value, max_value)?,
+                ))
+            }
             _ => Ok(CheatcodeHandlingResult::Forwarded),
         }
     }

crates/forge-runner/src/running.rs

@@ -202,6 +202,7 @@ pub fn run_test_case(
     let forge_extension = ForgeExtension {
         environment_variables: runtime_config.environment_variables,
         contracts_data: runtime_config.contracts_data,
+        fuzzer_rng: None,
     };
 
     let mut forge_runtime = ExtendedRuntime {

snforge_std/src/cheatcodes.cairo

@@ -8,6 +8,7 @@ pub mod storage;
 pub mod execution_info;
 pub mod message_to_l1;
 pub mod generate_random_felt;
+pub mod generate_arg;
 
 /// Enum used to specify how long the target should be cheated for.
 #[derive(Copy, Drop, Serde, PartialEq, Clone, Debug)]

snforge_std/src/cheatcodes/generate_arg.cairo

@@ -0,0 +1,8 @@
+use super::super::_cheatcode::execute_cheatcode_and_deserialize;
+
+// Generates a random number that is used for creating data for fuzz tests
+pub fn generate_arg<T, +Serde<T>, +Drop<T>, +Into<T, felt252>>(min_value: T, max_value: T) -> T {
+    execute_cheatcode_and_deserialize::<
+        'generate_arg'
+    >(array![min_value.into(), max_value.into()].span())
+}