anvil/eth/backend/mem/
state.rs

1//! Support for generating the state root for memdb storage
2
3use alloy_primitives::{
4    B256, Bytes, U256, keccak256,
5    map::{AddressMap, B256Map, HashSet, U256Map},
6};
7use alloy_rlp::Encodable;
8use alloy_trie::{
9    EMPTY_ROOT_HASH, HashBuilder, Nibbles, TrieMask,
10    nodes::{BranchNodeRef, ExtensionNodeRef, LeafNodeRef, RlpNode},
11};
12use revm::{
13    database::{AccountState, DbAccount},
14    state::{Account, AccountInfo},
15};
16use std::{array, mem};
17
18/// Incrementally maintains the state trie used for block headers.
19///
20/// The old state-root path rebuilt and sorted every account and storage trie after each block.
21/// That made EIP-2935's block-hash storage contract turn mining into linear work as its ring was
22/// populated. This cache keeps an in-memory Merkle Patricia trie and only rehashes paths changed
23/// by the latest block.
24#[derive(Debug, Default)]
25pub struct StateRootCache {
26    trie: Option<IncrementalStateTrie>,
27    dirty: AddressMap<DirtyAccount>,
28    /// Reused while encoding dirty trie nodes.
29    rlp_buf: Vec<u8>,
30}
31
32#[derive(Debug, Default)]
33struct DirtyAccount {
34    storage: HashSet<U256>,
35    reset_storage: bool,
36}
37
38impl StateRootCache {
39    /// Records changes that will be committed to the database.
40    pub fn record_changes(&mut self, changes: &AddressMap<Account>) {
41        for (address, account) in changes {
42            if !account.is_touched() {
43                continue;
44            }
45
46            let dirty = self.dirty.entry(*address).or_default();
47            dirty.reset_storage |= account.is_created() || account.is_selfdestructed();
48            dirty.storage.extend(account.changed_storage_slots().map(|(slot, _)| *slot));
49        }
50    }
51
52    /// Records an account-info change or a database load caused by `basic`.
53    pub fn record_account(&mut self, address: alloy_primitives::Address) {
54        self.dirty.entry(address).or_default();
55    }
56
57    /// Records a storage change or a database load caused by `storage`.
58    pub fn record_storage(&mut self, address: alloy_primitives::Address, slot: U256) {
59        self.dirty.entry(address).or_default().storage.insert(slot);
60    }
61
62    /// Records a cumulative overlay relative to the overlay used for the previous root.
63    /// Returns false and invalidates the cache if entries disappeared or storage no longer
64    /// replaces the base. The caller must then supply full, merged state to `root`.
65    /// The base must not contain storage-bearing `NotExisting` accounts: their storage is
66    /// omitted from the trie but can reappear when a full database merge touches them.
67    pub fn record_overlay(
68        &mut self,
69        accounts: &AddressMap<DbAccount>,
70        previous: &AddressMap<DbAccount>,
71    ) -> bool {
72        let incremental = previous.iter().all(|(address, previous)| {
73            accounts.get(address).is_some_and(|account| {
74                matches!(
75                    account.account_state,
76                    AccountState::StorageCleared | AccountState::NotExisting
77                ) || (!matches!(
78                    previous.account_state,
79                    AccountState::StorageCleared | AccountState::NotExisting
80                ) && previous.storage.keys().all(|slot| account.storage.contains_key(slot)))
81            })
82        });
83        if !incremental {
84            self.invalidate();
85            return false;
86        }
87
88        for (address, account) in accounts {
89            let dirty = self.dirty.entry(*address).or_default();
90            if account.account_state.is_storage_cleared() {
91                dirty.reset_storage = true;
92            } else {
93                dirty.storage.extend(account.storage.keys().copied());
94            }
95        }
96        true
97    }
98
99    /// Records every account and storage slot in `accounts` that differs from `previous`.
100    ///
101    /// `previous` must be the full state passed to the previous `root` call, and `accounts` must
102    /// be passed to the next one. Comparing complete states captures changes however they reached
103    /// `accounts`.
104    pub fn record_diff(
105        &mut self,
106        accounts: &AddressMap<DbAccount>,
107        previous: &AddressMap<DbAccount>,
108    ) {
109        for address in previous.keys().filter(|address| !accounts.contains_key(*address)) {
110            self.record_account(*address);
111        }
112        for (address, account) in accounts {
113            // New and recreated accounts have no storage trie left to update.
114            let Some(previous) = previous
115                .get(address)
116                .filter(|previous| previous.account_state != AccountState::NotExisting)
117            else {
118                if account.account_state != AccountState::NotExisting {
119                    self.dirty.entry(*address).or_default().reset_storage = true;
120                }
121                continue;
122            };
123            if previous.account_state != account.account_state || previous.info != account.info {
124                self.record_account(*address);
125            }
126            for (slot, value) in &account.storage {
127                if previous.storage.get(slot) != Some(value) {
128                    self.record_storage(*address, *slot);
129                }
130            }
131            for slot in previous.storage.keys() {
132                if !account.storage.contains_key(slot) {
133                    self.record_storage(*address, *slot);
134                }
135            }
136        }
137    }
138
139    /// Invalidates the trie after wholesale database replacement or clearing.
140    pub fn invalidate(&mut self) {
141        self.trie = None;
142        self.dirty.clear();
143    }
144
145    /// Returns the current root, applying only changes recorded since the previous call.
146    /// After initialization from full state, `accounts` may be a partial overlay containing
147    /// every dirty address; uncached accounts and storage slots retain their previous values.
148    pub fn root(&mut self, accounts: &AddressMap<DbAccount>) -> B256 {
149        let Self { trie, dirty, rlp_buf } = self;
150        if trie.is_none() {
151            *trie = Some(IncrementalStateTrie::from_accounts(accounts, rlp_buf));
152            dirty.clear();
153            return trie.as_mut().unwrap().root(rlp_buf);
154        }
155
156        let trie = trie.as_mut().unwrap();
157        for (address, dirty) in mem::take(dirty) {
158            let hashed_address = keccak256(address);
159            let Some(account) = accounts
160                .get(&address)
161                .filter(|account| account.account_state != AccountState::NotExisting)
162            else {
163                trie.accounts.remove(hashed_address);
164                trie.storage.remove(&hashed_address);
165                continue;
166            };
167
168            let storage_trie = if dirty.reset_storage {
169                trie.storage
170                    .entry(hashed_address)
171                    .insert_entry(IncrementalTrie::from_storage(&account.storage))
172                    .into_mut()
173            } else {
174                let storage_trie = trie.storage.entry(hashed_address).or_default();
175                for slot in dirty.storage {
176                    let key = keccak256(slot.to_be_bytes::<32>());
177                    if let Some(value) = account.storage.get(&slot).filter(|value| !value.is_zero())
178                    {
179                        storage_trie.insert(key, alloy_rlp::encode(value));
180                    } else {
181                        storage_trie.remove(key);
182                    }
183                }
184                storage_trie
185            };
186            let storage_root = storage_trie.root_with_buf(rlp_buf);
187            trie.accounts.insert(
188                hashed_address,
189                trie_account_rlp_with_storage_root(&account.info, storage_root),
190            );
191        }
192
193        trie.root(rlp_buf)
194    }
195}
196
197#[derive(Debug, Default)]
198struct IncrementalStateTrie {
199    accounts: IncrementalTrie,
200    storage: B256Map<IncrementalTrie>,
201}
202
203impl IncrementalStateTrie {
204    fn from_accounts(accounts: &AddressMap<DbAccount>, rlp_buf: &mut Vec<u8>) -> Self {
205        let mut trie = Self::default();
206        for (address, account) in accounts {
207            if account.account_state == AccountState::NotExisting {
208                continue;
209            }
210            let hashed_address = keccak256(address);
211            let mut storage_trie = IncrementalTrie::from_storage(&account.storage);
212            let storage_root = storage_trie.root_with_buf(rlp_buf);
213            trie.accounts.insert(
214                hashed_address,
215                trie_account_rlp_with_storage_root(&account.info, storage_root),
216            );
217            trie.storage.insert(hashed_address, storage_trie);
218        }
219        trie
220    }
221
222    fn root(&mut self, rlp_buf: &mut Vec<u8>) -> B256 {
223        self.accounts.root_with_buf(rlp_buf)
224    }
225}
226
227/// A mutable Merkle Patricia trie that caches the RLP reference for every unchanged node.
228#[derive(Debug, Default)]
229struct IncrementalTrie {
230    root: TrieNode,
231}
232
233impl IncrementalTrie {
234    fn from_storage(storage: &U256Map<U256>) -> Self {
235        let mut trie = Self::default();
236        for (slot, value) in storage.iter().filter(|(_, value)| !value.is_zero()) {
237            trie.insert(keccak256(slot.to_be_bytes::<32>()), alloy_rlp::encode(value));
238        }
239        trie
240    }
241
242    fn insert(&mut self, key: B256, value: Vec<u8>) {
243        self.root.insert(Nibbles::unpack(key), value);
244    }
245
246    fn remove(&mut self, key: B256) {
247        self.root.remove(Nibbles::unpack(key));
248    }
249
250    #[cfg(test)]
251    fn root(&mut self) -> B256 {
252        self.root_with_buf(&mut Vec::new())
253    }
254
255    fn root_with_buf(&mut self, rlp_buf: &mut Vec<u8>) -> B256 {
256        let Some(root) = self.root.rlp(rlp_buf) else { return EMPTY_ROOT_HASH };
257        root.as_hash().unwrap_or_else(|| keccak256(root.as_ref()))
258    }
259}
260
261#[derive(Debug, Default)]
262struct TrieNode {
263    kind: TrieNodeKind,
264    rlp: Option<RlpNode>,
265}
266
267#[derive(Debug, Default)]
268enum TrieNodeKind {
269    #[default]
270    Empty,
271    Leaf {
272        path: Nibbles,
273        value: Vec<u8>,
274    },
275    Extension {
276        path: Nibbles,
277        child: Box<TrieNode>,
278    },
279    Branch {
280        children: [Option<Box<TrieNode>>; 16],
281    },
282}
283
284impl TrieNode {
285    const fn leaf(path: Nibbles, value: Vec<u8>) -> Self {
286        Self { kind: TrieNodeKind::Leaf { path, value }, rlp: None }
287    }
288
289    fn extension(path: Nibbles, child: Self) -> Self {
290        debug_assert!(!path.is_empty());
291        Self { kind: TrieNodeKind::Extension { path, child: Box::new(child) }, rlp: None }
292    }
293
294    const fn branch(children: [Option<Box<Self>>; 16]) -> Self {
295        Self { kind: TrieNodeKind::Branch { children }, rlp: None }
296    }
297
298    fn empty_children() -> [Option<Box<Self>>; 16] {
299        Default::default()
300    }
301
302    fn insert(&mut self, key: Nibbles, value: Vec<u8>) {
303        let kind = mem::take(&mut self.kind);
304        self.rlp = None;
305        *self = match kind {
306            TrieNodeKind::Empty => Self::leaf(key, value),
307            TrieNodeKind::Leaf { path, value: old_value } => {
308                let common = path.common_prefix_length(&key);
309                if common == path.len() {
310                    debug_assert_eq!(common, key.len());
311                    Self::leaf(path, value)
312                } else {
313                    let mut children = Self::empty_children();
314                    let old_index = path.get(common).unwrap() as usize;
315                    let new_index = key.get(common).unwrap() as usize;
316                    children[old_index] =
317                        Some(Box::new(Self::leaf(path.slice(common + 1..), old_value)));
318                    children[new_index] =
319                        Some(Box::new(Self::leaf(key.slice(common + 1..), value)));
320                    let branch = Self::branch(children);
321                    if common == 0 { branch } else { Self::extension(path.slice(..common), branch) }
322                }
323            }
324            TrieNodeKind::Extension { path, mut child } => {
325                let common = path.common_prefix_length(&key);
326                if common == path.len() {
327                    child.insert(key.slice(common..), value);
328                    Self::extension(path, *child)
329                } else {
330                    let mut children = Self::empty_children();
331                    let old_index = path.get(common).unwrap() as usize;
332                    let old_path = path.slice(common + 1..);
333                    let old_child = if old_path.is_empty() {
334                        *child
335                    } else {
336                        Self::extension(old_path, *child)
337                    };
338                    children[old_index] = Some(Box::new(old_child));
339
340                    let new_index = key.get(common).unwrap() as usize;
341                    children[new_index] =
342                        Some(Box::new(Self::leaf(key.slice(common + 1..), value)));
343                    let branch = Self::branch(children);
344                    if common == 0 { branch } else { Self::extension(path.slice(..common), branch) }
345                }
346            }
347            TrieNodeKind::Branch { mut children } => {
348                let index = key.first().expect("trie keys have equal lengths") as usize;
349                children[index]
350                    .get_or_insert_with(|| Box::new(Self::default()))
351                    .insert(key.slice(1..), value);
352                Self::branch(children)
353            }
354        };
355    }
356
357    fn remove(&mut self, key: Nibbles) {
358        let kind = mem::take(&mut self.kind);
359        self.rlp = None;
360        *self = match kind {
361            TrieNodeKind::Empty => Self::default(),
362            TrieNodeKind::Leaf { path, value } => {
363                if path == key {
364                    Self::default()
365                } else {
366                    Self::leaf(path, value)
367                }
368            }
369            TrieNodeKind::Extension { path, mut child } => {
370                if key.starts_with(&path) {
371                    child.remove(key.slice(path.len()..));
372                    Self::normalize_extension(path, *child)
373                } else {
374                    Self::extension(path, *child)
375                }
376            }
377            TrieNodeKind::Branch { mut children } => {
378                let index = key.first().expect("trie keys have equal lengths") as usize;
379                if let Some(child) = &mut children[index] {
380                    child.remove(key.slice(1..));
381                    if matches!(child.kind, TrieNodeKind::Empty) {
382                        children[index] = None;
383                    }
384                }
385                Self::normalize_branch(children)
386            }
387        };
388    }
389
390    fn normalize_extension(path: Nibbles, child: Self) -> Self {
391        match child.kind {
392            TrieNodeKind::Empty => Self::default(),
393            TrieNodeKind::Leaf { path: child_path, value } => {
394                Self::leaf(path.join(&child_path), value)
395            }
396            TrieNodeKind::Extension { path: child_path, child } => {
397                Self::extension(path.join(&child_path), *child)
398            }
399            TrieNodeKind::Branch { children } => Self::extension(path, Self::branch(children)),
400        }
401    }
402
403    fn normalize_branch(mut children: [Option<Box<Self>>; 16]) -> Self {
404        let mut indexes =
405            children.iter().enumerate().filter_map(|(index, child)| child.as_ref().map(|_| index));
406        let Some(index) = indexes.next() else { return Self::default() };
407        if indexes.next().is_some() {
408            return Self::branch(children);
409        }
410
411        let child = *children[index].take().unwrap();
412        let prefix = Nibbles::from_nibbles([index as u8]);
413        match child.kind {
414            TrieNodeKind::Empty => unreachable!("empty branch children are removed"),
415            TrieNodeKind::Leaf { path, value } => Self::leaf(prefix.join(&path), value),
416            TrieNodeKind::Extension { path, child } => Self::extension(prefix.join(&path), *child),
417            TrieNodeKind::Branch { children } => Self::extension(prefix, Self::branch(children)),
418        }
419    }
420
421    fn rlp(&mut self, out: &mut Vec<u8>) -> Option<RlpNode> {
422        if let Some(rlp) = &self.rlp {
423            return Some(rlp.clone());
424        }
425
426        let rlp = self.encode(out)?;
427        self.rlp = Some(rlp.clone());
428        Some(rlp)
429    }
430
431    /// Encodes this node into `out`, returning its RLP reference without consulting the cache.
432    fn encode(&mut self, out: &mut Vec<u8>) -> Option<RlpNode> {
433        let rlp = match &mut self.kind {
434            TrieNodeKind::Empty => return None,
435            TrieNodeKind::Leaf { path, value } => {
436                out.clear();
437                LeafNodeRef::new(path, value).rlp(out)
438            }
439            TrieNodeKind::Extension { path, child } => {
440                let child = child.rlp(out).expect("extension nodes have a child");
441                out.clear();
442                ExtensionNodeRef::new(path, child.as_ref()).rlp(out)
443            }
444            TrieNodeKind::Branch { children } => {
445                let mut stack: [RlpNode; 16] = array::from_fn(|_| RlpNode::default());
446                let mut stack_len = 0;
447                let mut state_mask = TrieMask::default();
448                for (index, child) in children.iter_mut().enumerate() {
449                    if let Some(child) = child {
450                        stack[stack_len] = child.rlp(out).expect("branch children are not empty");
451                        stack_len += 1;
452                        state_mask.set_bit(index as u8);
453                    }
454                }
455                out.clear();
456                BranchNodeRef::new(&stack[..stack_len], state_mask).rlp(out)
457            }
458        };
459        Some(rlp)
460    }
461
462    /// Appends the RLP encoding of this node and all of its descendants to `nodes`.
463    fn collect_nodes(&mut self, rlp_buf: &mut Vec<u8>, nodes: &mut Vec<Bytes>) {
464        match &mut self.kind {
465            TrieNodeKind::Empty => return,
466            TrieNodeKind::Leaf { .. } => {}
467            TrieNodeKind::Extension { child, .. } => child.collect_nodes(rlp_buf, nodes),
468            TrieNodeKind::Branch { children } => {
469                for child in children.iter_mut().flatten() {
470                    child.collect_nodes(rlp_buf, nodes);
471                }
472            }
473        }
474        self.encode(rlp_buf);
475        nodes.push(Bytes::copy_from_slice(rlp_buf));
476    }
477}
478
479/// Builds the state trie for the given accounts and returns its root together with the RLP
480/// encoding of every node of the account trie and all storage tries.
481///
482/// The node set is a witness for any execution against this state: it is a strict superset of
483/// the nodes touched by any particular block.
484pub fn state_trie_witness(accounts: &AddressMap<DbAccount>) -> (B256, Vec<Bytes>) {
485    let mut rlp_buf = Vec::new();
486    let mut account_trie = IncrementalTrie::default();
487    let mut nodes = Vec::new();
488    for (address, account) in accounts {
489        if account.account_state == AccountState::NotExisting {
490            continue;
491        }
492
493        // Keep only one storage trie alive while collecting the witness.
494        let mut storage_trie = IncrementalTrie::from_storage(&account.storage);
495        let storage_root = storage_trie.root_with_buf(&mut rlp_buf);
496        storage_trie.root.collect_nodes(&mut rlp_buf, &mut nodes);
497        account_trie.insert(
498            keccak256(address),
499            trie_account_rlp_with_storage_root(&account.info, storage_root),
500        );
501    }
502    account_trie.root.collect_nodes(&mut rlp_buf, &mut nodes);
503    let root = account_trie.root_with_buf(&mut rlp_buf);
504    nodes.sort_unstable();
505    nodes.dedup();
506    (root, nodes)
507}
508
509pub fn build_root(values: impl IntoIterator<Item = (Nibbles, Vec<u8>)>) -> B256 {
510    let mut builder = HashBuilder::default();
511    for (key, value) in values {
512        builder.add_leaf(key, value.as_ref());
513    }
514    builder.root()
515}
516
517/// Builds state root from the given accounts
518pub fn state_root(accounts: &AddressMap<DbAccount>) -> B256 {
519    build_root(trie_accounts(accounts))
520}
521
522/// Builds storage root from the given storage
523pub fn storage_root(storage: &U256Map<U256>) -> B256 {
524    build_root(trie_storage(storage))
525}
526
527/// Builds iterator over stored key-value pairs ready for storage trie root calculation.
528pub fn trie_storage(storage: &U256Map<U256>) -> Vec<(Nibbles, Vec<u8>)> {
529    let mut storage = storage
530        .iter()
531        .filter(|(_, value)| !value.is_zero())
532        .map(|(key, value)| {
533            let data = alloy_rlp::encode(value);
534            (Nibbles::unpack(keccak256(key.to_be_bytes::<32>())), data)
535        })
536        .collect::<Vec<_>>();
537    storage.sort_by_key(|(key, _)| *key);
538
539    storage
540}
541
542/// Builds iterator over stored key-value pairs ready for account trie root calculation.
543pub fn trie_accounts(accounts: &AddressMap<DbAccount>) -> Vec<(Nibbles, Vec<u8>)> {
544    let mut accounts: Vec<(Nibbles, Vec<u8>)> = accounts
545        .iter()
546        .filter(|(_, account)| account.account_state != AccountState::NotExisting)
547        .map(|(address, account)| {
548            let data = trie_account_rlp(&account.info, &account.storage);
549            (Nibbles::unpack(keccak256(*address)), data)
550        })
551        .collect();
552    accounts.sort_by_key(|(key, _)| *key);
553
554    accounts
555}
556
557/// Returns the RLP for this account.
558pub fn trie_account_rlp(info: &AccountInfo, storage: &U256Map<U256>) -> Vec<u8> {
559    trie_account_rlp_with_storage_root(info, storage_root(storage))
560}
561
562/// Returns the RLP for this account with an already computed storage root.
563fn trie_account_rlp_with_storage_root(info: &AccountInfo, storage_root: B256) -> Vec<u8> {
564    let mut out: Vec<u8> = Vec::new();
565    let list: [&dyn Encodable; 4] = [&info.nonce, &info.balance, &storage_root, &info.code_hash()];
566
567    alloy_rlp::encode_list::<_, dyn Encodable>(&list, &mut out);
568
569    out
570}
571
572#[cfg(test)]
573mod tests {
574    use super::*;
575
576    #[test]
577    fn canonical_roots_omit_zero_storage_and_non_existing_accounts() {
578        let mut storage = U256Map::default();
579        storage.insert(U256::ONE, U256::ZERO);
580        assert_eq!(storage_root(&storage), EMPTY_ROOT_HASH);
581
582        let mut accounts = AddressMap::default();
583        accounts
584            .insert(alloy_primitives::Address::with_last_byte(1), DbAccount::new_not_existing());
585        assert_eq!(state_root(&accounts), EMPTY_ROOT_HASH);
586        assert_eq!(StateRootCache::default().root(&accounts), EMPTY_ROOT_HASH);
587    }
588
589    #[test]
590    fn incremental_roots_preserve_uncached_state_and_replace_cleared_storage() {
591        let address = alloy_primitives::Address::with_last_byte(1);
592        let untouched = alloy_primitives::Address::with_last_byte(2);
593        let [one, two, three, four] = [1u64, 2, 3, 4].map(U256::from);
594        let original = DbAccount {
595            info: AccountInfo { balance: U256::from(100), ..Default::default() },
596            account_state: AccountState::None,
597            storage: [(one, U256::from(11)), (two, U256::from(22))].into_iter().collect(),
598        };
599        let mut full = AddressMap::from_iter([(address, original.clone()), (untouched, original)]);
600        let mut cache = StateRootCache::default();
601        assert_eq!(cache.root(&full), state_root(&full));
602        let mut previous = AddressMap::default();
603        let mut check = |overlay: &AddressMap<DbAccount>, full: &AddressMap<DbAccount>| {
604            assert!(cache.record_overlay(overlay, &previous));
605            assert_eq!(cache.root(overlay), state_root(full));
606            previous.clone_from(overlay);
607        };
608
609        // An ordinary write must retain storage and accounts absent from the overlay.
610        let mut overlay = AddressMap::from_iter([(
611            address,
612            DbAccount {
613                info: full[&address].info.clone(),
614                account_state: AccountState::Touched,
615                storage: [(one, U256::from(33))].into_iter().collect(),
616            },
617        )]);
618        full.get_mut(&address).unwrap().storage.insert(one, U256::from(33));
619        check(&overlay, &full);
620
621        // Zero removes a leaf originally present only in the base; a new slot adds one.
622        overlay.get_mut(&address).unwrap().storage.extend([(two, U256::ZERO), (three, four)]);
623        full.get_mut(&address).unwrap().storage.remove(&two);
624        full.get_mut(&address).unwrap().storage.insert(three, four);
625        check(&overlay, &full);
626
627        // A balance-only override must not clear unread storage.
628        overlay.get_mut(&address).unwrap().info.balance = U256::from(101);
629        full.get_mut(&address).unwrap().info.balance = U256::from(101);
630        check(&overlay, &full);
631
632        // A deleted account must lose both its account leaf and its storage trie.
633        overlay.insert(address, DbAccount::new_not_existing());
634        full.remove(&address);
635        check(&overlay, &full);
636
637        // Recreate, then replace storage again without changing StorageCleared. Each replacement
638        // must discard the previous slot set, including slots introduced by earlier blocks.
639        for slot in [three, four] {
640            let account = DbAccount {
641                info: AccountInfo { balance: U256::from(102), ..Default::default() },
642                account_state: AccountState::StorageCleared,
643                storage: [(slot, U256::from(55))].into_iter().collect(),
644            };
645            overlay.insert(address, account.clone());
646            full.insert(address, account);
647            check(&overlay, &full);
648        }
649    }
650
651    #[test]
652    fn overlay_roots_rebuild_when_base_storage_is_exposed_again() {
653        let address = alloy_primitives::Address::with_last_byte(1);
654        let [one, two] = [1u64, 2].map(U256::from);
655        let base = DbAccount {
656            info: AccountInfo { balance: U256::from(100), ..Default::default() },
657            storage: [(one, U256::from(11)), (two, U256::from(22))].into_iter().collect(),
658            ..Default::default()
659        };
660        for account_state in
661            [AccountState::Touched, AccountState::StorageCleared, AccountState::NotExisting]
662        {
663            let storage = if account_state == AccountState::NotExisting {
664                U256Map::default()
665            } else {
666                [(one, U256::from(33))].into_iter().collect()
667            };
668            let previous = AddressMap::from_iter([(
669                address,
670                DbAccount { account_state: account_state.clone(), storage, ..base.clone() },
671            )]);
672            // Retain every overlay key when losing StorageCleared: the newly exposed base slot
673            // still requires a rebuild. NotExisting likewise has no keys to lose.
674            let storage = if account_state.is_storage_cleared() {
675                previous[&address].storage.clone()
676            } else {
677                U256Map::default()
678            };
679            let mut merged = base.clone();
680            merged.storage.extend(storage.clone());
681            let overlay = AddressMap::from_iter([(
682                address,
683                DbAccount { account_state: AccountState::Touched, storage, ..base.clone() },
684            )]);
685            let mut cache = StateRootCache::default();
686            cache.root(&previous);
687            assert!(!cache.record_overlay(&overlay, &previous));
688            let full = AddressMap::from_iter([(address, merged)]);
689            assert_eq!(cache.root(&full), state_root(&full));
690            assert!(!cache.record_overlay(&AddressMap::default(), &overlay));
691        }
692    }
693
694    fn rebuilt_root(values: &B256Map<Vec<u8>>) -> B256 {
695        let mut leaves = values
696            .iter()
697            .map(|(key, value)| (Nibbles::unpack(*key), value.clone()))
698            .collect::<Vec<_>>();
699        leaves.sort_by_key(|(key, _)| *key);
700        build_root(leaves)
701    }
702
703    #[test]
704    fn incremental_trie_matches_full_rebuild() {
705        let mut trie = IncrementalTrie::default();
706        let mut values = B256Map::default();
707        assert_eq!(trie.root(), EMPTY_ROOT_HASH);
708
709        for index in 0..128 {
710            let key = keccak256(U256::from(index).to_be_bytes::<32>());
711            let value = alloy_rlp::encode(U256::from(index + 1));
712            trie.insert(key, value.clone());
713            values.insert(key, value);
714            assert_eq!(trie.root(), rebuilt_root(&values));
715        }
716
717        for index in (0..128).step_by(3) {
718            let key = keccak256(U256::from(index).to_be_bytes::<32>());
719            let value = alloy_rlp::encode(U256::from(index + 1_000));
720            trie.insert(key, value.clone());
721            values.insert(key, value);
722            assert_eq!(trie.root(), rebuilt_root(&values));
723        }
724
725        for index in (0..128).rev() {
726            let key = keccak256(U256::from(index).to_be_bytes::<32>());
727            trie.remove(key);
728            values.remove(&key);
729            assert_eq!(trie.root(), rebuilt_root(&values));
730        }
731    }
732
733    #[test]
734    fn diff_roots_match_full_rebuild() {
735        let [address, other, empty] = [1, 2, 3].map(alloy_primitives::Address::with_last_byte);
736        let [one, two, three] = [1u64, 2, 3].map(U256::from);
737        let account = DbAccount {
738            info: AccountInfo { balance: U256::from(100), ..Default::default() },
739            storage: [(one, U256::from(11)), (two, U256::from(22))].into_iter().collect(),
740            ..Default::default()
741        };
742        let mut accounts = AddressMap::from_iter([(address, account.clone()), (other, account)]);
743        let mut cache = StateRootCache::default();
744        let mut previous = AddressMap::default();
745        let mut check = |accounts: &AddressMap<DbAccount>| {
746            cache.record_diff(accounts, &previous);
747            assert_eq!(cache.root(accounts), state_root(accounts));
748            previous.clone_from(accounts);
749        };
750        check(&accounts);
751
752        // Slot updates, additions, removals, and zero values.
753        let storage = &mut accounts.get_mut(&address).unwrap().storage;
754        storage.insert(one, U256::from(33));
755        storage.insert(three, U256::from(44));
756        storage.remove(&two);
757        check(&accounts);
758        accounts.get_mut(&address).unwrap().storage.insert(one, U256::ZERO);
759        check(&accounts);
760
761        // Account changes without storage changes.
762        accounts.get_mut(&address).unwrap().info.balance = U256::from(101);
763        check(&accounts);
764        accounts.get_mut(&address).unwrap().account_state = AccountState::Touched;
765        check(&accounts);
766        accounts.insert(
767            empty,
768            DbAccount { account_state: AccountState::Touched, ..Default::default() },
769        );
770        check(&accounts);
771        accounts.get_mut(&empty).unwrap().account_state = AccountState::NotExisting;
772        check(&accounts);
773
774        // A deleted account keeps cached storage, which must reappear when it is recreated.
775        accounts.get_mut(&other).unwrap().account_state = AccountState::NotExisting;
776        check(&accounts);
777        accounts.get_mut(&other).unwrap().storage.insert(three, U256::from(55));
778        check(&accounts);
779        accounts.get_mut(&other).unwrap().account_state = AccountState::Touched;
780        check(&accounts);
781
782        // Replaced storage drops every previous slot.
783        let replaced = accounts.get_mut(&other).unwrap();
784        replaced.account_state = AccountState::StorageCleared;
785        replaced.storage = [(two, U256::from(66))].into_iter().collect();
786        check(&accounts);
787
788        // Accounts missing from the next state lose their leaves.
789        accounts.remove(&address);
790        check(&accounts);
791    }
792}
anvil/eth/backend/mem/state.rs

anvil/eth/backend/mem/
state.rs
