fix(symbolic): replay setup arbitrary storage

Import setup-time vm.setArbitraryStorage targets into symbolic invariant execution and carry model values for concrete storage slots into replay. This lets invariant counterexamples that depend on symbolic storage behind deployed setup dependencies replay against Forge's concrete executor.

Author: figtracer

crates/cheatcodes/src/evm.rs

@@ -288,7 +288,10 @@ impl Cheatcode for loadCall {
             if ccx.state.has_arbitrary_storage(&target) {
                 // If storage slot is untouched and load from a target with arbitrary storage,
                 // then set random value for current slot.
-                let rand_value = ccx.state.rng().random();
+                let rand_value = ccx
+                    .state
+                    .cached_arbitrary_storage_value(target, slot.into())
+                    .unwrap_or_else(|| ccx.state.rng().random());
                 ccx.state.arbitrary_storage.as_mut().unwrap().save(
                     ccx.ecx,
                     target,

crates/cheatcodes/src/inspector.rs

@@ -403,11 +403,11 @@ pub struct ArbitraryStorage {
     /// Mapping of arbitrary storage addresses to generated values (slot, arbitrary value).
     /// (SLOADs return random value if storage slot wasn't accessed).
     /// Changed values are recorded and used to copy storage to different addresses.
-    pub values: HashMap<Address, HashMap<U256, U256>>,
+    values: HashMap<Address, HashMap<U256, U256>>,
     /// Mapping of address with storage copied to arbitrary storage address source.
-    pub copies: HashMap<Address, Address>,
+    copies: HashMap<Address, Address>,
     /// Address with storage slots that should be overwritten even if previously set.
-    pub overwrites: HashSet<Address>,
+    overwrites: HashSet<Address>,
 }
 
 impl ArbitraryStorage {
@@ -428,6 +428,23 @@ impl ArbitraryStorage {
         }
     }
 
+    /// Returns addresses explicitly marked with arbitrary storage.
+    fn targets(&self) -> impl Iterator<Item = Address> + '_ {
+        self.values.keys().copied()
+    }
+
+    /// Caches a concrete value for a slot on an arbitrary-storage address.
+    fn cache_value(&mut self, address: Address, slot: U256, data: U256) {
+        if let Some(values) = self.values.get_mut(&address) {
+            values.insert(slot, data);
+        }
+    }
+
+    /// Returns a cached arbitrary value for a slot.
+    fn cached_value(&self, address: Address, slot: U256) -> Option<U256> {
+        self.values.get(&address).and_then(|values| values.get(&slot)).copied()
+    }
+
     /// Saves arbitrary storage value for a given address:
     /// - store value in changed values cache.
     /// - update account's storage with given value.
@@ -1298,6 +1315,23 @@ impl<FEN: FoundryEvmNetwork> Cheatcodes<FEN> {
         self.arbitrary_storage.get_or_insert_with(ArbitraryStorage::default)
     }
 
+    /// Returns addresses explicitly marked with arbitrary storage.
+    pub fn arbitrary_storage_targets(&self) -> impl Iterator<Item = Address> + '_ {
+        self.arbitrary_storage.as_ref().into_iter().flat_map(ArbitraryStorage::targets)
+    }
+
+    /// Caches a concrete replay value for a slot on an arbitrary-storage address.
+    pub fn cache_arbitrary_storage_value(&mut self, address: Address, slot: U256, value: U256) {
+        if let Some(storage) = &mut self.arbitrary_storage {
+            storage.cache_value(address, slot, value);
+        }
+    }
+
+    /// Returns a cached arbitrary-storage replay value for a slot.
+    pub fn cached_arbitrary_storage_value(&self, address: Address, slot: U256) -> Option<U256> {
+        self.arbitrary_storage.as_ref().and_then(|storage| storage.cached_value(address, slot))
+    }
+
     /// Whether the given address has arbitrary storage.
     pub fn has_arbitrary_storage(&self, address: &Address) -> bool {
         match &self.arbitrary_storage {
@@ -2535,7 +2569,9 @@ impl<FEN: FoundryEvmNetwork> Cheatcodes<FEN> {
             || self.should_overwrite_arbitrary_storage(&target_address, key)
         {
             if self.has_arbitrary_storage(&target_address) {
-                let arbitrary_value = self.rng().random();
+                let arbitrary_value = self
+                    .cached_arbitrary_storage_value(target_address, key)
+                    .unwrap_or_else(|| self.rng().random());
                 self.arbitrary_storage.as_mut().unwrap().save(
                     ecx,
                     target_address,

crates/evm/symbolic/src/executor/invariant.rs

@@ -224,9 +224,9 @@ impl SymbolicExecutor {
         &mut self,
         steps: &[SequenceStepTemplate],
         state: &PathState,
-    ) -> Result<Vec<SymbolicInvariantStep>, SymbolicError> {
+    ) -> Result<(Vec<SymbolicInvariantStep>, Vec<SymbolicStorageAssignment>), SymbolicError> {
         let model = self.solver.model(&mut self.cx, &state.constraints)?;
-        steps
+        let sequence = steps
             .iter()
             .map(|step| {
                 let args = step.calldata.model_to_args(&mut self.cx, &model)?;
@@ -241,6 +241,8 @@ impl SymbolicExecutor {
                     calldata,
                 })
             })
-            .collect()
+            .collect::<Result<Vec<_>, SymbolicError>>()?;
+        let storage = state.world.replay_storage_assignments(&model);
+        Ok((sequence, storage))
     }
 }

crates/evm/symbolic/src/executor/run.rs

@@ -506,9 +506,11 @@ impl SymbolicExecutor {
             &mut completed_paths,
         )? {
             if outcome.failed {
-                let sequence = self.materialize_sequence(&initial.steps, &outcome.state)?;
+                let (sequence, storage) =
+                    self.materialize_sequence(&initial.steps, &outcome.state)?;
                 return Ok(SymbolicInvariantRunResult::Counterexample {
                     sequence,
+                    storage,
                     stats: self.stats_with_paths(completed_paths),
                 });
             }
@@ -558,20 +560,22 @@ impl SymbolicExecutor {
 
                                 match outcome.status {
                                     TopLevelCallStatus::Failure => {
-                                        let sequence =
+                                        let (sequence, storage) =
                                             self.materialize_sequence(&steps, &outcome.state)?;
                                         return Ok(SymbolicInvariantRunResult::Counterexample {
                                             sequence,
+                                            storage,
                                             stats: self.stats_with_paths(completed_paths),
                                         });
                                     }
                                     TopLevelCallStatus::Revert => {
                                         if input.fail_on_revert {
-                                            let sequence =
+                                            let (sequence, storage) =
                                                 self.materialize_sequence(&steps, &outcome.state)?;
                                             return Ok(
                                                 SymbolicInvariantRunResult::Counterexample {
                                                     sequence,
+                                                    storage,
                                                     stats: self.stats_with_paths(completed_paths),
                                                 },
                                             );
@@ -592,13 +596,15 @@ impl SymbolicExecutor {
                                             &mut completed_paths,
                                         )? {
                                             if invariant_outcome.failed {
-                                                let sequence = self.materialize_sequence(
-                                                    &steps,
-                                                    &invariant_outcome.state,
-                                                )?;
+                                                let (sequence, storage) = self
+                                                    .materialize_sequence(
+                                                        &steps,
+                                                        &invariant_outcome.state,
+                                                    )?;
                                                 return Ok(
                                                     SymbolicInvariantRunResult::Counterexample {
                                                         sequence,
+                                                        storage,
                                                         stats: self
                                                             .stats_with_paths(completed_paths),
                                                     },

crates/evm/symbolic/src/lib.rs

@@ -340,6 +340,17 @@ pub struct SymbolicInvariantRunInput<'a, FEN: FoundryEvmNetwork> {
     pub ffi_enabled: bool,
 }
 
+/// One concrete storage value required to replay a symbolic invariant candidate.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct SymbolicStorageAssignment {
+    /// Account whose storage slot should be initialized.
+    pub address: Address,
+    /// Concrete storage slot.
+    pub slot: U256,
+    /// Concrete value extracted from the solver model.
+    pub value: U256,
+}
+
 /// Outcome of bounded symbolic invariant execution.
 #[derive(Clone, Debug)]
 pub enum SymbolicInvariantRunResult {
@@ -349,6 +360,8 @@ pub enum SymbolicInvariantRunResult {
     Counterexample {
         /// Concrete sequence extracted from the solver model.
         sequence: Vec<SymbolicInvariantStep>,
+        /// Concrete setup-storage values needed for replay.
+        storage: Vec<SymbolicStorageAssignment>,
         /// Execution counters collected before the counterexample was returned.
         stats: SymbolicStats,
     },

crates/evm/symbolic/src/runtime/state.rs

@@ -143,6 +143,11 @@ impl PathState {
             .and_then(|cheats| cheats.gas_price)
             .unwrap_or_else(|| executor.tx_env().gas_price());
         self.gas_price = SymExpr::constant(cx, U256::from(gas_price));
+        if let Some(cheats) = executor.inspector().cheatcodes.as_ref() {
+            for target in cheats.arbitrary_storage_targets() {
+                self.world.enable_arbitrary_storage(target);
+            }
+        }
     }
 
     pub(crate) fn child(&self, frame: CallFrame) -> Self {
@@ -1475,6 +1480,7 @@ struct SymbolicWorldSnapshot {
     arbitrary_storage_all: bool,
     zero_init_symbolic_storage: bool,
     symbolic_address_aliases: HashMap<SymExpr, Address>,
+    replay_storage_slots: HashMap<Symbol, SymbolicReplayStorageSlot>,
 }
 
 impl From<&SymbolicWorld> for SymbolicWorldSnapshot {
@@ -1494,10 +1500,17 @@ impl From<&SymbolicWorld> for SymbolicWorldSnapshot {
             arbitrary_storage_all: world.arbitrary_storage_all,
             zero_init_symbolic_storage: world.zero_init_symbolic_storage,
             symbolic_address_aliases: world.symbolic_address_aliases.clone(),
+            replay_storage_slots: world.replay_storage_slots.clone(),
         }
     }
 }
 
+#[derive(Clone, Debug)]
+struct SymbolicReplayStorageSlot {
+    address: Address,
+    slot: U256,
+}
+
 #[derive(Clone, Debug, Default)]
 pub(crate) struct SymbolicWorld {
     storage: Vec<StorageWrite>,
@@ -1512,6 +1525,7 @@ pub(crate) struct SymbolicWorld {
     arbitrary_storage_all: bool,
     zero_init_symbolic_storage: bool,
     symbolic_address_aliases: HashMap<SymExpr, Address>,
+    replay_storage_slots: HashMap<Symbol, SymbolicReplayStorageSlot>,
     snapshots: HashMap<U256, SymbolicWorldSnapshot>,
     next_snapshot_id: u64,
 }
@@ -1547,7 +1561,7 @@ impl SymbolicWorld {
     }
 
     pub(crate) fn sload<FEN: FoundryEvmNetwork>(
-        &self,
+        &mut self,
         cx: &mut SymCx,
         executor: &Executor<FEN>,
         address: Address,
@@ -1591,6 +1605,22 @@ impl SymbolicWorld {
         self.arbitrary_storage_accounts.insert(address);
     }
 
+    pub(crate) fn replay_storage_assignments(
+        &self,
+        model: &SymbolicModel,
+    ) -> Vec<SymbolicStorageAssignment> {
+        self.replay_storage_slots
+            .iter()
+            .filter_map(|(symbol, slot)| {
+                model.get(symbol).copied().map(|value| SymbolicStorageAssignment {
+                    address: slot.address,
+                    slot: slot.slot,
+                    value,
+                })
+            })
+            .collect()
+    }
+
     pub(crate) fn resolve_address(&self, expr: &SymExpr) -> Option<Address> {
         expr.as_const().map(word_to_address).or_else(|| {
             self.symbolic_address_aliases.get(expr).copied().or_else(|| {
@@ -1639,6 +1669,7 @@ impl SymbolicWorld {
         self.arbitrary_storage_all = snapshot.arbitrary_storage_all;
         self.zero_init_symbolic_storage = snapshot.zero_init_symbolic_storage;
         self.symbolic_address_aliases = snapshot.symbolic_address_aliases;
+        self.replay_storage_slots = snapshot.replay_storage_slots;
         true
     }
 
@@ -1651,15 +1682,21 @@ impl SymbolicWorld {
     }
 
     pub(crate) fn storage_base<FEN: FoundryEvmNetwork>(
-        &self,
+        &mut self,
         cx: &mut SymCx,
         executor: &Executor<FEN>,
         address: Address,
         key: &SymExpr,
         concrete_key: Option<U256>,
     ) -> Result<SymExpr, SymbolicError> {
-        if self.arbitrary_storage_all || self.arbitrary_storage_accounts.contains(&address) {
-            let name = stable_symbol("storage", format!("{address:?}:{key:?}").as_bytes());
+        let has_arbitrary_storage = self.arbitrary_storage_accounts.contains(&address);
+        if self.arbitrary_storage_all || has_arbitrary_storage {
+            let name = symbolic_storage_symbol(address, key);
+            if has_arbitrary_storage && let Some(slot) = concrete_key.or_else(|| key.as_const()) {
+                self.replay_storage_slots
+                    .entry(name.clone())
+                    .or_insert(SymbolicReplayStorageSlot { address, slot });
+            }
             return Ok(SymExpr::var_symbol(cx, name));
         }
         if let Some(key) = concrete_key {
@@ -1678,7 +1715,7 @@ impl SymbolicWorld {
         } else if self.zero_init_symbolic_storage {
             Ok(SymExpr::zero(cx))
         } else {
-            let name = stable_symbol("storage", format!("{address:?}:{key:?}").as_bytes());
+            let name = symbolic_storage_symbol(address, key);
             Ok(SymExpr::var_symbol(cx, name))
         }
     }
@@ -2092,6 +2129,10 @@ impl SymbolicWorld {
     }
 }
 
+fn symbolic_storage_symbol(address: Address, key: &SymExpr) -> Symbol {
+    stable_symbol("storage", format!("{address:?}:{key:?}").as_bytes())
+}
+
 #[derive(Clone, Debug)]
 pub(crate) struct SymbolicBlock {
     pub(crate) chain_id: SymExpr,