perf: compute literals only once (#12716)

* perf: compute literals only once

* init on setup

* fix: deadlock

* fix: Default for FuzzDictionary

* chore: clippy

* fix: defaults

* fix: usize::min to allow disabling with 0. still unintuitive

* chore: fast path max > 0

* chore: break if break

* chore: consider only fuzz dictionary config

Author: DaniPopes

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

@@ -99,19 +99,19 @@ impl FuzzedExecutor {
         &mut self,
         func: &Function,
         fuzz_fixtures: &FuzzFixtures,
-        deployed_libs: &[Address],
+        state: EvmFuzzState,
         address: Address,
         rd: &RevertDecoder,
         progress: Option<&ProgressBar>,
         early_exit: &EarlyExit,
     ) -> Result<FuzzTestResult> {
+        let state = &state;
         // Stores the fuzz test execution data.
         let mut test_data = FuzzTestData::default();
-        let state = self.build_fuzz_state(deployed_libs);
         let dictionary_weight = self.config.dictionary.dictionary_weight.min(100);
         let strategy = proptest::prop_oneof![
             100 - dictionary_weight => fuzz_calldata(func.clone(), fuzz_fixtures),
-            dictionary_weight => fuzz_calldata_from_state(func.clone(), &state),
+            dictionary_weight => fuzz_calldata_from_state(func.clone(), state),
         ]
         .prop_map(move |calldata| BasicTxDetails {
             warp: None,
@@ -173,7 +173,7 @@ impl FuzzedExecutor {
 
                 test_data.runs += 1;
 
-                match corpus_manager.new_input(&mut self.runner, &state, func) {
+                match corpus_manager.new_input(&mut self.runner, state, func) {
                     Ok(input) => input,
                     Err(err) => {
                         test_data.failure = Some(TestCaseError::fail(format!(
@@ -382,27 +382,4 @@ impl FuzzedExecutor {
             }))
         }
     }
-
-    /// Stores fuzz state for use with [fuzz_calldata_from_state]
-    pub fn build_fuzz_state(&self, deployed_libs: &[Address]) -> EvmFuzzState {
-        let inspector = self.executor.inspector();
-
-        if let Some(fork_db) = self.executor.backend().active_fork_db() {
-            EvmFuzzState::new(
-                fork_db,
-                self.config.dictionary,
-                deployed_libs,
-                inspector.analysis.as_ref(),
-                inspector.paths_config(),
-            )
-        } else {
-            EvmFuzzState::new(
-                self.executor.backend().mem_db(),
-                self.config.dictionary,
-                deployed_libs,
-                inspector.analysis.as_ref(),
-                inspector.paths_config(),
-            )
-        }
-    }
 }

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

@@ -331,7 +331,7 @@ impl<'a> InvariantExecutor<'a> {
         &mut self,
         invariant_contract: InvariantContract<'_>,
         fuzz_fixtures: &FuzzFixtures,
-        deployed_libs: &[Address],
+        fuzz_state: EvmFuzzState,
         progress: Option<&ProgressBar>,
         early_exit: &EarlyExit,
     ) -> Result<InvariantFuzzTestResult> {
@@ -341,7 +341,7 @@ impl<'a> InvariantExecutor<'a> {
         }
 
         let (mut invariant_test, mut corpus_manager) =
-            self.prepare_test(&invariant_contract, fuzz_fixtures, deployed_libs)?;
+            self.prepare_test(&invariant_contract, fuzz_fixtures, fuzz_state)?;
 
         // Start timer for this invariant test.
         let mut runs = 0;
@@ -552,23 +552,13 @@ impl<'a> InvariantExecutor<'a> {
         &mut self,
         invariant_contract: &InvariantContract<'_>,
         fuzz_fixtures: &FuzzFixtures,
-        deployed_libs: &[Address],
+        fuzz_state: EvmFuzzState,
     ) -> Result<(InvariantTest, CorpusManager)> {
         // Finds out the chosen deployed contracts and/or senders.
         self.select_contract_artifacts(invariant_contract.address)?;
         let (targeted_senders, targeted_contracts) =
             self.select_contracts_and_senders(invariant_contract.address)?;
 
-        // Stores fuzz state for use with [fuzz_calldata_from_state].
-        let inspector = self.executor.inspector();
-        let fuzz_state = EvmFuzzState::new(
-            self.executor.backend().mem_db(),
-            self.config.dictionary,
-            deployed_libs,
-            inspector.analysis.as_ref(),
-            inspector.paths_config(),
-        );
-
         // Creates the invariant strategy.
         let strategy = invariant_strat(
             fuzz_state.clone(),

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

@@ -9,17 +9,20 @@ use solar::{
     ast::{self, Visit},
     interface::source_map::FileName,
 };
-use std::{ops::ControlFlow, sync::OnceLock};
+use std::{
+    ops::ControlFlow,
+    sync::{Arc, OnceLock},
+};
 
-#[derive(Debug, Default)]
+#[derive(Clone, Debug)]
 pub struct LiteralsDictionary {
-    /// Data required for initialization, captured from `EvmFuzzState::new`.
-    analysis: Option<Analysis>,
-    paths_config: Option<ProjectPathsConfig>,
-    max_values: usize,
+    maps: Arc<OnceLock<LiteralMaps>>,
+}
 
-    /// Lazy initialized literal maps.
-    maps: OnceLock<LiteralMaps>,
+impl Default for LiteralsDictionary {
+    fn default() -> Self {
+        Self::new(None, None, usize::MAX)
+    }
 }
 
 impl LiteralsDictionary {
@@ -28,42 +31,42 @@ impl LiteralsDictionary {
         paths_config: Option<ProjectPathsConfig>,
         max_values: usize,
     ) -> Self {
-        Self { analysis, paths_config, max_values, maps: OnceLock::default() }
+        let maps = Arc::new(OnceLock::<LiteralMaps>::new());
+        if let Some(analysis) = analysis
+            && max_values > 0
+        {
+            let maps = maps.clone();
+            // This can't be done in a rayon task (including inside of `get`) because it can cause a
+            // deadlock, since internally `solar` also uses rayon.
+            let _ = std::thread::Builder::new().name("literal-collector".into()).spawn(move || {
+                let _ = maps.get_or_init(|| {
+                    let literals =
+                        LiteralsCollector::process(&analysis, paths_config.as_ref(), max_values);
+                    debug!(
+                        words = literals.words.values().map(|set| set.len()).sum::<usize>(),
+                        strings = literals.strings.len(),
+                        bytes = literals.bytes.len(),
+                        "collected source code literals for fuzz dictionary"
+                    );
+                    literals
+                });
+            });
+        } else {
+            maps.set(Default::default()).unwrap();
+        }
+        Self { maps }
     }
 
-    /// Returns a reference to the `LiteralMaps`, initializing them on the first call.
+    /// Returns a reference to the `LiteralMaps`.
     pub fn get(&self) -> &LiteralMaps {
-        self.maps.get_or_init(|| {
-            if let Some(analysis) = &self.analysis {
-                let literals = LiteralsCollector::process(
-                    analysis,
-                    self.paths_config.as_ref(),
-                    self.max_values,
-                );
-                trace!(
-                    words = literals.words.values().map(|set| set.len()).sum::<usize>(),
-                    strings = literals.strings.len(),
-                    bytes = literals.bytes.len(),
-                    "collected source code literals for fuzz dictionary"
-                );
-                literals
-            } else {
-                LiteralMaps::default()
-            }
-        })
-    }
-
-    /// Takes ownership of the dictionary words, leaving an empty map in their place.
-    /// Ensures the map is initialized before taking its contents.
-    pub fn take_words(&mut self) -> HashMap<DynSolType, B256IndexSet> {
-        let _ = self.get();
-        self.maps.get_mut().map(|m| std::mem::take(&mut m.words)).unwrap_or_default()
+        self.maps.wait()
     }
 
-    #[cfg(test)]
     /// Test-only helper to seed the dictionary with literal values.
+    #[cfg(test)]
     pub(crate) fn set(&mut self, map: super::LiteralMaps) {
-        let _ = self.maps.set(map);
+        self.maps = Arc::new(OnceLock::new());
+        self.maps.set(map).unwrap();
     }
 }
 
@@ -102,8 +105,10 @@ impl LiteralsCollector {
                     continue;
                 }
 
-                if let Some(ref ast) = source.ast {
-                    let _ = literals_collector.visit_source_unit(ast);
+                if let Some(ast) = &source.ast
+                    && literals_collector.visit_source_unit(ast).is_break()
+                {
+                    break;
                 }
             }
 

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

@@ -20,4 +20,4 @@ mod mutators;
 pub use mutators::BoundMutator;
 
 mod literals;
-pub use literals::{LiteralMaps, LiteralsCollector};
+pub use literals::{LiteralMaps, LiteralsCollector, LiteralsDictionary};

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

@@ -433,16 +433,13 @@ mod tests {
     };
     use alloy_primitives::B256;
     use foundry_common::abi::get_func;
-    use foundry_config::FuzzDictionaryConfig;
-    use revm::database::{CacheDB, EmptyDB};
     use std::collections::HashSet;
 
     #[test]
     fn can_fuzz_array() {
         let f = "testArray(uint64[2] calldata values)";
         let func = get_func(f).unwrap();
-        let db = CacheDB::new(EmptyDB::default());
-        let state = EvmFuzzState::new(&db, FuzzDictionaryConfig::default(), &[], None, None);
+        let state = EvmFuzzState::test();
         let strategy = proptest::prop_oneof![
             60 => fuzz_calldata(func.clone(), &FuzzFixtures::default()),
             40 => fuzz_calldata_from_state(func, &state),
@@ -467,8 +464,7 @@ mod tests {
         literals.words.entry(DynSolType::FixedBytes(32)).or_default().insert(keccak256("hello"));
         literals.words.entry(DynSolType::FixedBytes(32)).or_default().insert(keccak256("world"));
 
-        let db = CacheDB::new(EmptyDB::default());
-        let state = EvmFuzzState::new(&db, FuzzDictionaryConfig::default(), &[], None, None);
+        let state = EvmFuzzState::test();
         state.seed_literals(literals);
 
         let cfg = proptest::test_runner::Config { failure_persistence: None, ..Default::default() };

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

@@ -8,10 +8,9 @@ use alloy_primitives::{
     map::{AddressIndexSet, AddressMap, B256IndexSet, HashMap, IndexSet},
 };
 use foundry_common::{
-    compile::Analysis, ignore_metadata_hash, mapping_slots::MappingSlots,
-    slot_identifier::SlotIdentifier,
+    ignore_metadata_hash, mapping_slots::MappingSlots, slot_identifier::SlotIdentifier,
 };
-use foundry_compilers::{ProjectPathsConfig, artifacts::StorageLayout};
+use foundry_compilers::artifacts::StorageLayout;
 use foundry_config::FuzzDictionaryConfig;
 use foundry_evm_core::{bytecode::InstIter, utils::StateChangeset};
 use parking_lot::{RawRwLock, RwLock, lock_api::RwLockReadGuard};
@@ -43,12 +42,21 @@ pub struct EvmFuzzState {
 }
 
 impl EvmFuzzState {
+    #[cfg(test)]
+    pub(crate) fn test() -> Self {
+        Self::new(
+            &[],
+            &CacheDB::<revm::database::EmptyDB>::default(),
+            FuzzDictionaryConfig::default(),
+            None,
+        )
+    }
+
     pub fn new<DB: DatabaseRef>(
+        deployed_libs: &[Address],
         db: &CacheDB<DB>,
         config: FuzzDictionaryConfig,
-        deployed_libs: &[Address],
-        analysis: Option<&Analysis>,
-        paths_config: Option<&ProjectPathsConfig>,
+        literals: Option<&LiteralsDictionary>,
     ) -> Self {
         // Sort accounts to ensure deterministic dictionary generation from the same setUp state.
         let mut accs = db.cache.accounts.iter().collect::<Vec<_>>();
@@ -57,11 +65,9 @@ impl EvmFuzzState {
         // Create fuzz dictionary and insert values from db state.
         let mut dictionary = FuzzDictionary::new(config);
         dictionary.insert_db_values(accs);
-        dictionary.literal_values = LiteralsDictionary::new(
-            analysis.cloned(),
-            paths_config.cloned(),
-            config.max_fuzz_dictionary_literals,
-        );
+        if let Some(literals) = literals {
+            dictionary.literal_values = literals.clone();
+        }
 
         Self {
             inner: Arc::new(RwLock::new(dictionary)),
@@ -126,16 +132,15 @@ impl EvmFuzzState {
         self.inner.read().log_stats();
     }
 
-    #[cfg(test)]
     /// Test-only helper to seed the dictionary with literal values.
+    #[cfg(test)]
     pub(crate) fn seed_literals(&self, map: super::LiteralMaps) {
         self.inner.write().seed_literals(map);
     }
 }
 
 // We're using `IndexSet` to have a stable element order when restoring persisted state, as well as
 // for performance when iterating over the sets.
-#[derive(Default)]
 pub struct FuzzDictionary {
     /// Collected state values.
     state_values: B256IndexSet,
@@ -173,9 +178,27 @@ impl fmt::Debug for FuzzDictionary {
     }
 }
 
+impl Default for FuzzDictionary {
+    fn default() -> Self {
+        Self::new(Default::default())
+    }
+}
+
 impl FuzzDictionary {
     pub fn new(config: FuzzDictionaryConfig) -> Self {
-        let mut dictionary = Self { config, samples_seeded: false, ..Default::default() };
+        let mut dictionary = Self {
+            config,
+            samples_seeded: false,
+
+            state_values: Default::default(),
+            addresses: Default::default(),
+            db_state_values: Default::default(),
+            db_addresses: Default::default(),
+            sample_values: Default::default(),
+            literal_values: Default::default(),
+            misses: Default::default(),
+            hits: Default::default(),
+        };
         dictionary.prefill();
         dictionary
     }
@@ -190,7 +213,8 @@ impl FuzzDictionary {
     #[cold]
     fn seed_samples(&mut self) {
         trace!("seeding `sample_values` from literal dictionary");
-        self.sample_values.extend(self.literal_values.take_words());
+        self.sample_values
+            .extend(self.literal_values.get().words.iter().map(|(k, v)| (k.clone(), v.clone())));
         self.samples_seeded = true;
     }