perf(fuzz): improve Fuzz Dictionary size enforcement (#4680)

* perf(fuzz): improve Fuzz Dictionary size enforcment

* update tests

* flaky test

config/src/fuzz.rs

@@ -19,6 +19,25 @@ pub struct FuzzConfig {
         deserialize_with = "ethers_core::types::serde_helpers::deserialize_stringified_numeric_opt"
     )]
     pub seed: Option<U256>,
+    /// The fuzz dictionary configuration
+    #[serde(flatten)]
+    pub dictionary: FuzzDictionaryConfig,
+}
+
+impl Default for FuzzConfig {
+    fn default() -> Self {
+        FuzzConfig {
+            runs: 256,
+            max_test_rejects: 65536,
+            seed: None,
+            dictionary: FuzzDictionaryConfig::default(),
+        }
+    }
+}
+
+/// Contains for fuzz testing
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+pub struct FuzzDictionaryConfig {
     /// The weight of the dictionary
     #[serde(deserialize_with = "crate::deserialize_stringified_percent")]
     pub dictionary_weight: u32,
@@ -38,17 +57,14 @@ pub struct FuzzConfig {
     pub max_fuzz_dictionary_values: usize,
 }
 
-impl Default for FuzzConfig {
+impl Default for FuzzDictionaryConfig {
     fn default() -> Self {
-        FuzzConfig {
-            runs: 256,
-            max_test_rejects: 65536,
-            seed: None,
+        FuzzDictionaryConfig {
             dictionary_weight: 40,
             include_storage: true,
             include_push_bytes: true,
-            // limit this to 200MB
-            max_fuzz_dictionary_addresses: (200 * 1024 * 1024) / 20,
+            // limit this to 300MB
+            max_fuzz_dictionary_addresses: (300 * 1024 * 1024) / 20,
             // limit this to 200MB
             max_fuzz_dictionary_values: (200 * 1024 * 1024) / 32,
         }

config/src/invariant.rs

@@ -1,5 +1,6 @@
 //! Configuration for invariant testing
 
+use crate::fuzz::FuzzDictionaryConfig;
 use serde::{Deserialize, Serialize};
 
 /// Contains for invariant testing
@@ -14,13 +15,9 @@ pub struct InvariantConfig {
     /// Allows overriding an unsafe external call when running invariant tests. eg. reentrancy
     /// checks
     pub call_override: bool,
-    /// The weight of the dictionary
-    #[serde(deserialize_with = "crate::deserialize_stringified_percent")]
-    pub dictionary_weight: u32,
-    /// The flag indicating whether to include values from storage
-    pub include_storage: bool,
-    /// The flag indicating whether to include push bytes values
-    pub include_push_bytes: bool,
+    /// The fuzz dictionary configuration
+    #[serde(flatten)]
+    pub dictionary: FuzzDictionaryConfig,
 }
 
 impl Default for InvariantConfig {
@@ -30,9 +27,7 @@ impl Default for InvariantConfig {
             depth: 15,
             fail_on_revert: false,
             call_override: false,
-            dictionary_weight: 80,
-            include_storage: true,
-            include_push_bytes: true,
+            dictionary: FuzzDictionaryConfig { dictionary_weight: 80, ..Default::default() },
         }
     }
 }

config/src/lib.rs

@@ -85,7 +85,7 @@ use crate::{
 use providers::*;
 
 mod fuzz;
-pub use fuzz::FuzzConfig;
+pub use fuzz::{FuzzConfig, FuzzDictionaryConfig};
 
 mod invariant;
 use crate::fs_permissions::PathPermission;
@@ -3638,18 +3638,33 @@ mod tests {
             assert_ne!(config.invariant.runs, config.fuzz.runs);
             assert_eq!(config.invariant.runs, 420);
 
-            assert_ne!(config.fuzz.include_storage, invariant_default.include_storage);
-            assert_eq!(config.invariant.include_storage, config.fuzz.include_storage);
+            assert_ne!(
+                config.fuzz.dictionary.include_storage,
+                invariant_default.dictionary.include_storage
+            );
+            assert_eq!(
+                config.invariant.dictionary.include_storage,
+                config.fuzz.dictionary.include_storage
+            );
 
-            assert_ne!(config.fuzz.dictionary_weight, invariant_default.dictionary_weight);
-            assert_eq!(config.invariant.dictionary_weight, config.fuzz.dictionary_weight);
+            assert_ne!(
+                config.fuzz.dictionary.dictionary_weight,
+                invariant_default.dictionary.dictionary_weight
+            );
+            assert_eq!(
+                config.invariant.dictionary.dictionary_weight,
+                config.fuzz.dictionary.dictionary_weight
+            );
 
             jail.set_env("FOUNDRY_PROFILE", "ci");
             let ci_config = Config::load();
             assert_eq!(ci_config.fuzz.runs, 1);
             assert_eq!(ci_config.invariant.runs, 400);
-            assert_eq!(ci_config.fuzz.dictionary_weight, 5);
-            assert_eq!(ci_config.invariant.dictionary_weight, config.fuzz.dictionary_weight);
+            assert_eq!(ci_config.fuzz.dictionary.dictionary_weight, 5);
+            assert_eq!(
+                ci_config.invariant.dictionary.dictionary_weight,
+                config.fuzz.dictionary.dictionary_weight
+            );
 
             Ok(())
         })
@@ -4086,7 +4101,7 @@ mod tests {
 
             let config = Config::load();
             assert_eq!(config.fmt.line_length, 95);
-            assert_eq!(config.fuzz.dictionary_weight, 99);
+            assert_eq!(config.fuzz.dictionary.dictionary_weight, 99);
             assert_eq!(config.invariant.depth, 5);
 
             Ok(())

evm/src/executor/inspector/fuzzer.rs

@@ -20,6 +20,20 @@ impl<DB> Inspector<DB> for Fuzzer
 where
     DB: Database,
 {
+    fn step(
+        &mut self,
+        interpreter: &mut Interpreter,
+        _: &mut EVMData<'_, DB>,
+        _is_static: bool,
+    ) -> Return {
+        // We only collect `stack` and `memory` data before and after calls.
+        if self.collect {
+            self.collect_data(interpreter);
+            self.collect = false;
+        }
+        Return::Continue
+    }
+
     fn call(
         &mut self,
         data: &mut EVMData<'_, DB>,
@@ -38,20 +52,6 @@ where
         (Return::Continue, Gas::new(call.gas_limit), Bytes::new())
     }
 
-    fn step(
-        &mut self,
-        interpreter: &mut Interpreter,
-        _: &mut EVMData<'_, DB>,
-        _is_static: bool,
-    ) -> Return {
-        // We only collect `stack` and `memory` data before and after calls.
-        if self.collect {
-            self.collect_data(interpreter);
-            self.collect = false;
-        }
-        Return::Continue
-    }
-
     fn call_end(
         &mut self,
         _: &mut EVMData<'_, DB>,
@@ -79,7 +79,7 @@ impl Fuzzer {
         let mut state = self.fuzz_state.write();
 
         for slot in interpreter.stack().data() {
-            state.insert(utils::u256_to_h256_be(*slot).into());
+            state.values_mut().insert(utils::u256_to_h256_be(*slot).into());
         }
 
         // TODO: disabled for now since it's flooding the dictionary

evm/src/fuzz/invariant/executor.rs

@@ -24,7 +24,7 @@ use ethers::{
 };
 use eyre::ContextCompat;
 use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
-use foundry_config::InvariantConfig;
+use foundry_config::{FuzzDictionaryConfig, InvariantConfig};
 use hashbrown::HashMap;
 use parking_lot::{Mutex, RwLock};
 use proptest::{
@@ -157,8 +157,7 @@ impl<'a> InvariantExecutor<'a> {
                         sender,
                         &call_result,
                         fuzz_state.clone(),
-                        self.config.include_storage,
-                        self.config.include_push_bytes,
+                        &self.config.dictionary,
                     );
 
                     if let Err(error) = collect_created_contracts(
@@ -220,7 +219,7 @@ impl<'a> InvariantExecutor<'a> {
             });
         }
 
-        tracing::trace!(target: "forge::test::invariant::dictionary", "{:?}", fuzz_state.read().iter().map(hex::encode).collect::<Vec<_>>());
+        tracing::trace!(target: "forge::test::invariant::dictionary", "{:?}", fuzz_state.read().values().iter().map(hex::encode).collect::<Vec<_>>());
 
         let (reverts, invariants) = failures.into_inner().into_inner();
 
@@ -250,11 +249,8 @@ impl<'a> InvariantExecutor<'a> {
         }
 
         // Stores fuzz state for use with [fuzz_calldata_from_state].
-        let fuzz_state: EvmFuzzState = build_initial_state(
-            self.executor.backend().mem_db(),
-            self.config.include_storage,
-            self.config.include_push_bytes,
-        );
+        let fuzz_state: EvmFuzzState =
+            build_initial_state(self.executor.backend().mem_db(), &self.config.dictionary);
 
         // During execution, any newly created contract is added here and used through the rest of
         // the fuzz run.
@@ -266,7 +262,7 @@ impl<'a> InvariantExecutor<'a> {
             fuzz_state.clone(),
             targeted_senders,
             targeted_contracts.clone(),
-            self.config.dictionary_weight,
+            self.config.dictionary.dictionary_weight,
         )
         .no_shrink()
         .boxed();
@@ -529,8 +525,7 @@ fn collect_data(
     sender: &Address,
     call_result: &RawCallResult,
     fuzz_state: EvmFuzzState,
-    include_storage: bool,
-    include_push_bytes: bool,
+    config: &FuzzDictionaryConfig,
 ) {
     // Verify it has no code.
     let mut has_code = false;
@@ -545,13 +540,7 @@ fn collect_data(
         sender_changeset = state_changeset.remove(sender);
     }
 
-    collect_state_from_call(
-        &call_result.logs,
-        &*state_changeset,
-        fuzz_state,
-        include_storage,
-        include_push_bytes,
-    );
+    collect_state_from_call(&call_result.logs, &*state_changeset, fuzz_state, config);
 
     // Re-add changes
     if let Some(changed) = sender_changeset {

evm/src/fuzz/mod.rs

@@ -81,22 +81,17 @@ impl<'a> FuzzedExecutor<'a> {
 
         // Stores fuzz state for use with [fuzz_calldata_from_state]
         let state: EvmFuzzState = if let Some(fork_db) = self.executor.backend().active_fork_db() {
-            build_initial_state(
-                fork_db,
-                self.config.include_storage,
-                self.config.include_push_bytes,
-            )
+            build_initial_state(fork_db, &self.config.dictionary)
         } else {
-            build_initial_state(
-                self.executor.backend().mem_db(),
-                self.config.include_storage,
-                self.config.include_push_bytes,
-            )
+            build_initial_state(self.executor.backend().mem_db(), &self.config.dictionary)
         };
 
         let strat = proptest::strategy::Union::new_weighted(vec![
-            (100 - self.config.dictionary_weight, fuzz_calldata(func.clone())),
-            (self.config.dictionary_weight, fuzz_calldata_from_state(func.clone(), state.clone())),
+            (100 - self.config.dictionary.dictionary_weight, fuzz_calldata(func.clone())),
+            (
+                self.config.dictionary.dictionary_weight,
+                fuzz_calldata_from_state(func.clone(), state.clone()),
+            ),
         ]);
         tracing::debug!(func = ?func.name, should_fail, "fuzzing");
         let run_result = self.runner.clone().run(&strat, |calldata| {
@@ -109,22 +104,12 @@ impl<'a> FuzzedExecutor<'a> {
                 .as_ref()
                 .ok_or_else(|| TestCaseError::fail(FuzzError::EmptyChangeset))?;
 
-            // keep memory in check
-            {
-                let mut state = state.write();
-                state.enforce_limit(
-                    self.config.max_fuzz_dictionary_addresses,
-                    self.config.max_fuzz_dictionary_values,
-                );
-            }
-
             // Build fuzzer state
             collect_state_from_call(
                 &call.logs,
                 state_changeset,
                 state.clone(),
-                self.config.include_storage,
-                self.config.include_push_bytes,
+                &self.config.dictionary,
             );
 
             // When assume cheat code is triggered return a special string "FOUNDRY::ASSUME"

evm/src/fuzz/strategies/param.rs

@@ -56,13 +56,13 @@ pub fn fuzz_param(param: &ParamType) -> impl Strategy<Value = Token> {
 /// Works with ABI Encoder v2 tuples.
 pub fn fuzz_param_from_state(param: &ParamType, arc_state: EvmFuzzState) -> BoxedStrategy<Token> {
     // These are to comply with lifetime requirements
-    let state_len = arc_state.read().len();
+    let state_len = arc_state.read().values().len();
 
     // Select a value from the state
     let st = arc_state.clone();
     let value = any::<prop::sample::Index>()
         .prop_map(move |index| index.index(state_len))
-        .prop_map(move |index| *st.read().iter().nth(index).unwrap());
+        .prop_map(move |index| *st.read().values().iter().nth(index).unwrap());
 
     // Convert the value based on the parameter type
     match param {
@@ -131,6 +131,7 @@ pub fn fuzz_param_from_state(param: &ParamType, arc_state: EvmFuzzState) -> Boxe
 mod tests {
     use crate::fuzz::strategies::{build_initial_state, fuzz_calldata, fuzz_calldata_from_state};
     use ethers::abi::HumanReadableParser;
+    use foundry_config::FuzzDictionaryConfig;
     use revm::db::{CacheDB, EmptyDB};
 
     #[test]
@@ -139,7 +140,7 @@ mod tests {
         let func = HumanReadableParser::parse_function(f).unwrap();
 
         let db = CacheDB::new(EmptyDB());
-        let state = build_initial_state(&db, true, true);
+        let state = build_initial_state(&db, &FuzzDictionaryConfig::default());
 
         let strat = proptest::strategy::Union::new_weighted(vec![
             (60, fuzz_calldata(func.clone())),