foundry_evm_fuzz/invariant/
mod.rs

1use alloy_json_abi::{Function, JsonAbi};
2use alloy_primitives::{Address, Selector, map::HashMap};
3use foundry_compilers::artifacts::StorageLayout;
4use itertools::Either;
5use parking_lot::Mutex;
6use std::{collections::BTreeMap, sync::Arc};
7
8mod call_override;
9pub use call_override::RandomCallGenerator;
10
11mod filters;
12use crate::BasicTxDetails;
13pub use filters::{ArtifactFilters, SenderFilters};
14use foundry_common::{ContractsByAddress, ContractsByArtifact};
15use foundry_evm_core::utils::{StateChangeset, get_function};
16
17/// Returns true if the function returns `int256`, indicating optimization mode.
18/// In optimization mode, the fuzzer maximizes the return value instead of checking invariants.
19pub fn is_optimization_invariant(func: &Function) -> bool {
20    func.outputs.len() == 1 && func.outputs[0].ty == "int256"
21}
22
23/// Contracts identified as targets during a fuzz run.
24///
25/// During execution, any newly created contract is added as target and used through the rest of
26/// the fuzz run if the collection is updatable (no `targetContract` specified in `setUp`).
27#[derive(Clone, Debug)]
28pub struct FuzzRunIdentifiedContracts {
29    /// Contracts identified as targets during a fuzz run.
30    pub targets: Arc<Mutex<TargetedContracts>>,
31    /// Whether target contracts are updatable or not.
32    pub is_updatable: bool,
33}
34
35impl FuzzRunIdentifiedContracts {
36    /// Creates a new `FuzzRunIdentifiedContracts` instance.
37    pub fn new(targets: TargetedContracts, is_updatable: bool) -> Self {
38        Self { targets: Arc::new(Mutex::new(targets)), is_updatable }
39    }
40
41    /// If targets are updatable, collect all contracts created during an invariant run (which
42    /// haven't been discovered yet).
43    pub fn collect_created_contracts(
44        &self,
45        state_changeset: &StateChangeset,
46        project_contracts: &ContractsByArtifact,
47        setup_contracts: &ContractsByAddress,
48        artifact_filters: &ArtifactFilters,
49        created_contracts: &mut Vec<Address>,
50    ) -> eyre::Result<()> {
51        if !self.is_updatable {
52            return Ok(());
53        }
54
55        let mut targets = self.targets.lock();
56        for (address, account) in state_changeset {
57            if setup_contracts.contains_key(address) {
58                continue;
59            }
60            if !account.is_touched() {
61                continue;
62            }
63            let Some(code) = &account.info.code else {
64                continue;
65            };
66            if code.is_empty() {
67                continue;
68            }
69            let Some((artifact, contract)) =
70                project_contracts.find_by_deployed_code(code.original_byte_slice())
71            else {
72                continue;
73            };
74            let Some(functions) =
75                artifact_filters.get_targeted_functions(artifact, &contract.abi)?
76            else {
77                continue;
78            };
79            created_contracts.push(*address);
80            let contract = TargetedContract {
81                identifier: artifact.name.clone(),
82                abi: contract.abi.clone(),
83                targeted_functions: functions,
84                excluded_functions: Vec::new(),
85                storage_layout: contract.storage_layout.as_ref().map(Arc::clone),
86            };
87            targets.insert(*address, contract);
88        }
89        Ok(())
90    }
91
92    /// Clears targeted contracts created during an invariant run.
93    pub fn clear_created_contracts(&self, created_contracts: Vec<Address>) {
94        if !created_contracts.is_empty() {
95            let mut targets = self.targets.lock();
96            for addr in &created_contracts {
97                targets.remove(addr);
98            }
99        }
100    }
101}
102
103/// A collection of contracts identified as targets for invariant testing.
104#[derive(Clone, Debug, Default)]
105pub struct TargetedContracts {
106    /// The inner map of targeted contracts.
107    pub inner: BTreeMap<Address, TargetedContract>,
108}
109
110impl TargetedContracts {
111    /// Returns a new `TargetedContracts` instance.
112    pub fn new() -> Self {
113        Self::default()
114    }
115
116    /// Returns fuzzed contract abi and fuzzed function from address and provided calldata.
117    ///
118    /// Used to decode return values and logs in order to add values into fuzz dictionary.
119    pub fn fuzzed_artifacts(&self, tx: &BasicTxDetails) -> (Option<&JsonAbi>, Option<&Function>) {
120        match self.inner.get(&tx.call_details.target) {
121            Some(c) => (
122                Some(&c.abi),
123                c.abi.functions().find(|f| f.selector() == tx.call_details.calldata[..4]),
124            ),
125            None => (None, None),
126        }
127    }
128
129    /// Returns flatten target contract address and functions to be fuzzed.
130    /// Includes contract targeted functions if specified, else all mutable contract functions.
131    pub fn fuzzed_functions(&self) -> impl Iterator<Item = (&Address, &Function)> {
132        self.inner
133            .iter()
134            .filter(|(_, c)| !c.abi.functions.is_empty())
135            .flat_map(|(contract, c)| c.abi_fuzzed_functions().map(move |f| (contract, f)))
136    }
137
138    /// Returns whether the given transaction can be replayed or not with known contracts.
139    pub fn can_replay(&self, tx: &BasicTxDetails) -> bool {
140        match self.inner.get(&tx.call_details.target) {
141            Some(c) => c.abi.functions().any(|f| f.selector() == tx.call_details.calldata[..4]),
142            None => false,
143        }
144    }
145
146    /// Identifies fuzzed contract and function based on given tx details and returns unique metric
147    /// key composed from contract identifier and function name.
148    pub fn fuzzed_metric_key(&self, tx: &BasicTxDetails) -> Option<String> {
149        self.inner.get(&tx.call_details.target).and_then(|contract| {
150            contract
151                .abi
152                .functions()
153                .find(|f| f.selector() == tx.call_details.calldata[..4])
154                .map(|function| format!("{}.{}", contract.identifier.clone(), function.name))
155        })
156    }
157
158    /// Returns a map of contract addresses to their storage layouts.
159    pub fn get_storage_layouts(&self) -> HashMap<Address, Arc<StorageLayout>> {
160        self.inner
161            .iter()
162            .filter_map(|(addr, c)| {
163                c.storage_layout.as_ref().map(|layout| (*addr, Arc::clone(layout)))
164            })
165            .collect()
166    }
167}
168
169impl std::ops::Deref for TargetedContracts {
170    type Target = BTreeMap<Address, TargetedContract>;
171
172    fn deref(&self) -> &Self::Target {
173        &self.inner
174    }
175}
176
177impl std::ops::DerefMut for TargetedContracts {
178    fn deref_mut(&mut self) -> &mut Self::Target {
179        &mut self.inner
180    }
181}
182
183/// A contract identified as target for invariant testing.
184#[derive(Clone, Debug)]
185pub struct TargetedContract {
186    /// The contract identifier. This is only used in error messages.
187    pub identifier: String,
188    /// The contract's ABI.
189    pub abi: JsonAbi,
190    /// The targeted functions of the contract.
191    pub targeted_functions: Vec<Function>,
192    /// The excluded functions of the contract.
193    pub excluded_functions: Vec<Function>,
194    /// The contract's storage layout, if available.
195    pub storage_layout: Option<Arc<StorageLayout>>,
196}
197
198impl TargetedContract {
199    /// Returns a new `TargetedContract` instance.
200    pub fn new(identifier: String, abi: JsonAbi) -> Self {
201        Self {
202            identifier,
203            abi,
204            targeted_functions: Vec::new(),
205            excluded_functions: Vec::new(),
206            storage_layout: None,
207        }
208    }
209
210    /// Determines contract storage layout from project contracts. Needs `storageLayout` to be
211    /// enabled as extra output in project configuration.
212    pub fn with_project_contracts(mut self, project_contracts: &ContractsByArtifact) -> Self {
213        if let Some((src, name)) = self.identifier.split_once(':')
214            && let Some((_, contract_data)) = project_contracts.iter().find(|(artifact, _)| {
215                artifact.name == name && artifact.source.as_path().ends_with(src)
216            })
217        {
218            self.storage_layout = contract_data.storage_layout.as_ref().map(Arc::clone);
219        }
220        self
221    }
222
223    /// Helper to retrieve functions to fuzz for specified abi.
224    /// Returns specified targeted functions if any, else mutable abi functions that are not
225    /// marked as excluded.
226    pub fn abi_fuzzed_functions(&self) -> impl Iterator<Item = &Function> {
227        if !self.targeted_functions.is_empty() {
228            Either::Left(self.targeted_functions.iter())
229        } else {
230            Either::Right(self.abi.functions().filter(|&func| {
231                !matches!(
232                    func.state_mutability,
233                    alloy_json_abi::StateMutability::Pure | alloy_json_abi::StateMutability::View
234                ) && !self.excluded_functions.contains(func)
235            }))
236        }
237    }
238
239    /// Returns the function for the given selector.
240    pub fn get_function(&self, selector: Selector) -> eyre::Result<&Function> {
241        get_function(&self.identifier, selector, &self.abi)
242    }
243
244    /// Adds the specified selectors to the targeted functions.
245    pub fn add_selectors(
246        &mut self,
247        selectors: impl IntoIterator<Item = Selector>,
248        should_exclude: bool,
249    ) -> eyre::Result<()> {
250        for selector in selectors {
251            if should_exclude {
252                self.excluded_functions.push(self.get_function(selector)?.clone());
253            } else {
254                self.targeted_functions.push(self.get_function(selector)?.clone());
255            }
256        }
257        Ok(())
258    }
259}
260
261/// Test contract which is testing its invariants.
262#[derive(Clone, Debug)]
263pub struct InvariantContract<'a> {
264    /// Address of the test contract.
265    pub address: Address,
266    /// Invariant function present in the test contract.
267    pub invariant_function: &'a Function,
268    /// If true, `afterInvariant` function is called after each invariant run.
269    pub call_after_invariant: bool,
270    /// ABI of the test contract.
271    pub abi: &'a JsonAbi,
272}
273
274impl<'a> InvariantContract<'a> {
275    /// Creates a new invariant contract.
276    pub fn new(
277        address: Address,
278        invariant_function: &'a Function,
279        call_after_invariant: bool,
280        abi: &'a JsonAbi,
281    ) -> Self {
282        Self { address, invariant_function, call_after_invariant, abi }
283    }
284
285    /// Returns true if this is an optimization mode invariant (returns int256).
286    pub fn is_optimization(&self) -> bool {
287        is_optimization_invariant(self.invariant_function)
288    }
289}
foundry_evm_fuzz/invariant/mod.rs

foundry_evm_fuzz/invariant/
mod.rs
