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

anvil/eth/backend/mem/
state.rs
