Changes after review - comments, wrap Arc as CalldataFuzzDictionary.i…

…nner, code cleanup

crates/evm/evm/src/executors/invariant/mod.rs

@@ -33,7 +33,7 @@ use std::{cell::RefCell, collections::BTreeMap, sync::Arc};
 
 mod error;
 pub use error::{InvariantFailures, InvariantFuzzError, InvariantFuzzTestResult};
-use foundry_evm_fuzz::strategies::{CalldataFuzzDictionary, CalldataFuzzDictionaryConfig};
+use foundry_evm_fuzz::strategies::CalldataFuzzDictionary;
 
 mod funcs;
 pub use funcs::{assert_invariants, replay_run};
@@ -251,7 +251,7 @@ impl<'a> InvariantExecutor<'a> {
             Ok(())
         });
 
-        trace!(target: "forge::test::invariant::calldata_address_fuzz_dictionary", "{:?}", calldata_fuzz_dictionary.clone().addresses);
+        trace!(target: "forge::test::invariant::calldata_address_fuzz_dictionary", "{:?}", calldata_fuzz_dictionary.inner.addresses);
         trace!(target: "forge::test::invariant::dictionary", "{:?}", fuzz_state.read().values().iter().map(hex::encode).collect::<Vec<_>>());
 
         let (reverts, error) = failures.into_inner().into_inner();
@@ -290,9 +290,8 @@ impl<'a> InvariantExecutor<'a> {
         let targeted_contracts: FuzzRunIdentifiedContracts =
             Arc::new(Mutex::new(targeted_contracts));
 
-        let calldata_fuzz_config: CalldataFuzzDictionary = Arc::new(
-            CalldataFuzzDictionaryConfig::new(&self.config.dictionary, fuzz_state.clone()),
-        );
+        let calldata_fuzz_config =
+            CalldataFuzzDictionary::new(&self.config.dictionary, fuzz_state.clone());
 
         // Creates the invariant strategy.
         let strat = invariant_strat(

crates/evm/fuzz/src/strategies/calldata.rs

@@ -7,7 +7,17 @@ use hashbrown::HashSet;
 use proptest::prelude::{BoxedStrategy, Strategy};
 use std::{fmt, sync::Arc};
 
-pub type CalldataFuzzDictionary = Arc<CalldataFuzzDictionaryConfig>;
+/// Clonable wrapper around [CalldataFuzzDictionary].
+#[derive(Debug, Clone)]
+pub struct CalldataFuzzDictionary {
+    pub inner: Arc<CalldataFuzzDictionaryConfig>,
+}
+
+impl CalldataFuzzDictionary {
+    pub fn new(config: &FuzzDictionaryConfig, state: EvmFuzzState) -> Self {
+        Self { inner: Arc::new(CalldataFuzzDictionaryConfig::new(config, state)) }
+    }
+}
 
 #[derive(Clone)]
 pub struct CalldataFuzzDictionaryConfig {
@@ -21,7 +31,16 @@ impl fmt::Debug for CalldataFuzzDictionaryConfig {
     }
 }
 
+/// Represents custom configuration for invariant fuzzed calldata strategies.
+///
+/// At the moment only the dictionary of addresses to be used for a fuzzed `function(address)` can
+/// be configured, but support for other types can be added.
 impl CalldataFuzzDictionaryConfig {
+    /// Creates config with the set of addresses that can be used for fuzzing invariant calldata (if
+    /// `max_calldata_fuzz_dictionary_addresses` configured).
+    /// The set of addresses contains a number of `max_calldata_fuzz_dictionary_addresses` random
+    /// addresses plus all addresses that already had their PUSH bytes collected (retrieved from
+    /// `EvmFuzzState`, if `include_push_bytes` config enabled).
     pub fn new(config: &FuzzDictionaryConfig, state: EvmFuzzState) -> Self {
         let mut addresses: HashSet<Address> = HashSet::new();
         let dict_size = config.max_calldata_fuzz_dictionary_addresses;
@@ -34,8 +53,7 @@ impl CalldataFuzzDictionaryConfig {
                 addresses.insert(Address::random());
             }
 
-            // add any state address calldata fuzz dictionary, in addition to random generated
-            // addresses
+            // Add all addresses that already had their PUSH bytes collected.
             let mut state = state.write();
             addresses.extend(state.addresses());
         }
@@ -50,6 +68,8 @@ pub fn fuzz_calldata(func: Function) -> BoxedStrategy<Bytes> {
     fuzz_calldata_with_config(func, None)
 }
 
+/// Given a function, it returns a strategy which generates valid calldata
+/// for that function's input types, following custom configuration rules.
 pub fn fuzz_calldata_with_config(
     func: Function,
     config: Option<CalldataFuzzDictionary>,

crates/evm/fuzz/src/strategies/param.rs

@@ -18,22 +18,24 @@ pub fn fuzz_param(
     let param = param.to_owned();
     match param {
         DynSolType::Address => {
-            let cfg = config.clone();
-            if cfg.is_some() && !cfg.unwrap().addresses.is_empty() {
-                let dict_len = config.clone().unwrap().addresses.len();
-                any::<prop::sample::Index>()
-                    .prop_map(move |index| index.index(dict_len))
-                    .prop_map(move |index| {
-                        DynSolValue::Address(
-                            config.clone().unwrap().addresses.get(index).cloned().unwrap(),
-                        )
-                    })
-                    .boxed()
-            } else {
-                any::<[u8; 32]>()
-                    .prop_map(|x| DynSolValue::Address(Address::from_word(x.into())))
-                    .boxed()
+            if config.is_some() {
+                let fuzz_config = config.unwrap().inner;
+                let address_dict_len = fuzz_config.addresses.len();
+                if address_dict_len > 0 {
+                    // Create strategy to return random address from configured dictionary.
+                    return any::<prop::sample::Index>()
+                        .prop_map(move |index| index.index(address_dict_len))
+                        .prop_map(move |index| {
+                            DynSolValue::Address(fuzz_config.addresses.get(index).cloned().unwrap())
+                        })
+                        .boxed()
+                }
             }
+
+            // If no config for addresses dictionary then create unbounded addresses strategy.
+            any::<[u8; 32]>()
+                .prop_map(|x| DynSolValue::Address(Address::from_word(x.into())))
+                .boxed()
         }
         DynSolType::Int(n) => {
             let strat = super::IntStrategy::new(n, vec![]);