evm_state_diff/
filter.rs

1use std::collections::BTreeSet;
2use std::str::FromStr;
3
4use crate::diff::{AccountDiff, StateDiff};
5use crate::primitives::{Address, B256, U256};
6
7#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
8pub enum Kind {
9    Balance,
10    Nonce,
11    Code,
12    Storage,
13}
14
15impl FromStr for Kind {
16    type Err = String;
17
18    fn from_str(s: &str) -> Result<Self, String> {
19        match s.to_ascii_lowercase().as_str() {
20            "balance" => Ok(Kind::Balance),
21            "nonce" => Ok(Kind::Nonce),
22            "code" => Ok(Kind::Code),
23            "storage" => Ok(Kind::Storage),
24            other => Err(format!("unknown change kind `{other}` (expected balance, nonce, code or storage)")),
25        }
26    }
27}
28
29/// Selects which accounts and fields of a diff to keep. An empty set means
30/// "no restriction".
31#[derive(Clone, Debug, Default)]
32pub struct Filter {
33    pub include: BTreeSet<Address>,
34    pub exclude: BTreeSet<Address>,
35    pub kinds: BTreeSet<Kind>,
36    pub slots: BTreeSet<B256>,
37    /// Drop balance changes whose absolute delta is below this many wei.
38    pub min_balance_delta: Option<U256>,
39}
40
41impl Filter {
42    pub fn is_noop(&self) -> bool {
43        self.include.is_empty()
44            && self.exclude.is_empty()
45            && self.kinds.is_empty()
46            && self.slots.is_empty()
47            && self.min_balance_delta.is_none()
48    }
49
50    fn wants(&self, kind: Kind) -> bool {
51        let kind_ok = self.kinds.is_empty() || self.kinds.contains(&kind);
52        // Selecting slots narrows the output to storage.
53        kind_ok && (self.slots.is_empty() || kind == Kind::Storage)
54    }
55
56    fn account(&self, mut acct: AccountDiff) -> Option<AccountDiff> {
57        if !self.wants(Kind::Balance) {
58            acct.balance = None;
59        }
60        if let (Some(min), Some(delta)) = (self.min_balance_delta, acct.balance_delta()) {
61            if delta.magnitude < min {
62                acct.balance = None;
63            }
64        }
65        if !self.wants(Kind::Nonce) {
66            acct.nonce = None;
67        }
68        if !self.wants(Kind::Code) {
69            acct.code = None;
70            acct.code_hash = None;
71        }
72        if !self.wants(Kind::Storage) {
73            acct.storage.clear();
74        } else if !self.slots.is_empty() {
75            acct.storage.retain(|slot, _| self.slots.contains(slot));
76        }
77        acct.has_changes().then_some(acct)
78    }
79
80    pub fn apply(&self, diff: &StateDiff) -> StateDiff {
81        if self.is_noop() {
82            return diff.clone();
83        }
84        let accounts = diff
85            .accounts
86            .iter()
87            .filter(|(addr, _)| self.include.is_empty() || self.include.contains(addr))
88            .filter(|(addr, _)| !self.exclude.contains(addr))
89            .filter_map(|(addr, acct)| self.account(acct.clone()).map(|a| (*addr, a)))
90            .collect();
91        StateDiff { accounts }
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use crate::trace::parse_traces_str;
99
100    fn sample() -> StateDiff {
101        let json = r#"{"pre":{
102            "0x0101010101010101010101010101010101010101":{"balance":"0x100","nonce":1,"storage":{"0x1":"0x1","0x2":"0x2"}},
103            "0x0202020202020202020202020202020202020202":{"balance":"0x100"}},
104          "post":{
105            "0x0101010101010101010101010101010101010101":{"balance":"0x90","nonce":2,"storage":{"0x1":"0x5","0x2":"0x6"}},
106            "0x0202020202020202020202020202020202020202":{"balance":"0x101"}}}"#;
107        StateDiff::from_prestate(&parse_traces_str(json).unwrap()[0].diff)
108    }
109
110    #[test]
111    fn noop_keeps_everything() {
112        assert_eq!(Filter::default().apply(&sample()), sample());
113    }
114
115    #[test]
116    fn address_filters() {
117        let a = Address([1; 20]);
118        let only_a = Filter { include: [a].into(), ..Filter::default() }.apply(&sample());
119        assert_eq!(only_a.accounts.keys().collect::<Vec<_>>(), vec![&a]);
120        let not_a = Filter { exclude: [a].into(), ..Filter::default() }.apply(&sample());
121        assert!(!not_a.accounts.contains_key(&a));
122        assert_eq!(not_a.accounts.len(), 1);
123    }
124
125    #[test]
126    fn kind_slot_and_threshold_filters() {
127        let storage_only = Filter { kinds: [Kind::Storage].into(), ..Filter::default() }.apply(&sample());
128        assert_eq!(storage_only.accounts.len(), 1, "account without storage changes is dropped");
129        let acct = storage_only.accounts.values().next().unwrap();
130        assert!(acct.balance.is_none() && acct.nonce.is_none());
131        assert_eq!(acct.storage.len(), 2);
132
133        let slot = B256::from_word("0x2").unwrap();
134        let one_slot = Filter { slots: [slot].into(), ..Filter::default() }.apply(&sample());
135        assert_eq!(
136            one_slot.accounts.values().next().unwrap().storage.keys().collect::<Vec<_>>(),
137            vec![&slot]
138        );
139
140        let big = Filter {
141            min_balance_delta: Some(U256::from_u128(2)),
142            kinds: [Kind::Balance].into(),
143            ..Filter::default()
144        }
145        .apply(&sample());
146        assert_eq!(big.accounts.len(), 1, "the 1 wei change is below the threshold");
147    }
148
149    #[test]
150    fn kind_parsing() {
151        assert_eq!("Storage".parse::<Kind>(), Ok(Kind::Storage));
152        assert!("gas".parse::<Kind>().is_err());
153    }
154}
evm_state_diff/filter.rs

evm_state_diff/
filter.rs
