foundry_evm_fuzz/strategies/
state.rs

1use crate::{BasicTxDetails, invariant::FuzzRunIdentifiedContracts};
2use alloy_dyn_abi::{DynSolType, DynSolValue, EventExt, FunctionExt};
3use alloy_json_abi::{Function, JsonAbi};
4use alloy_primitives::{
5    Address, B256, Bytes, Log, U256,
6    map::{AddressIndexSet, B256IndexSet, HashMap},
7};
8use foundry_common::ignore_metadata_hash;
9use foundry_compilers::artifacts::StorageLayout;
10use foundry_config::FuzzDictionaryConfig;
11use foundry_evm_core::utils::StateChangeset;
12use parking_lot::{RawRwLock, RwLock, lock_api::RwLockReadGuard};
13use revm::{
14    bytecode::opcode,
15    database::{CacheDB, DatabaseRef, DbAccount},
16    state::AccountInfo,
17};
18use std::{collections::BTreeMap, fmt, sync::Arc};
19
20/// The maximum number of bytes we will look at in bytecodes to find push bytes (24 KiB).
21///
22/// This is to limit the performance impact of fuzz tests that might deploy arbitrarily sized
23/// bytecode (as is the case with Solmate).
24const PUSH_BYTE_ANALYSIS_LIMIT: usize = 24 * 1024;
25
26/// A set of arbitrary 32 byte data from the VM used to generate values for the strategy.
27///
28/// Wrapped in a shareable container.
29#[derive(Clone, Debug)]
30pub struct EvmFuzzState {
31    inner: Arc<RwLock<FuzzDictionary>>,
32    /// Addresses of external libraries deployed in test setup, excluded from fuzz test inputs.
33    pub deployed_libs: Vec<Address>,
34}
35
36impl EvmFuzzState {
37    pub fn new<DB: DatabaseRef>(
38        db: &CacheDB<DB>,
39        config: FuzzDictionaryConfig,
40        deployed_libs: &[Address],
41    ) -> Self {
42        // Sort accounts to ensure deterministic dictionary generation from the same setUp state.
43        let mut accs = db.cache.accounts.iter().collect::<Vec<_>>();
44        accs.sort_by_key(|(address, _)| *address);
45
46        // Create fuzz dictionary and insert values from db state.
47        let mut dictionary = FuzzDictionary::new(config);
48        dictionary.insert_db_values(accs);
49        Self { inner: Arc::new(RwLock::new(dictionary)), deployed_libs: deployed_libs.to_vec() }
50    }
51
52    pub fn collect_values(&self, values: impl IntoIterator<Item = B256>) {
53        let mut dict = self.inner.write();
54        for value in values {
55            dict.insert_value(value);
56        }
57    }
58
59    /// Collects state changes from a [StateChangeset] and logs into an [EvmFuzzState] according to
60    /// the given [FuzzDictionaryConfig].
61    pub fn collect_values_from_call(
62        &self,
63        fuzzed_contracts: &FuzzRunIdentifiedContracts,
64        tx: &BasicTxDetails,
65        result: &Bytes,
66        logs: &[Log],
67        state_changeset: &StateChangeset,
68        run_depth: u32,
69    ) {
70        let mut dict = self.inner.write();
71        {
72            let targets = fuzzed_contracts.targets.lock();
73            let (target_abi, target_function) = targets.fuzzed_artifacts(tx);
74            dict.insert_logs_values(target_abi, logs, run_depth);
75            dict.insert_result_values(target_function, result, run_depth);
76            // Get storage layouts for contracts in the state changeset
77            let storage_layouts = targets.get_storage_layouts();
78            dict.insert_new_state_values(state_changeset, &storage_layouts);
79        }
80    }
81
82    /// Removes all newly added entries from the dictionary.
83    ///
84    /// Should be called between fuzz/invariant runs to avoid accumulating data derived from fuzz
85    /// inputs.
86    pub fn revert(&self) {
87        self.inner.write().revert();
88    }
89
90    pub fn dictionary_read(&self) -> RwLockReadGuard<'_, RawRwLock, FuzzDictionary> {
91        self.inner.read()
92    }
93
94    /// Logs stats about the current state.
95    pub fn log_stats(&self) {
96        self.inner.read().log_stats();
97    }
98}
99
100// We're using `IndexSet` to have a stable element order when restoring persisted state, as well as
101// for performance when iterating over the sets.
102#[derive(Default)]
103pub struct FuzzDictionary {
104    /// Collected state values.
105    state_values: B256IndexSet,
106    /// Addresses that already had their PUSH bytes collected.
107    addresses: AddressIndexSet,
108    /// Configuration for the dictionary.
109    config: FuzzDictionaryConfig,
110    /// Number of state values initially collected from db.
111    /// Used to revert new collected values at the end of each run.
112    db_state_values: usize,
113    /// Number of address values initially collected from db.
114    /// Used to revert new collected addresses at the end of each run.
115    db_addresses: usize,
116    /// Sample typed values that are collected from call result and used across invariant runs.
117    sample_values: HashMap<DynSolType, B256IndexSet>,
118
119    misses: usize,
120    hits: usize,
121}
122
123impl fmt::Debug for FuzzDictionary {
124    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
125        f.debug_struct("FuzzDictionary")
126            .field("state_values", &self.state_values.len())
127            .field("addresses", &self.addresses)
128            .finish()
129    }
130}
131
132impl FuzzDictionary {
133    pub fn new(config: FuzzDictionaryConfig) -> Self {
134        let mut dictionary = Self { config, ..Default::default() };
135        dictionary.prefill();
136        dictionary
137    }
138
139    /// Insert common values into the dictionary at initialization.
140    fn prefill(&mut self) {
141        self.insert_value(B256::ZERO);
142    }
143
144    /// Insert values from initial db state into fuzz dictionary.
145    /// These values are persisted across invariant runs.
146    fn insert_db_values(&mut self, db_state: Vec<(&Address, &DbAccount)>) {
147        for (address, account) in db_state {
148            // Insert basic account information
149            self.insert_value(address.into_word());
150            // Insert push bytes
151            self.insert_push_bytes_values(address, &account.info);
152            // Insert storage values.
153            if self.config.include_storage {
154                // Sort storage values before inserting to ensure deterministic dictionary.
155                let values = account.storage.iter().collect::<BTreeMap<_, _>>();
156                for (slot, value) in values {
157                    self.insert_storage_value(slot, value, None);
158                }
159            }
160        }
161
162        // We need at least some state data if DB is empty,
163        // otherwise we can't select random data for state fuzzing.
164        if self.values().is_empty() {
165            // Prefill with a random address.
166            self.insert_value(Address::random().into_word());
167        }
168
169        // Record number of values and addresses inserted from db to be used for reverting at the
170        // end of each run.
171        self.db_state_values = self.state_values.len();
172        self.db_addresses = self.addresses.len();
173    }
174
175    /// Insert values collected from call result into fuzz dictionary.
176    fn insert_result_values(
177        &mut self,
178        function: Option<&Function>,
179        result: &Bytes,
180        run_depth: u32,
181    ) {
182        if let Some(function) = function
183            && !function.outputs.is_empty()
184        {
185            // Decode result and collect samples to be used in subsequent fuzz runs.
186            if let Ok(decoded_result) = function.abi_decode_output(result) {
187                self.insert_sample_values(decoded_result, run_depth);
188            }
189        }
190    }
191
192    /// Insert values from call log topics and data into fuzz dictionary.
193    fn insert_logs_values(&mut self, abi: Option<&JsonAbi>, logs: &[Log], run_depth: u32) {
194        let mut samples = Vec::new();
195        // Decode logs with known events and collect samples from indexed fields and event body.
196        for log in logs {
197            let mut log_decoded = false;
198            // Try to decode log with events from contract abi.
199            if let Some(abi) = abi {
200                for event in abi.events() {
201                    if let Ok(decoded_event) = event.decode_log(log) {
202                        samples.extend(decoded_event.indexed);
203                        samples.extend(decoded_event.body);
204                        log_decoded = true;
205                        break;
206                    }
207                }
208            }
209
210            // If we weren't able to decode event then we insert raw data in fuzz dictionary.
211            if !log_decoded {
212                for &topic in log.topics() {
213                    self.insert_value(topic);
214                }
215                let chunks = log.data.data.chunks_exact(32);
216                let rem = chunks.remainder();
217                for chunk in chunks {
218                    self.insert_value(chunk.try_into().unwrap());
219                }
220                if !rem.is_empty() {
221                    self.insert_value(B256::right_padding_from(rem));
222                }
223            }
224        }
225
226        // Insert samples collected from current call in fuzz dictionary.
227        self.insert_sample_values(samples, run_depth);
228    }
229
230    /// Insert values from call state changeset into fuzz dictionary.
231    /// These values are removed at the end of current run.
232    fn insert_new_state_values(
233        &mut self,
234        state_changeset: &StateChangeset,
235        storage_layouts: &HashMap<Address, Arc<StorageLayout>>,
236    ) {
237        for (address, account) in state_changeset {
238            // Insert basic account information.
239            self.insert_value(address.into_word());
240            // Insert push bytes.
241            self.insert_push_bytes_values(address, &account.info);
242            // Insert storage values.
243            if self.config.include_storage {
244                let storage_layout = storage_layouts.get(address).map(|arc| arc.as_ref());
245                for (slot, value) in &account.storage {
246                    self.insert_storage_value(slot, &value.present_value, storage_layout);
247                }
248            }
249        }
250    }
251
252    /// Insert values from push bytes into fuzz dictionary.
253    /// Values are collected only once for a given address.
254    /// If values are newly collected then they are removed at the end of current run.
255    fn insert_push_bytes_values(&mut self, address: &Address, account_info: &AccountInfo) {
256        if self.config.include_push_bytes && !self.addresses.contains(address) {
257            // Insert push bytes
258            if let Some(code) = &account_info.code {
259                self.insert_address(*address);
260                self.collect_push_bytes(ignore_metadata_hash(code.original_byte_slice()));
261            }
262        }
263    }
264
265    fn collect_push_bytes(&mut self, code: &[u8]) {
266        let mut i = 0;
267        let len = code.len().min(PUSH_BYTE_ANALYSIS_LIMIT);
268        while i < len {
269            let op = code[i];
270            if (opcode::PUSH1..=opcode::PUSH32).contains(&op) {
271                let push_size = (op - opcode::PUSH1 + 1) as usize;
272                let push_start = i + 1;
273                let push_end = push_start + push_size;
274                // As a precaution, if a fuzz test deploys malformed bytecode (such as using
275                // `CREATE2`) this will terminate the loop early.
276                if push_start > code.len() || push_end > code.len() {
277                    break;
278                }
279
280                let push_value = U256::try_from_be_slice(&code[push_start..push_end]).unwrap();
281                if push_value != U256::ZERO {
282                    // Never add 0 to the dictionary as it's always present.
283                    self.insert_value(push_value.into());
284
285                    // Also add the value below and above the push value to the dictionary.
286                    self.insert_value((push_value - U256::from(1)).into());
287
288                    if push_value != U256::MAX {
289                        self.insert_value((push_value + U256::from(1)).into());
290                    }
291                }
292
293                i += push_size;
294            }
295            i += 1;
296        }
297    }
298
299    /// Resolves storage types from a storage layout for a given slot and all mapping types.
300    /// Returns a tuple of (slot_type, mapping_types) where slot_type is the specific type
301    /// for the storage slot and mapping_types are all mapping value types found in the layout.
302    fn resolve_storage_types(
303        &self,
304        storage_layout: Option<&StorageLayout>,
305        storage_slot: &U256,
306    ) -> (Option<DynSolType>, Vec<DynSolType>) {
307        let Some(layout) = storage_layout else {
308            return (None, Vec::new());
309        };
310
311        // Try to determine the type of this specific storage slot
312        let slot_type =
313            layout.storage.iter().find(|s| s.slot == storage_slot.to_string()).and_then(
314                |storage| {
315                    layout
316                        .types
317                        .get(&storage.storage_type)
318                        .and_then(|t| DynSolType::parse(&t.label).ok())
319                },
320            );
321
322        // Collect all mapping value types from the layout
323        let mapping_types = layout
324            .types
325            .values()
326            .filter_map(|type_info| {
327                if type_info.encoding == "mapping"
328                    && let Some(t_value) = type_info.value.as_ref()
329                    && let Some(mapping_value) = t_value.strip_prefix("t_")
330                {
331                    DynSolType::parse(mapping_value).ok()
332                } else {
333                    None
334                }
335            })
336            .collect();
337
338        (slot_type, mapping_types)
339    }
340
341    /// Insert values from single storage slot and storage value into fuzz dictionary.
342    /// If storage values are newly collected then they are removed at the end of current run.
343    fn insert_storage_value(
344        &mut self,
345        storage_slot: &U256,
346        storage_value: &U256,
347        storage_layout: Option<&StorageLayout>,
348    ) {
349        // Always insert the slot itself
350        self.insert_value(B256::from(*storage_slot));
351
352        let (slot_type, mapping_types) = self.resolve_storage_types(storage_layout, storage_slot);
353
354        if let Some(sol_type) = slot_type {
355            self.insert_decoded_storage_value(sol_type, storage_value);
356        } else if !mapping_types.is_empty() {
357            self.insert_mapping_storage_values(mapping_types, storage_value);
358        } else {
359            // No type information available, insert as raw values (old behavior)
360            self.insert_value(B256::from(*storage_value));
361            // also add the value below and above the storage value to the dictionary.
362            if *storage_value != U256::ZERO {
363                let below_value = storage_value - U256::from(1);
364                self.insert_value(B256::from(below_value));
365            }
366            if *storage_value != U256::MAX {
367                let above_value = storage_value + U256::from(1);
368                self.insert_value(B256::from(above_value));
369            }
370        }
371    }
372
373    /// Insert decoded storage values into the fuzz dictionary.
374    /// Only simple static type values are inserted as sample values.
375    /// Complex types (dynamic arrays, structs) are inserted as raw values
376    fn insert_decoded_storage_value(&mut self, sol_type: DynSolType, storage_value: &U256) {
377        // Only insert values for types that can be represented as a single word
378        match &sol_type {
379            DynSolType::Address
380            | DynSolType::Uint(_)
381            | DynSolType::Int(_)
382            | DynSolType::Bool
383            | DynSolType::FixedBytes(_)
384            | DynSolType::Bytes => {
385                // Insert as a typed sample value
386                self.sample_values.entry(sol_type).or_default().insert(B256::from(*storage_value));
387            }
388            _ => {
389                // For complex types (arrays, mappings, structs), insert as raw value
390                self.insert_value(B256::from(*storage_value));
391            }
392        }
393    }
394
395    /// Insert storage values of mapping value types as sample values in the fuzz dictionary.
396    ///
397    /// ```solidity
398    /// mapping(uint256 => address) public myMapping;
399    /// // `address` is the mapping value type here.
400    /// // `uint256` is the mapping key type.
401    /// ```
402    ///
403    /// A storage value is inserted if and only if it can be decoded into one of the mapping
404    /// [`DynSolType`] value types found in the [`StorageLayout`].
405    ///
406    /// If decoding fails, the value is inserted as a raw value.
407    fn insert_mapping_storage_values(
408        &mut self,
409        mapping_types: Vec<DynSolType>,
410        storage_value: &U256,
411    ) {
412        for sol_type in mapping_types {
413            match sol_type.abi_decode(storage_value.as_le_slice()) {
414                Ok(_) => {
415                    self.sample_values
416                        .entry(sol_type)
417                        .or_default()
418                        .insert(B256::from(*storage_value));
419                }
420                Err(_) => {
421                    // If decoding fails, insert as raw value
422                    self.insert_value(B256::from(*storage_value));
423                }
424            }
425        }
426    }
427
428    /// Insert address into fuzz dictionary.
429    /// If address is newly collected then it is removed by index at the end of current run.
430    fn insert_address(&mut self, address: Address) {
431        if self.addresses.len() < self.config.max_fuzz_dictionary_addresses {
432            self.addresses.insert(address);
433        }
434    }
435
436    /// Insert raw value into fuzz dictionary.
437    /// If value is newly collected then it is removed by index at the end of current run.
438    fn insert_value(&mut self, value: B256) {
439        if self.state_values.len() < self.config.max_fuzz_dictionary_values {
440            let new_value = self.state_values.insert(value);
441            let counter = if new_value { &mut self.misses } else { &mut self.hits };
442            *counter += 1;
443        }
444    }
445
446    /// Insert sample values that are reused across multiple runs.
447    /// The number of samples is limited to invariant run depth.
448    /// If collected samples limit is reached then values are inserted as regular values.
449    pub fn insert_sample_values(
450        &mut self,
451        sample_values: impl IntoIterator<Item = DynSolValue>,
452        limit: u32,
453    ) {
454        for sample in sample_values {
455            if let (Some(sample_type), Some(sample_value)) = (sample.as_type(), sample.as_word()) {
456                if let Some(values) = self.sample_values.get_mut(&sample_type) {
457                    if values.len() < limit as usize {
458                        values.insert(sample_value);
459                    } else {
460                        // Insert as state value (will be removed at the end of the run).
461                        self.insert_value(sample_value);
462                    }
463                } else {
464                    self.sample_values.entry(sample_type).or_default().insert(sample_value);
465                }
466            }
467        }
468    }
469
470    pub fn values(&self) -> &B256IndexSet {
471        &self.state_values
472    }
473
474    pub fn len(&self) -> usize {
475        self.state_values.len()
476    }
477
478    pub fn is_empty(&self) -> bool {
479        self.state_values.is_empty()
480    }
481
482    #[inline]
483    pub fn samples(&self, param_type: &DynSolType) -> Option<&B256IndexSet> {
484        self.sample_values.get(param_type)
485    }
486
487    #[inline]
488    pub fn addresses(&self) -> &AddressIndexSet {
489        &self.addresses
490    }
491
492    /// Revert values and addresses collected during the run by truncating to initial db len.
493    pub fn revert(&mut self) {
494        self.state_values.truncate(self.db_state_values);
495        self.addresses.truncate(self.db_addresses);
496    }
497
498    pub fn log_stats(&self) {
499        trace!(
500            addresses.len = self.addresses.len(),
501            sample.len = self.sample_values.len(),
502            state.len = self.state_values.len(),
503            state.misses = self.misses,
504            state.hits = self.hits,
505            "FuzzDictionary stats",
506        );
507    }
508}