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}