foundry_evm_fuzz/invariant/
filters.rs

1use alloy_json_abi::{Function, JsonAbi};
2use alloy_primitives::{Address, Selector};
3use foundry_compilers::ArtifactId;
4use foundry_evm_core::utils::get_function;
5use std::collections::BTreeMap;
6
7/// Contains which contracts are to be targeted or excluded on an invariant test through their
8/// artifact identifiers.
9#[derive(Default)]
10pub struct ArtifactFilters {
11    /// List of `contract_path:contract_name` along with selectors, which are to be targeted. If
12    /// list of functions is not empty, target only those.
13    pub targeted: BTreeMap<String, Vec<Selector>>,
14    /// List of `contract_path:contract_name` which are to be excluded.
15    pub excluded: Vec<String>,
16}
17
18impl ArtifactFilters {
19    /// Returns `true` if the given identifier matches this filter.
20    pub fn matches(&self, identifier: &str) -> bool {
21        (self.targeted.is_empty() || self.targeted.contains_key(identifier)) &&
22            (self.excluded.is_empty() || !self.excluded.iter().any(|id| id == identifier))
23    }
24
25    /// Gets all the targeted functions from `artifact`. Returns error, if selectors do not match
26    /// the `artifact`.
27    ///
28    /// An empty vector means that it targets any mutable function.
29    pub fn get_targeted_functions(
30        &self,
31        artifact: &ArtifactId,
32        abi: &JsonAbi,
33    ) -> eyre::Result<Option<Vec<Function>>> {
34        if let Some(selectors) = self.targeted.get(&artifact.identifier()) {
35            let functions = selectors
36                .iter()
37                .map(|selector| get_function(&artifact.name, *selector, abi).cloned())
38                .collect::<eyre::Result<Vec<_>>>()?;
39            // targetArtifactSelectors > excludeArtifacts > targetArtifacts
40            if functions.is_empty() && self.excluded.contains(&artifact.identifier()) {
41                return Ok(None)
42            }
43            return Ok(Some(functions))
44        }
45        // If no contract is specifically targeted, and this contract is not excluded, then accept
46        // all functions.
47        if self.targeted.is_empty() && !self.excluded.contains(&artifact.identifier()) {
48            return Ok(Some(vec![]))
49        }
50        Ok(None)
51    }
52}
53
54/// Filter for acceptable senders to use for invariant testing. Exclusion takes priority if
55/// clashing.
56///
57/// `address(0)` is excluded by default.
58#[derive(Default)]
59pub struct SenderFilters {
60    pub targeted: Vec<Address>,
61    pub excluded: Vec<Address>,
62}
63
64impl SenderFilters {
65    pub fn new(mut targeted: Vec<Address>, mut excluded: Vec<Address>) -> Self {
66        let addr_0 = Address::ZERO;
67        if !excluded.contains(&addr_0) {
68            excluded.push(addr_0);
69        }
70        targeted.retain(|addr| !excluded.contains(addr));
71        Self { targeted, excluded }
72    }
73}