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;
67/// 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.
13pub targeted: BTreeMap<String, Vec<Selector>>,
14/// List of `contract_path:contract_name` which are to be excluded.
15pub excluded: Vec<String>,
16}
1718impl ArtifactFilters {
19/// Returns `true` if the given identifier matches this filter.
20pub 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 }
2425/// 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.
29pub fn get_targeted_functions(
30&self,
31 artifact: &ArtifactId,
32 abi: &JsonAbi,
33 ) -> eyre::Result<Option<Vec<Function>>> {
34if let Some(selectors) = self.targeted.get(&artifact.identifier()) {
35let functions = selectors36 .iter()
37 .map(|selector| get_function(&artifact.name, *selector, abi).cloned())
38 .collect::<eyre::Result<Vec<_>>>()?;
39// targetArtifactSelectors > excludeArtifacts > targetArtifacts
40if functions.is_empty() && self.excluded.contains(&artifact.identifier()) {
41return Ok(None)
42 }
43return Ok(Some(functions))
44 }
45// If no contract is specifically targeted, and this contract is not excluded, then accept
46 // all functions.
47if self.targeted.is_empty() && !self.excluded.contains(&artifact.identifier()) {
48return Ok(Some(vec![]))
49 }
50Ok(None)
51 }
52}
5354/// 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 {
60pub targeted: Vec<Address>,
61pub excluded: Vec<Address>,
62}
6364impl SenderFilters {
65pub fn new(mut targeted: Vec<Address>, mut excluded: Vec<Address>) -> Self {
66let addr_0 = Address::ZERO;
67if !excluded.contains(&addr_0) {
68excluded.push(addr_0);
69 }
70targeted.retain(|addr| !excluded.contains(addr));
71Self { targeted, excluded }
72 }
73}