feat(forge): warp and roll on invariant tx (#12616)

grandizzy · web-flow · commit 0584a581bf0c · 2025-11-21T19:41:32.000Z
* feat(forge): warp and roll on invariant tx

* Target test
diff --git a/crates/config/src/invariant.rs b/crates/config/src/invariant.rs
@@ -37,6 +37,10 @@ pub struct InvariantConfig {
     pub timeout: Option<u32>,
     /// Display counterexample as solidity calls.
     pub show_solidity: bool,
+    /// Maximum time (in seconds) between generated txs.
+    pub max_time_delay: Option<u32>,
+    /// Maximum number of blocks elapsed between generated txs.
+    pub max_block_delay: Option<u32>,
 }
 
 impl Default for InvariantConfig {
@@ -55,6 +59,8 @@ impl Default for InvariantConfig {
             show_metrics: true,
             timeout: None,
             show_solidity: false,
+            max_time_delay: None,
+            max_block_delay: None,
         }
     }
 }
diff --git a/crates/evm/evm/src/executors/corpus.rs b/crates/evm/evm/src/executors/corpus.rs
@@ -1,7 +1,7 @@
-use crate::executors::{Executor, RawCallResult};
+use crate::executors::{Executor, RawCallResult, invariant::execute_tx};
 use alloy_dyn_abi::JsonAbiExt;
 use alloy_json_abi::Function;
