anvil/eth/backend/mem/
state.rs

1//! Support for generating the state root for memdb storage
2
3use alloy_primitives::{
4    B256, 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 mined 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    /// Invalidates the trie after wholesale database replacement or clearing.
63    pub fn invalidate(&mut self) {
64        self.trie = None;
65        self.dirty.clear();
66    }
67
68    /// Returns the current root, applying only changes recorded since the previous call.
69    pub fn root(&mut self, accounts: &AddressMap<DbAccount>) -> B256 {
70        let Self { trie, dirty, rlp_buf } = self;
71        if trie.is_none() {
72            *trie = Some(IncrementalStateTrie::from_accounts(accounts, rlp_buf));
73            dirty.clear();
74            return trie.as_mut().unwrap().root(rlp_buf);
75        }
76
77        let trie = trie.as_mut().unwrap();
78        for (address, dirty) in mem::take(dirty) {
79            let hashed_address = keccak256(address);
80            let Some(account) = accounts
81                .get(&address)
82                .filter(|account| account.account_state != AccountState::NotExisting)
83            else {
84                trie.accounts.remove(hashed_address);
85                trie.storage.remove(&hashed_address);
86                continue;
87            };
88
89            let storage_trie = if dirty.reset_storage {
90                trie.storage
91                    .entry(hashed_address)
92                    .insert_entry(IncrementalTrie::from_storage(&account.storage))
93                    .into_mut()
94            } else {
95                let storage_trie = trie.storage.entry(hashed_address).or_default();
96                for slot in dirty.storage {
97                    let key = keccak256(slot.to_be_bytes::<32>());
98                    if let Some(value) = account.storage.get(&slot).filter(|value| !value.is_zero())
99                    {
100                        storage_trie.insert(key, alloy_rlp::encode(value));
101                    } else {
102                        storage_trie.remove(key);
103                    }
104                }
105                storage_trie
106            };
107            let storage_root = storage_trie.root_with_buf(rlp_buf);
108            trie.accounts.insert(
109                hashed_address,
110                trie_account_rlp_with_storage_root(&account.info, storage_root),
111            );
112        }
113
114        trie.root(rlp_buf)
115    }
116}
117
118#[derive(Debug, Default)]
119struct IncrementalStateTrie {
120    accounts: IncrementalTrie,
121    storage: B256Map<IncrementalTrie>,
122}
123
124impl IncrementalStateTrie {
125    fn from_accounts(accounts: &AddressMap<DbAccount>, rlp_buf: &mut Vec<u8>) -> Self {
126        let mut trie = Self::default();
127        for (address, account) in accounts {
128            if account.account_state == AccountState::NotExisting {
129                continue;
130            }
131            let hashed_address = keccak256(address);
132            let mut storage_trie = IncrementalTrie::from_storage(&account.storage);
133            let storage_root = storage_trie.root_with_buf(rlp_buf);
134            trie.accounts.insert(
135                hashed_address,
136                trie_account_rlp_with_storage_root(&account.info, storage_root),
137            );
138            trie.storage.insert(hashed_address, storage_trie);
139        }
140        trie
141    }
142
143    fn root(&mut self, rlp_buf: &mut Vec<u8>) -> B256 {
144        self.accounts.root_with_buf(rlp_buf)
145    }
146}
147
148/// A mutable Merkle Patricia trie that caches the RLP reference for every unchanged node.
149#[derive(Debug, Default)]
150struct IncrementalTrie {
151    root: TrieNode,
152}
153
154impl IncrementalTrie {
155    fn from_storage(storage: &U256Map<U256>) -> Self {
156        let mut trie = Self::default();
157        for (slot, value) in storage.iter().filter(|(_, value)| !value.is_zero()) {
158            trie.insert(keccak256(slot.to_be_bytes::<32>()), alloy_rlp::encode(value));
159        }
160        trie
161    }
162
163    fn insert(&mut self, key: B256, value: Vec<u8>) {
164        self.root.insert(Nibbles::unpack(key), value);
165    }
166
167    fn remove(&mut self, key: B256) {
168        self.root.remove(Nibbles::unpack(key));
169    }
170
171    #[cfg(test)]
172    fn root(&mut self) -> B256 {
173        self.root_with_buf(&mut Vec::new())
174    }
175
176    fn root_with_buf(&mut self, rlp_buf: &mut Vec<u8>) -> B256 {
177        let Some(root) = self.root.rlp(rlp_buf) else { return EMPTY_ROOT_HASH };
178        root.as_hash().unwrap_or_else(|| keccak256(root.as_ref()))
179    }
180}
181
182#[derive(Debug, Default)]
183struct TrieNode {
184    kind: TrieNodeKind,
185    rlp: Option<RlpNode>,
186}
187
188#[derive(Debug, Default)]
189enum TrieNodeKind {
190    #[default]
191    Empty,
192    Leaf {
193        path: Nibbles,
194        value: Vec<u8>,
195    },
196    Extension {
197        path: Nibbles,
198        child: Box<TrieNode>,
199    },
200    Branch {
201        children: [Option<Box<TrieNode>>; 16],
202    },
203}
204
205impl TrieNode {
206    const fn leaf(path: Nibbles, value: Vec<u8>) -> Self {
207        Self { kind: TrieNodeKind::Leaf { path, value }, rlp: None }
208    }
209
210    fn extension(path: Nibbles, child: Self) -> Self {
211        debug_assert!(!path.is_empty());
212        Self { kind: TrieNodeKind::Extension { path, child: Box::new(child) }, rlp: None }
213    }
214
215    const fn branch(children: [Option<Box<Self>>; 16]) -> Self {
216        Self { kind: TrieNodeKind::Branch { children }, rlp: None }
217    }
218
219    fn empty_children() -> [Option<Box<Self>>; 16] {
220        Default::default()
221    }
222
223    fn insert(&mut self, key: Nibbles, value: Vec<u8>) {
224        let kind = mem::take(&mut self.kind);
225        self.rlp = None;
226        *self = match kind {
227            TrieNodeKind::Empty => Self::leaf(key, value),
228            TrieNodeKind::Leaf { path, value: old_value } => {
229                let common = path.common_prefix_length(&key);
230                if common == path.len() {
231                    debug_assert_eq!(common, key.len());
232                    Self::leaf(path, value)
233                } else {
234                    let mut children = Self::empty_children();
235                    let old_index = path.get(common).unwrap() as usize;
236                    let new_index = key.get(common).unwrap() as usize;
237                    children[old_index] =
238                        Some(Box::new(Self::leaf(path.slice(common + 1..), old_value)));
239                    children[new_index] =
240                        Some(Box::new(Self::leaf(key.slice(common + 1..), value)));
241                    let branch = Self::branch(children);
242                    if common == 0 { branch } else { Self::extension(path.slice(..common), branch) }
243                }
244            }
245            TrieNodeKind::Extension { path, mut child } => {
246                let common = path.common_prefix_length(&key);
247                if common == path.len() {
248                    child.insert(key.slice(common..), value);
249                    Self::extension(path, *child)
250                } else {
251                    let mut children = Self::empty_children();
252                    let old_index = path.get(common).unwrap() as usize;
253                    let old_path = path.slice(common + 1..);
254                    let old_child = if old_path.is_empty() {
255                        *child
256                    } else {
257                        Self::extension(old_path, *child)
258                    };
259                    children[old_index] = Some(Box::new(old_child));
260
261                    let new_index = key.get(common).unwrap() as usize;
262                    children[new_index] =
263                        Some(Box::new(Self::leaf(key.slice(common + 1..), value)));
264                    let branch = Self::branch(children);
265                    if common == 0 { branch } else { Self::extension(path.slice(..common), branch) }
266                }
267            }
268            TrieNodeKind::Branch { mut children } => {
269                let index = key.first().expect("trie keys have equal lengths") as usize;
270                children[index]
271                    .get_or_insert_with(|| Box::new(Self::default()))
272                    .insert(key.slice(1..), value);
273                Self::branch(children)
274            }
275        };
276    }
277
278    fn remove(&mut self, key: Nibbles) {
279        let kind = mem::take(&mut self.kind);
280        self.rlp = None;
281        *self = match kind {
282            TrieNodeKind::Empty => Self::default(),
283            TrieNodeKind::Leaf { path, value } => {
284                if path == key {
285                    Self::default()
286                } else {
287                    Self::leaf(path, value)
288                }
289            }
290            TrieNodeKind::Extension { path, mut child } => {
291                if key.starts_with(&path) {
292                    child.remove(key.slice(path.len()..));
293                    Self::normalize_extension(path, *child)
294                } else {
295                    Self::extension(path, *child)
296                }
297            }
298            TrieNodeKind::Branch { mut children } => {
299                let index = key.first().expect("trie keys have equal lengths") as usize;
300                if let Some(child) = &mut children[index] {
301                    child.remove(key.slice(1..));
302                    if matches!(child.kind, TrieNodeKind::Empty) {
303                        children[index] = None;
304                    }
305                }
306                Self::normalize_branch(children)
307            }
308        };
309    }
310
311    fn normalize_extension(path: Nibbles, child: Self) -> Self {
312        match child.kind {
313            TrieNodeKind::Empty => Self::default(),
314            TrieNodeKind::Leaf { path: child_path, value } => {
315                Self::leaf(path.join(&child_path), value)
316            }
317            TrieNodeKind::Extension { path: child_path, child } => {
318                Self::extension(path.join(&child_path), *child)
319            }
320            TrieNodeKind::Branch { children } => Self::extension(path, Self::branch(children)),
321        }
322    }
323
324    fn normalize_branch(mut children: [Option<Box<Self>>; 16]) -> Self {
325        let mut indexes =
326            children.iter().enumerate().filter_map(|(index, child)| child.as_ref().map(|_| index));
327        let Some(index) = indexes.next() else { return Self::default() };
328        if indexes.next().is_some() {
329            return Self::branch(children);
330        }
331
332        let child = *children[index].take().unwrap();
333        let prefix = Nibbles::from_nibbles([index as u8]);
334        match child.kind {
335            TrieNodeKind::Empty => unreachable!("empty branch children are removed"),
336            TrieNodeKind::Leaf { path, value } => Self::leaf(prefix.join(&path), value),
337            TrieNodeKind::Extension { path, child } => Self::extension(prefix.join(&path), *child),
338            TrieNodeKind::Branch { children } => Self::extension(prefix, Self::branch(children)),
339        }
340    }
341
342    fn rlp(&mut self, out: &mut Vec<u8>) -> Option<RlpNode> {
343        if let Some(rlp) = &self.rlp {
344            return Some(rlp.clone());
345        }
346
347        let rlp = match &mut self.kind {
348            TrieNodeKind::Empty => return None,
349            TrieNodeKind::Leaf { path, value } => {
350                out.clear();
351                LeafNodeRef::new(path, value).rlp(out)
352            }
353            TrieNodeKind::Extension { path, child } => {
354                let child = child.rlp(out).expect("extension nodes have a child");
355                out.clear();
356                ExtensionNodeRef::new(path, child.as_ref()).rlp(out)
357            }
358            TrieNodeKind::Branch { children } => {
359                let mut stack: [RlpNode; 16] = array::from_fn(|_| RlpNode::default());
360                let mut stack_len = 0;
361                let mut state_mask = TrieMask::default();
362                for (index, child) in children.iter_mut().enumerate() {
363                    if let Some(child) = child {
364                        stack[stack_len] = child.rlp(out).expect("branch children are not empty");
365                        stack_len += 1;
366                        state_mask.set_bit(index as u8);
367                    }
368                }
369                out.clear();
370                BranchNodeRef::new(&stack[..stack_len], state_mask).rlp(out)
371            }
372        };
373        self.rlp = Some(rlp.clone());
374        Some(rlp)
375    }
376}
377
378pub fn build_root(values: impl IntoIterator<Item = (Nibbles, Vec<u8>)>) -> B256 {
379    let mut builder = HashBuilder::default();
380    for (key, value) in values {
381        builder.add_leaf(key, value.as_ref());
382    }
383    builder.root()
384}
385
386/// Builds state root from the given accounts
387pub fn state_root(accounts: &AddressMap<DbAccount>) -> B256 {
388    build_root(trie_accounts(accounts))
389}
390
391/// Builds storage root from the given storage
392pub fn storage_root(storage: &U256Map<U256>) -> B256 {
393    build_root(trie_storage(storage))
394}
395
396/// Builds iterator over stored key-value pairs ready for storage trie root calculation.
397pub fn trie_storage(storage: &U256Map<U256>) -> Vec<(Nibbles, Vec<u8>)> {
398    let mut storage = storage
399        .iter()
400        .filter(|(_, value)| !value.is_zero())
401        .map(|(key, value)| {
402            let data = alloy_rlp::encode(value);
403            (Nibbles::unpack(keccak256(key.to_be_bytes::<32>())), data)
404        })
405        .collect::<Vec<_>>();
406    storage.sort_by_key(|(key, _)| *key);
407
408    storage
409}
410
411/// Builds iterator over stored key-value pairs ready for account trie root calculation.
412pub fn trie_accounts(accounts: &AddressMap<DbAccount>) -> Vec<(Nibbles, Vec<u8>)> {
413    let mut accounts: Vec<(Nibbles, Vec<u8>)> = accounts
414        .iter()
415        .filter(|(_, account)| account.account_state != AccountState::NotExisting)
416        .map(|(address, account)| {
417            let data = trie_account_rlp(&account.info, &account.storage);
418            (Nibbles::unpack(keccak256(*address)), data)
419        })
420        .collect();
421    accounts.sort_by_key(|(key, _)| *key);
422
423    accounts
424}
425
426/// Returns the RLP for this account.
427pub fn trie_account_rlp(info: &AccountInfo, storage: &U256Map<U256>) -> Vec<u8> {
428    trie_account_rlp_with_storage_root(info, storage_root(storage))
429}
430
431/// Returns the RLP for this account with an already computed storage root.
432fn trie_account_rlp_with_storage_root(info: &AccountInfo, storage_root: B256) -> Vec<u8> {
433    let mut out: Vec<u8> = Vec::new();
434    let list: [&dyn Encodable; 4] = [&info.nonce, &info.balance, &storage_root, &info.code_hash];
435
436    alloy_rlp::encode_list::<_, dyn Encodable>(&list, &mut out);
437
438    out
439}
440
441#[cfg(test)]
442mod tests {
443    use super::*;
444
445    #[test]
446    fn canonical_roots_omit_zero_storage_and_non_existing_accounts() {
447        let mut storage = U256Map::default();
448        storage.insert(U256::from(1), U256::ZERO);
449        assert_eq!(storage_root(&storage), EMPTY_ROOT_HASH);
450
451        let mut accounts = AddressMap::default();
452        accounts.insert(
453            alloy_primitives::Address::with_last_byte(1),
454            DbAccount { account_state: AccountState::NotExisting, ..Default::default() },
455        );
456        assert_eq!(state_root(&accounts), EMPTY_ROOT_HASH);
457        assert_eq!(StateRootCache::default().root(&accounts), EMPTY_ROOT_HASH);
458    }
459
460    fn rebuilt_root(values: &B256Map<Vec<u8>>) -> B256 {
461        let mut leaves = values
462            .iter()
463            .map(|(key, value)| (Nibbles::unpack(*key), value.clone()))
464            .collect::<Vec<_>>();
465        leaves.sort_by_key(|(key, _)| *key);
466        build_root(leaves)
467    }
468
469    #[test]
470    fn incremental_trie_matches_full_rebuild() {
471        let mut trie = IncrementalTrie::default();
472        let mut values = B256Map::default();
473        assert_eq!(trie.root(), EMPTY_ROOT_HASH);
474
475        for index in 0..128 {
476            let key = keccak256(U256::from(index).to_be_bytes::<32>());
477            let value = alloy_rlp::encode(U256::from(index + 1));
478            trie.insert(key, value.clone());
479            values.insert(key, value);
480            assert_eq!(trie.root(), rebuilt_root(&values));
481        }
482
483        for index in (0..128).step_by(3) {
484            let key = keccak256(U256::from(index).to_be_bytes::<32>());
485            let value = alloy_rlp::encode(U256::from(index + 1_000));
486            trie.insert(key, value.clone());
487            values.insert(key, value);
488            assert_eq!(trie.root(), rebuilt_root(&values));
489        }
490
491        for index in (0..128).rev() {
492            let key = keccak256(U256::from(index).to_be_bytes::<32>());
493            trie.remove(key);
494            values.remove(&key);
495            assert_eq!(trie.root(), rebuilt_root(&values));
496        }
497    }
498}
anvil/eth/backend/mem/state.rs

anvil/eth/backend/mem/
state.rs