-use alloy_primitives::{Bytes, U256};
+use alloy_primitives::Bytes;
 use eyre::eyre;
 use foundry_config::FuzzCorpusConfig;
 use foundry_evm_fuzz::{
@@ -225,15 +225,7 @@ impl CorpusManager {
                 let mut executor = executor.clone();
                 for tx in &tx_seq {
                     if can_replay_tx(tx) {
-                        let mut call_result = executor
-                            .call_raw(
-                                tx.sender,
-                                tx.call_details.target,
-                                tx.call_details.calldata.clone(),
-                                U256::ZERO,
-                            )
-                            .map_err(|e| eyre!(format!("Could not make raw evm call: {e}")))?;
-
+                        let mut call_result = execute_tx(&mut executor, tx)?;
                         let (new_coverage, is_edge) =
                             call_result.merge_edge_coverage(&mut history_map);
                         if new_coverage {
@@ -648,6 +640,8 @@ mod tests {
 
     fn basic_tx() -> BasicTxDetails {
         BasicTxDetails {
+            warp: None,
+            roll: None,
             sender: Address::ZERO,
             call_details: foundry_evm_fuzz::CallDetails {
                 target: Address::ZERO,
diff --git a/crates/evm/evm/src/executors/fuzz/mod.rs b/crates/evm/evm/src/executors/fuzz/mod.rs
@@ -113,6 +113,8 @@ impl FuzzedExecutor {
             dictionary_weight => fuzz_calldata_from_state(func.clone(), &state),
         ]
         .prop_map(move |calldata| BasicTxDetails {
+            warp: None,
+            roll: None,
             sender: Default::default(),
             call_details: CallDetails { target: Default::default(), calldata },
         });
@@ -321,6 +323,8 @@ impl FuzzedExecutor {
         let new_coverage = coverage_metrics.merge_edge_coverage(&mut call);
         coverage_metrics.process_inputs(
             &[BasicTxDetails {
+                warp: None,
+                roll: None,
                 sender: self.sender,
                 call_details: CallDetails { target: address, calldata: calldata.clone() },
             }],
diff --git a/crates/evm/evm/src/executors/invariant/mod.rs b/crates/evm/evm/src/executors/invariant/mod.rs
@@ -395,16 +395,7 @@ impl<'a> InvariantExecutor<'a> {
 
                 // Execute call from the randomly generated sequence without committing state.
                 // State is committed only if call is not a magic assume.
-                let mut call_result = current_run
-                    .executor
-                    .call_raw(
-                        tx.sender,
-                        tx.call_details.target,
-                        tx.call_details.calldata.clone(),
-                        U256::ZERO,
-                    )
-                    .map_err(|e| eyre!(format!("Could not make raw evm call: {e}")))?;
-
+                let mut call_result = execute_tx(&mut current_run.executor, tx)?;
                 let discarded = call_result.result.as_ref() == MAGIC_ASSUME;
                 if self.config.show_metrics {
                     invariant_test.record_metrics(tx, call_result.reverted, discarded);
@@ -583,7 +574,7 @@ impl<'a> InvariantExecutor<'a> {
             fuzz_state.clone(),
             targeted_senders,
             targeted_contracts.clone(),
-            self.config.dictionary.dictionary_weight,
+            self.config.clone(),
             fuzz_fixtures.clone(),
         )
         .no_shrink();
@@ -1037,3 +1028,29 @@ pub(crate) fn call_invariant_function(
     let success = executor.is_raw_call_mut_success(address, &mut call_result, false);
     Ok((call_result, success))
 }
+
+/// Calls the invariant selector and returns call result and if succeeded.
+/// Updates the block number and block timestamp if configured.
+pub(crate) fn execute_tx(executor: &mut Executor, tx: &BasicTxDetails) -> Result<RawCallResult> {
+    // Apply pre-call block adjustments.
+    if let Some(warp) = tx.warp {
+        executor.env_mut().evm_env.block_env.timestamp += warp;
+    }
+    if let Some(roll) = tx.roll {
+        executor.env_mut().evm_env.block_env.number += roll;
+    }
+
+    // Perform the raw call.
+    let mut call_result = executor
+        .call_raw(tx.sender, tx.call_details.target, tx.call_details.calldata.clone(), U256::ZERO)
+        .map_err(|e| eyre!(format!("Could not make raw evm call: {e}")))?;
+
+    // Propagate block adjustments to call result which will be committed.
+    if let Some(warp) = tx.warp {
+        call_result.env.evm_env.block_env.timestamp += warp;
+    }
+    if let Some(roll) = tx.roll {
+        call_result.env.evm_env.block_env.number += roll;
+    }
+    Ok(call_result)
+}
diff --git a/crates/evm/evm/src/executors/invariant/replay.rs b/crates/evm/evm/src/executors/invariant/replay.rs
@@ -1,7 +1,7 @@
-use super::{call_after_invariant_function, call_invariant_function};
+use super::{call_after_invariant_function, call_invariant_function, execute_tx};
 use crate::executors::{EarlyExit, Executor, invariant::shrink::shrink_sequence};
 use alloy_dyn_abi::JsonAbiExt;
-use alloy_primitives::{Log, U256, map::HashMap};
+use alloy_primitives::{Log, map::HashMap};
 use eyre::Result;
 use foundry_common::{ContractsByAddress, ContractsByArtifact};
 use foundry_config::InvariantConfig;
@@ -36,13 +36,7 @@ pub fn replay_run(
 
     // Replay each call from the sequence, collect logs, traces and coverage.
     for tx in inputs {
-        let call_result = executor.transact_raw(
-            tx.sender,
-            tx.call_details.target,
-            tx.call_details.calldata.clone(),
-            U256::ZERO,
-        )?;
-
+        let call_result = execute_tx(&mut executor, tx)?;
         logs.extend(call_result.logs);
         traces.push((TraceKind::Execution, call_result.traces.clone().unwrap()));
         HitMaps::merge_opt(line_coverage, call_result.line_coverage);
@@ -53,9 +47,7 @@ pub fn replay_run(
 
         // Create counter example to be used in failed case.
         counterexample_sequence.push(BaseCounterExample::from_invariant_call(
-            tx.sender,
-            tx.call_details.target,
-            &tx.call_details.calldata,
+            tx,
             &ided_contracts,
             call_result.traces,
             show_solidity,
diff --git a/crates/evm/evm/src/executors/invariant/shrink.rs b/crates/evm/evm/src/executors/invariant/shrink.rs
@@ -1,8 +1,8 @@
 use crate::executors::{
     EarlyExit, Executor,
-    invariant::{call_after_invariant_function, call_invariant_function},
+    invariant::{call_after_invariant_function, call_invariant_function, execute_tx},
 };
-use alloy_primitives::{Address, Bytes, U256};
+use alloy_primitives::{Address, Bytes};
 use foundry_config::InvariantConfig;
 use foundry_evm_core::constants::MAGIC_ASSUME;
 use foundry_evm_fuzz::{BasicTxDetails, invariant::InvariantContract};
@@ -118,12 +118,8 @@ pub fn check_sequence(
     // Apply the call sequence.
     for call_index in sequence {
         let tx = &calls[call_index];
-        let call_result = executor.transact_raw(
-            tx.sender,
-            tx.call_details.target,
-            tx.call_details.calldata.clone(),
-            U256::ZERO,
-        )?;
+        let mut call_result = execute_tx(&mut executor, tx)?;
+        executor.commit(&mut call_result);
         // Ignore calls reverted with `MAGIC_ASSUME`. This is needed to handle failed scenarios that
         // are replayed with a modified version of test driver (that use new `vm.assume`
         // cheatcodes).
diff --git a/crates/evm/fuzz/src/invariant/call_override.rs b/crates/evm/fuzz/src/invariant/call_override.rs
@@ -75,12 +75,9 @@ impl RandomCallGenerator {
             *self.target_reference.write() = original_caller;
 
             // `original_caller` has a 80% chance of being the `new_target`.
-            let choice = self
-                .strategy
-                .new_tree(&mut self.runner.lock())
-                .unwrap()
-                .current()
-                .map(|call_details| BasicTxDetails { sender, call_details });
+            let choice = self.strategy.new_tree(&mut self.runner.lock()).unwrap().current().map(
+                |call_details| BasicTxDetails { warp: None, roll: None, sender, call_details },
+            );
 
             self.last_sequence.write().push(choice.clone());
             choice
diff --git a/crates/evm/fuzz/src/lib.rs b/crates/evm/fuzz/src/lib.rs
@@ -10,7 +10,7 @@ extern crate tracing;
 
 use alloy_dyn_abi::{DynSolValue, JsonAbiExt};
 use alloy_primitives::{
-    Address, Bytes, Log,
+    Address, Bytes, Log, U256,
     map::{AddressHashMap, HashMap},
 };
 use foundry_common::{calc, contracts::ContractsByAddress};
@@ -36,6 +36,10 @@ pub use inspector::Fuzzer;
 /// Details of a transaction generated by fuzz strategy for fuzzing a target.
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct BasicTxDetails {
+    // Time (in seconds) to increase block timestamp before executing the tx.
+    pub warp: Option<U256>,
+    // Number to increase block number before executing the tx.
+    pub roll: Option<U256>,
     // Transaction sender address.
     pub sender: Address,
     // Transaction call details.
@@ -62,6 +66,10 @@ pub enum CounterExample {
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct BaseCounterExample {
+    // Amount to increase block timestamp.
+    pub warp: Option<U256>,
+    // Amount to increase block number.
+    pub roll: Option<U256>,
     /// Address which makes the call.
     pub sender: Option<Address>,
     /// Address to which to call to.
@@ -89,21 +97,26 @@ pub struct BaseCounterExample {
 impl BaseCounterExample {
     /// Creates counter example representing a step from invariant call sequence.
     pub fn from_invariant_call(
-        sender: Address,
-        addr: Address,
-        bytes: &Bytes,
+        tx: &BasicTxDetails,
         contracts: &ContractsByAddress,
         traces: Option<SparsedTraceArena>,
         show_solidity: bool,
     ) -> Self {
-        if let Some((name, abi)) = &contracts.get(&addr)
+        let sender = tx.sender;
+        let target = tx.call_details.target;
+        let bytes = &tx.call_details.calldata;
+        let warp = tx.warp;
+        let roll = tx.roll;
+        if let Some((name, abi)) = &contracts.get(&target)
             && let Some(func) = abi.functions().find(|f| f.selector() == bytes[..4])
         {
             // skip the function selector when decoding
             if let Ok(args) = func.abi_decode_input(&bytes[4..]) {
                 return Self {
+                    warp,
+                    roll,
                     sender: Some(sender),
-                    addr: Some(addr),
+                    addr: Some(target),
                     calldata: bytes.clone(),
                     contract_name: Some(name.clone()),
                     func_name: Some(func.name.clone()),
@@ -119,8 +132,10 @@ impl BaseCounterExample {
         }
 
         Self {
+            warp,
+            roll,
             sender: Some(sender),
-            addr: Some(addr),
+            addr: Some(target),
             calldata: bytes.clone(),
             contract_name: None,
             func_name: None,
@@ -139,6 +154,8 @@ impl BaseCounterExample {
         traces: Option<SparsedTraceArena>,
     ) -> Self {
         Self {
+            warp: None,
+            roll: None,
             sender: None,
             addr: None,
             calldata: bytes,
@@ -160,6 +177,12 @@ impl fmt::Display for BaseCounterExample {
             && let (Some(sender), Some(contract), Some(address), Some(func_name), Some(args)) =
                 (&self.sender, &self.contract_name, &self.addr, &self.func_name, &self.raw_args)
         {
+            if let Some(warp) = &self.warp {
+                writeln!(f, "\t\tvm.warp(block.timestamp + {warp});")?;
+            }
+            if let Some(roll) = &self.roll {
+                writeln!(f, "\t\tvm.roll(block.number + {roll});")?;
+            }
             writeln!(f, "\t\tvm.prank({sender});")?;
             write!(
                 f,
@@ -186,6 +209,13 @@ impl fmt::Display for BaseCounterExample {
             write!(f, "{addr} ")?
         }
 
+        if let Some(warp) = &self.warp {
+            write!(f, "warp={warp} ")?;
+        }
+        if let Some(roll) = &self.roll {
+            write!(f, "roll={roll} ")?;
+        }
+
         if let Some(sig) = &self.signature {
             write!(f, "calldata={sig}")?
         } else {
diff --git a/crates/evm/fuzz/src/strategies/invariants.rs b/crates/evm/fuzz/src/strategies/invariants.rs
@@ -5,7 +5,8 @@ use crate::{
     strategies::{EvmFuzzState, fuzz_calldata_from_state, fuzz_param},
 };
 use alloy_json_abi::Function;
-use alloy_primitives::Address;
+use alloy_primitives::{Address, U256};
+use foundry_config::InvariantConfig;
 use parking_lot::RwLock;
 use proptest::prelude::*;
 use rand::seq::IteratorRandom;
@@ -60,25 +61,44 @@ pub fn invariant_strat(
     fuzz_state: EvmFuzzState,
     senders: SenderFilters,
     contracts: FuzzRunIdentifiedContracts,
-    dictionary_weight: u32,
+    config: InvariantConfig,
     fuzz_fixtures: FuzzFixtures,
 ) -> impl Strategy<Value = BasicTxDetails> {
     let senders = Rc::new(senders);
+    let dictionary_weight = config.dictionary.dictionary_weight;
+
+    // Strategy to generate values for tx warp and roll.
+    let warp_roll_strat = |cond: bool| {
+        if cond { any::<U256>().prop_map(Some).boxed() } else { Just(None).boxed() }
+    };
+
     any::<prop::sample::Selector>()
         .prop_flat_map(move |selector| {
             let contracts = contracts.targets.lock();
             let functions = contracts.fuzzed_functions();
             let (target_address, target_function) = selector.select(functions);
+
             let sender = select_random_sender(&fuzz_state, senders.clone(), dictionary_weight);
+
             let call_details = fuzz_contract_with_calldata(
                 &fuzz_state,
                 &fuzz_fixtures,
                 *target_address,
                 target_function.clone(),
             );
-            (sender, call_details)
+
+            let warp = warp_roll_strat(config.max_time_delay.is_some());
+            let roll = warp_roll_strat(config.max_block_delay.is_some());
+
+            (warp, roll, sender, call_details)
+        })
+        .prop_map(move |(warp, roll, sender, call_details)| {
+            let warp =
+                warp.map(|time| time % U256::from(config.max_time_delay.unwrap_or_default()));
+            let roll =
+                roll.map(|block| block % U256::from(config.max_block_delay.unwrap_or_default()));
+            BasicTxDetails { warp, roll, sender, call_details }
         })
-        .prop_map(|(sender, call_details)| BasicTxDetails { sender, call_details })
 }
 
 /// Strategy to select a sender address:
diff --git a/crates/forge/src/runner.rs b/crates/forge/src/runner.rs
@@ -780,6 +780,8 @@ impl<'a> FunctionRunner<'a> {
                 .map(|seq| {
                     seq.show_solidity = show_solidity;
                     BasicTxDetails {
+                        warp: seq.warp,
+                        roll: seq.roll,
                         sender: seq.sender.unwrap_or_default(),
                         call_details: CallDetails {
                             target: seq.addr.unwrap_or_default(),
diff --git a/crates/forge/tests/cli/config.rs b/crates/forge/tests/cli/config.rs
@@ -1272,7 +1272,9 @@ forgetest_init!(test_default_config, |prj, cmd| {
     "failure_persist_dir": "cache/invariant",
     "show_metrics": true,
     "timeout": null,
-    "show_solidity": false
+    "show_solidity": false,
+    "max_time_delay": null,
+    "max_block_delay": null
   },
   "ffi": false,
   "allow_internal_expect_revert": false,
diff --git a/crates/forge/tests/cli/test_cmd/invariant/common.rs b/crates/forge/tests/cli/test_cmd/invariant/common.rs

Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,10 @@ pub struct InvariantConfig {`
`37`	`37`	`pub timeout: Option<u32>,`
`38`	`38`	`/// Display counterexample as solidity calls.`
`39`	`39`	`pub show_solidity: bool,`
	`40`	`+ /// Maximum time (in seconds) between generated txs.`
	`41`	`+ pub max_time_delay: Option<u32>,`
	`42`	`+ /// Maximum number of blocks elapsed between generated txs.`
	`43`	`+ pub max_block_delay: Option<u32>,`
`40`	`44`	`}`
`41`	`45`
`42`	`46`	`impl Default for InvariantConfig {`
`@@ -55,6 +59,8 @@ impl Default for InvariantConfig {`
`55`	`59`	`show_metrics: true,`
`56`	`60`	`timeout: None,`
`57`	`61`	`show_solidity: false,`
	`62`	`+ max_time_delay: None,`
	`63`	`+ max_block_delay: None,`
`58`	`64`	`}`
`59`	`65`	`}`
`60`	`66`	`}`