ethrex_trie/
trie.rs

1pub mod db;
2pub mod error;
3pub mod logger;
4mod nibbles;
5pub mod node;
6mod node_hash;
7pub mod rkyv_utils;
8mod rlp;
9#[cfg(test)]
10mod test_utils;
11pub mod threadpool;
12mod trie_iter;
13pub mod trie_sorted;
14mod verify_range;
15use ethereum_types::H256;
16use ethrex_crypto::keccak::keccak_hash;
17use ethrex_crypto::{Crypto, NativeCrypto};
18use ethrex_rlp::constants::RLP_NULL;
19use ethrex_rlp::encode::RLPEncode;
20use rustc_hash::{FxHashMap, FxHashSet};
21use std::collections::BTreeMap;
22use std::sync::{Arc, Mutex};
23
24pub use self::db::{InMemoryTrieDB, TrieDB};
25pub use self::logger::{TrieLogger, TrieWitness};
26pub use self::nibbles::Nibbles;
27pub use self::threadpool::ThreadPool;
28pub use self::verify_range::verify_range;
29pub use self::{
30    node::{Node, NodeRef, OnceLock},
31    node_hash::NodeHash,
32};
33
34pub use self::error::{ExtensionNodeErrorData, InconsistentTreeError, TrieError};
35use self::{node::LeafNode, trie_iter::TrieIterator};
36
37use ethrex_rlp::decode::RLPDecode;
38use lazy_static::lazy_static;
39
40lazy_static! {
41    // Hash value for an empty trie, equal to keccak(RLP_NULL)
42    pub static ref EMPTY_TRIE_HASH: H256 = H256(
43        keccak_hash([RLP_NULL]),
44    );
45}
46
47/// RLP-encoded trie path
48pub type PathRLP = Vec<u8>;
49/// RLP-encoded trie value
50pub type ValueRLP = Vec<u8>;
51/// RLP-encoded trie node
52pub type NodeRLP = Vec<u8>;
53/// Represents a node in the Merkle Patricia Trie.
54pub type TrieNode = (Nibbles, NodeRLP);
55
56/// Ethereum-compatible Merkle Patricia Trie
57pub struct Trie {
58    db: Box<dyn TrieDB>,
59    pub root: NodeRef,
60    pending_removal: FxHashSet<Nibbles>,
61    dirty: FxHashSet<Nibbles>,
62}
63
64impl Default for Trie {
65    fn default() -> Self {
66        Self::new_temp()
67    }
68}
69
70impl Trie {
71    /// Creates a new Trie from a clean DB
72    pub fn new(db: Box<dyn TrieDB>) -> Self {
73        Self {
74            db,
75            root: NodeRef::default(),
76            pending_removal: Default::default(),
77            dirty: Default::default(),
78        }
79    }
80
81    /// Creates a trie from an already-initialized DB and sets root as the root node of the trie
82    pub fn open(db: Box<dyn TrieDB>, root: H256) -> Self {
83        Self {
84            db,
85            root: if root != *EMPTY_TRIE_HASH {
86                NodeHash::from(root).into()
87            } else {
88                Default::default()
89            },
90            pending_removal: Default::default(),
91            dirty: Default::default(),
92        }
93    }
94
95    /// Return a reference to the internal database.
96    ///
97    /// Warning: All changes made to the db will bypass the trie and may cause the trie to suddenly
98    ///   become inconsistent.
99    pub fn db(&self) -> &dyn TrieDB {
100        self.db.as_ref()
101    }
102
103    /// Retrieve an RLP-encoded value from the trie given its RLP-encoded path.
104    pub fn get(&self, pathrlp: &[u8]) -> Result<Option<ValueRLP>, TrieError> {
105        let path = Nibbles::from_bytes(pathrlp);
106
107        if !self.dirty.contains(&path) && self.db().flatkeyvalue_computed(path.clone()) {
108            let Some(value_rlp) = self.db.get(path)? else {
109                return Ok(None);
110            };
111            if value_rlp.is_empty() {
112                return Ok(None);
113            }
114            return Ok(Some(value_rlp));
115        }
116
117        Ok(match self.root {
118            NodeRef::Node(ref node, _) => node.get(self.db.as_ref(), path)?,
119            NodeRef::Hash(hash) if hash.is_valid() => {
120                Node::decode(&self.db.get(Nibbles::default())?.ok_or_else(|| {
121                    TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFound(
122                        hash.finalize(&NativeCrypto),
123                    )))
124                })?)
125                .map_err(TrieError::RLPDecode)?
126                .get(self.db.as_ref(), path)?
127            }
128            _ => None,
129        })
130    }
131
132    /// Insert an RLP-encoded value into the trie.
133    pub fn insert(&mut self, path: PathRLP, value: ValueRLP) -> Result<(), TrieError> {
134        let path = Nibbles::from_bytes(&path);
135        self.pending_removal.remove(&path);
136        self.dirty.insert(path.clone());
137
138        if self.root.is_valid() {
139            // If the trie is not empty, call the root node's insertion logic.
140            self.root
141                .get_node_mut(self.db.as_ref(), Nibbles::default())?
142                .ok_or_else(|| {
143                    TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFoundNoHash))
144                })?
145                .insert(self.db.as_ref(), path, value)?
146        } else {
147            // If the trie is empty, just add a leaf.
148            self.root = Node::from(LeafNode::new(path, value)).into()
149        };
150        self.root.clear_hash();
151
152        Ok(())
153    }
154
155    /// Remove a value from the trie given its RLP-encoded path.
156    /// Returns the value if it was succesfully removed or None if it wasn't part of the trie
157    pub fn remove(&mut self, path: &[u8]) -> Result<Option<ValueRLP>, TrieError> {
158        self.dirty.insert(Nibbles::from_bytes(path));
159        if !self.root.is_valid() {
160            return Ok(None);
161        }
162        self.pending_removal.insert(Nibbles::from_bytes(path));
163
164        // If the trie is not empty, call the root node's removal logic.
165        let (is_trie_empty, value) = self
166            .root
167            .get_node_mut(self.db.as_ref(), Nibbles::default())?
168            .ok_or_else(|| {
169                TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFoundNoHash))
170            })?
171            .remove(self.db.as_ref(), Nibbles::from_bytes(path))?;
172        if is_trie_empty {
173            self.root = NodeRef::default();
174        } else {
175            self.root.clear_hash();
176        }
177
178        Ok(value)
179    }
180
181    /// Return the hash of the trie's root node.
182    /// Returns keccak(RLP_NULL) if the trie is empty
183    /// Also commits changes to the DB
184    pub fn hash(&mut self, crypto: &dyn Crypto) -> Result<H256, TrieError> {
185        self.commit(crypto)?;
186        Ok(self.hash_no_commit(crypto))
187    }
188
189    /// Return the hash of the trie's root node.
190    /// Returns keccak(RLP_NULL) if the trie is empty
191    pub fn hash_no_commit(&self, crypto: &dyn Crypto) -> H256 {
192        if self.root.is_valid() {
193            // 512 is the maximum size of an encoded node
194            let mut buf = Vec::with_capacity(512);
195            self.root
196                .compute_hash_no_alloc(&mut buf, crypto)
197                .finalize(crypto)
198        } else {
199            *EMPTY_TRIE_HASH
200        }
201    }
202
203    pub fn get_root_node(&self, path: Nibbles) -> Result<Arc<Node>, TrieError> {
204        self.root
205            .get_node_checked(self.db.as_ref(), path)?
206            .ok_or_else(|| {
207                TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFound(
208                    self.root
209                        .compute_hash(&NativeCrypto)
210                        .finalize(&NativeCrypto),
211                )))
212            })
213    }
214
215    /// Returns a list of changes in a TrieNode format since last root hash processed.
216    ///
217    /// # Returns
218    ///
219    /// A tuple containing the hash and the list of changes.
220    pub fn collect_changes_since_last_hash(
221        &mut self,
222        crypto: &dyn Crypto,
223    ) -> (H256, Vec<TrieNode>) {
224        let updates = self.commit_without_storing(crypto);
225        let ret_hash = self.hash_no_commit(crypto);
226        (ret_hash, updates)
227    }
228
229    /// Compute the hash of the root node and flush any changes into the database.
230    ///
231    /// This method will also compute the hash of all internal nodes indirectly. It will not clear
232    /// the cached nodes.
233    pub fn commit(&mut self, crypto: &dyn Crypto) -> Result<(), TrieError> {
234        let acc = self.commit_without_storing(crypto);
235        self.db.put_batch(acc)?;
236
237        // Commit the underlying transaction
238        self.db.commit()?;
239
240        Ok(())
241    }
242
243    /// Computes the nodes that would be added if updating the trie.
244    /// Nodes are given with their hash pre-calculated.
245    pub fn commit_without_storing(&mut self, crypto: &dyn Crypto) -> Vec<TrieNode> {
246        let mut acc = Vec::new();
247        if self.root.is_valid() {
248            self.root.commit(Nibbles::default(), &mut acc, crypto);
249        }
250        if self.root.compute_hash(crypto) == NodeHash::Hashed(*EMPTY_TRIE_HASH) {
251            acc.push((Nibbles::default(), vec![RLP_NULL]))
252        }
253        acc.extend(self.pending_removal.drain().map(|nib| (nib, vec![])));
254
255        acc
256    }
257
258    /// Obtain a merkle proof for the given path.
259    /// The proof will contain all the encoded nodes traversed until reaching the node where the path is stored (including this last node).
260    /// The proof will still be constructed even if the path is not stored in the trie, proving its absence.
261    ///
262    /// Note: This method has a different behavior in regard to non-existent trie root nodes. Normal
263    ///   behavior is to return `Err(InconsistentTrie)`, but this method will return
264    ///   `Ok(Vec::new())` instead.
265    pub fn get_proof(&self, path: &[u8]) -> Result<Vec<NodeRLP>, TrieError> {
266        if self.root.is_valid() {
267            let hash = self.root.compute_hash(&NativeCrypto);
268
269            let mut node_path = Vec::new();
270            if let NodeHash::Inline((data, len)) = hash {
271                node_path.push(data[..len as usize].to_vec());
272            }
273
274            let root = match self
275                .root
276                .get_node_checked(self.db.as_ref(), Nibbles::default())?
277            {
278                Some(x) => x,
279                None => return Ok(Vec::new()),
280            };
281            root.get_path(self.db.as_ref(), Nibbles::from_bytes(path), &mut node_path)?;
282
283            Ok(node_path)
284        } else {
285            Ok(Vec::new())
286        }
287    }
288
289    /// Obtains all encoded nodes traversed until reaching the node where every path is stored.
290    /// The list doesn't include the root node, this is returned separately.
291    /// Will still be constructed even if some path is not stored in the trie.
292    pub fn get_proofs(
293        &self,
294        paths: &[PathRLP],
295    ) -> Result<(Option<NodeRLP>, Vec<NodeRLP>), TrieError> {
296        if self.root.is_valid() {
297            let encoded_root = self.get_root_node(Nibbles::default())?.encode_to_vec();
298
299            let mut node_path: FxHashSet<_> = Default::default();
300            for path in paths {
301                let mut nodes = self.get_proof(path)?;
302                nodes.swap_remove(0);
303                node_path.extend(nodes);
304            }
305
306            Ok((Some(encoded_root), node_path.into_iter().collect()))
307        } else {
308            Ok((None, Vec::new()))
309        }
310    }
311
312    pub fn empty_in_memory() -> Self {
313        Self::new(Box::new(InMemoryTrieDB::new(Arc::new(Mutex::new(
314            BTreeMap::new(),
315        )))))
316    }
317
318    /// Gets node with embedded references to child nodes, all in just one `Node`.
319    pub fn get_embedded_root(
320        all_nodes: &BTreeMap<H256, Node>,
321        root_hash: H256,
322    ) -> Result<NodeRef, TrieError> {
323        // If the root hash is of the empty trie then we can get away by setting the NodeRef to default
324        if root_hash == *EMPTY_TRIE_HASH {
325            return Ok(NodeRef::default());
326        }
327
328        let root_rlp = all_nodes.get(&root_hash).ok_or_else(|| {
329            TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFound(root_hash)))
330        })?;
331
332        fn get_embedded_node(
333            all_nodes: &BTreeMap<H256, Node>,
334            cur_node: &Node,
335        ) -> Result<Node, TrieError> {
336            Ok(match cur_node.clone() {
337                Node::Branch(mut node) => {
338                    for choice in &mut node.choices {
339                        let NodeRef::Hash(hash) = *choice else {
340                            continue;
341                        };
342
343                        if hash.is_valid() {
344                            *choice = match all_nodes.get(&hash.finalize(&NativeCrypto)) {
345                                Some(node) => get_embedded_node(all_nodes, node)?.into(),
346                                None => hash.into(),
347                            };
348                        }
349                    }
350
351                    (*node).into()
352                }
353                Node::Extension(mut node) => {
354                    let NodeRef::Hash(hash) = node.child else {
355                        return Ok(node.into());
356                    };
357
358                    node.child = match all_nodes.get(&hash.finalize(&NativeCrypto)) {
359                        Some(node) => get_embedded_node(all_nodes, node)?.into(),
360                        None => hash.into(),
361                    };
362
363                    node.into()
364                }
365                Node::Leaf(node) => node.into(),
366            })
367        }
368
369        let root = get_embedded_node(all_nodes, root_rlp)?;
370        Ok(root.into())
371    }
372
373    /// Gets node with embedded references to child nodes, all in just one `Node`.
374    ///
375    /// Note that this method caches the hash of each node, so it assumes the provided `all_nodes` are well-formed.
376    pub fn get_embedded_root_committed(
377        all_nodes: &FxHashMap<H256, Node>,
378        root_hash: H256,
379        crypto: &dyn Crypto,
380    ) -> Result<NodeRef, TrieError> {
381        // If the root hash is of the empty trie then we can get away by setting the NodeRef to default
382        if root_hash == *EMPTY_TRIE_HASH {
383            return Ok(NodeRef::default());
384        }
385
386        let root_rlp = all_nodes.get(&root_hash).ok_or_else(|| {
387            TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFound(root_hash)))
388        })?;
389
390        /// Creates an embedded node reference with its hash slot pre-seeded.
391        ///
392        /// The caller must guarantee that `hash` is the hash of the referenced
393        /// node, for example because the node was just resolved by looking that
394        /// hash up. Seeding lets later hash computations over the subtree reuse
395        /// the known value instead of re-encoding and re-hashing it.
396        fn node_with_hash(node: Node, hash: NodeHash) -> NodeRef {
397            NodeRef::Node(Arc::new(node), OnceLock::from(hash))
398        }
399
400        fn get_embedded_node_committed(
401            all_nodes: &FxHashMap<H256, Node>,
402            cur_node: &Node,
403            crypto: &dyn Crypto,
404        ) -> Result<Node, TrieError> {
405            Ok(match cur_node.clone() {
406                Node::Branch(mut node) => {
407                    for choice in &mut node.choices {
408                        let NodeRef::Hash(hash) = *choice else {
409                            continue;
410                        };
411
412                        if hash.is_valid() {
413                            *choice = match all_nodes.get(&hash.finalize(crypto)) {
414                                Some(node) => node_with_hash(
415                                    get_embedded_node_committed(all_nodes, node, crypto)?,
416                                    hash,
417                                ),
418                                None => hash.into(),
419                            };
420                        }
421                    }
422
423                    (*node).into()
424                }
425                Node::Extension(mut node) => {
426                    let NodeRef::Hash(hash) = node.child else {
427                        return Ok(node.into());
428                    };
429
430                    node.child = match all_nodes.get(&hash.finalize(crypto)) {
431                        Some(node) => node_with_hash(
432                            get_embedded_node_committed(all_nodes, node, crypto)?,
433                            hash,
434                        ),
435                        None => hash.into(),
436                    };
437
438                    node.into()
439                }
440                Node::Leaf(node) => node.into(),
441            })
442        }
443
444        let root = get_embedded_node_committed(all_nodes, root_rlp, crypto)?;
445        Ok(node_with_hash(root, NodeHash::Hashed(root_hash)))
446    }
447
448    /// Builds a trie from a set of nodes with an empty InMemoryTrieDB as a backend because the nodes are embedded in the root.
449    ///
450    /// Note: This method will not ensure that all node references are valid. Invalid references
451    ///   will cause other methods (including, but not limited to `Trie::get`, `Trie::insert` and
452    ///   `Trie::remove`) to return `Err(InconsistentTrie)`.
453    /// Note: This method will ignore any dangling nodes. All nodes that are not accessible from the
454    ///   root node are considered dangling.
455    pub fn from_nodes(
456        root_hash: H256,
457        state_nodes: &BTreeMap<H256, Node>,
458    ) -> Result<Self, TrieError> {
459        let mut trie = Trie::new(Box::new(InMemoryTrieDB::default()));
460        let root = Self::get_embedded_root(state_nodes, root_hash)?;
461        trie.root = root;
462
463        Ok(trie)
464    }
465
466    /// Builds an in-memory trie from the given elements and returns its hash
467    pub fn compute_hash_from_unsorted_iter(
468        iter: impl Iterator<Item = (PathRLP, ValueRLP)>,
469        crypto: &dyn Crypto,
470    ) -> H256 {
471        let mut trie = Trie::stateless();
472        for (path, value) in iter {
473            // Unwraping here won't panic as our in_memory trie DB won't fail
474            trie.insert(path, value).unwrap();
475        }
476
477        trie.hash_no_commit(crypto)
478    }
479
480    /// Creates a new stateless trie. This trie won't be able to store any nodes so all data will be lost after calculating the hash
481    /// Only use it for proof verification or computing a hash from an iterator
482    pub(crate) fn stateless() -> Trie {
483        // We will only be using the trie's cache so we don't need a working DB
484        struct NullTrieDB;
485
486        impl TrieDB for NullTrieDB {
487            fn get(&self, _key: Nibbles) -> Result<Option<Vec<u8>>, TrieError> {
488                Ok(None)
489            }
490
491            fn put_batch(&self, _key_values: Vec<TrieNode>) -> Result<(), TrieError> {
492                Ok(())
493            }
494        }
495
496        Trie::new(Box::new(NullTrieDB))
497    }
498
499    /// Obtain the encoded node given its path.
500    /// Allows usage of full paths (byte slice of 32 bytes) or compact-encoded nibble slices (with length lower than 32)
501    pub fn get_node(&self, partial_path: &PathRLP) -> Result<Vec<u8>, TrieError> {
502        // Convert compact-encoded nibbles into a byte slice if necessary
503        let partial_path = match partial_path.len() {
504            // Compact-encoded nibbles
505            n if n < 32 => Nibbles::decode_compact(partial_path),
506            // Full path (No conversion needed)
507            32 => Nibbles::from_bytes(partial_path),
508            // We won't handle paths with length over 32
509            _ => return Ok(vec![]),
510        };
511
512        fn get_node_inner(
513            db: &dyn TrieDB,
514            current_path: Nibbles,
515            node: &Node,
516            mut partial_path: Nibbles,
517        ) -> Result<Vec<u8>, TrieError> {
518            // If we reached the end of the partial path, return the current node
519            if partial_path.is_empty() {
520                return Ok(node.encode_to_vec());
521            }
522            match node {
523                Node::Branch(branch_node) => match partial_path.next_choice() {
524                    Some(idx) => {
525                        let child_ref = &branch_node.choices[idx];
526                        if child_ref.is_valid() {
527                            let child_path = current_path.append_new(idx as u8);
528                            let child_node = child_ref
529                                .get_node_checked(db, child_path.clone())?
530                                .ok_or_else(|| {
531                                    TrieError::InconsistentTree(Box::new(
532                                        InconsistentTreeError::NodeNotFoundOnBranchNode(
533                                            child_ref
534                                                .compute_hash(&NativeCrypto)
535                                                .finalize(&NativeCrypto),
536                                            branch_node
537                                                .compute_hash(&NativeCrypto)
538                                                .finalize(&NativeCrypto),
539                                            child_path.clone(),
540                                        ),
541                                    ))
542                                })?;
543                            get_node_inner(db, child_path, &child_node, partial_path)
544                        } else {
545                            Ok(vec![])
546                        }
547                    }
548                    _ => Ok(vec![]),
549                },
550                Node::Extension(extension_node) => {
551                    if partial_path.skip_prefix(&extension_node.prefix)
552                        && extension_node.child.is_valid()
553                    {
554                        let child_path = partial_path.concat(&extension_node.prefix);
555                        let child_node = extension_node
556                            .child
557                            .get_node_checked(db, child_path.clone())?
558                            .ok_or_else(|| {
559                                TrieError::InconsistentTree(Box::new(
560                                    InconsistentTreeError::ExtensionNodeChildNotFound(
561                                        ExtensionNodeErrorData {
562                                            node_hash: extension_node
563                                                .child
564                                                .compute_hash(&NativeCrypto)
565                                                .finalize(&NativeCrypto),
566                                            extension_node_hash: extension_node
567                                                .compute_hash(&NativeCrypto)
568                                                .finalize(&NativeCrypto),
569                                            extension_node_prefix: extension_node.prefix.clone(),
570                                            node_path: child_path.clone(),
571                                        },
572                                    ),
573                                ))
574                            })?;
575                        get_node_inner(db, child_path, &child_node, partial_path)
576                    } else {
577                        Ok(vec![])
578                    }
579                }
580                Node::Leaf(_) => Ok(vec![]),
581            }
582        }
583
584        // Fetch node
585        if self.root.is_valid() {
586            let root_node = self.get_root_node(Default::default())?;
587            get_node_inner(
588                self.db.as_ref(),
589                Default::default(),
590                &root_node,
591                partial_path,
592            )
593        } else {
594            Ok(Vec::new())
595        }
596    }
597
598    pub fn root_node(&self) -> Result<Option<Arc<Node>>, TrieError> {
599        if self.root.is_valid() {
600            self.root.get_node(self.db.as_ref(), Nibbles::default())
601        } else {
602            Ok(None)
603        }
604    }
605
606    /// Creates a new Trie based on a temporary InMemory DB
607    pub fn new_temp() -> Self {
608        let db = InMemoryTrieDB::new(Default::default());
609        Trie::new(Box::new(db))
610    }
611
612    /// Creates a new Trie based on a temporary InMemory DB, with a specified root
613    ///
614    /// This is usually used to create a Trie from a root that was embedded with the rest of the nodes.
615    pub fn new_temp_with_root(root: NodeRef) -> Self {
616        let db = InMemoryTrieDB::new(Default::default());
617        let mut trie = Trie::new(Box::new(db));
618        trie.root = root;
619        trie
620    }
621
622    /// Validates that the Trie isn't missing any nodes expected in the branches
623    ///
624    /// This is used internally with debug assertions to check the status of the trie
625    /// after syncing operations.
626    /// Note: this operation validates the hashes because the iterator uses
627    /// get_node_checked. We shouldn't downgrade that to the unchecked version
628    pub fn validate(self) -> Result<(), TrieError> {
629        let mut expected_count = if self.root.is_valid() { 1 } else { 0 };
630        for (_, node) in self.into_iter() {
631            expected_count -= 1;
632            match node {
633                Node::Branch(branch_node) => {
634                    expected_count += branch_node
635                        .choices
636                        .iter()
637                        .filter(|child| child.is_valid())
638                        .count();
639                }
640                Node::Extension(_) => {
641                    expected_count += 1;
642                }
643                Node::Leaf(_) => {}
644            }
645        }
646        if expected_count != 0 {
647            return Err(TrieError::Verify(format!(
648                "Node count mismatch, expected {expected_count} more"
649            )));
650        }
651        Ok(())
652    }
653
654    /// Validate the trie structure in parallel by splitting at the root branch node.
655    /// Each of the root's 16 subtrees is validated independently using rayon.
656    pub fn validate_parallel(self) -> Result<(), TrieError> {
657        use rayon::prelude::*;
658
659        if !self.root.is_valid() {
660            return Ok(());
661        }
662
663        let db = &*self.db;
664        let root_node = self
665            .root
666            .get_node_checked(db, Nibbles::default())?
667            .ok_or_else(|| TrieError::Verify("Root node not found".to_string()))?;
668
669        match &*root_node {
670            Node::Branch(branch_node) => {
671                let children: Vec<(Nibbles, NodeRef)> = branch_node
672                    .choices
673                    .iter()
674                    .enumerate()
675                    .filter(|(_, child)| child.is_valid())
676                    .map(|(i, child)| {
677                        let path = Nibbles::default().append_new(i as u8);
678                        (path, child.clone())
679                    })
680                    .collect();
681
682                children.par_iter().try_for_each(|(start_path, start_ref)| {
683                    validate_subtree(db, start_path.clone(), start_ref.clone())
684                })
685            }
686            _ => {
687                // Non-branch root (rare): validate sequentially
688                validate_subtree(db, Nibbles::default(), self.root.clone())
689            }
690        }
691    }
692}
693
694/// Validate a subtree rooted at `start_ref`, checking that all referenced nodes exist
695/// and their hashes match.
696fn validate_subtree(
697    db: &dyn TrieDB,
698    start_path: Nibbles,
699    start_ref: NodeRef,
700) -> Result<(), TrieError> {
701    let mut expected_count: isize = 1;
702    let mut stack = vec![(start_path, start_ref)];
703
704    while let Some((path, node_ref)) = stack.pop() {
705        let node = node_ref
706            .get_node_checked(db, path.clone())?
707            .ok_or_else(|| TrieError::Verify(format!("Missing node at path {path:?}")))?;
708
709        expected_count -= 1;
710        match &*node {
711            Node::Branch(branch) => {
712                for (choice, child) in branch.choices.iter().enumerate().rev() {
713                    if child.is_valid() {
714                        expected_count += 1;
715                        stack.push((path.append_new(choice as u8), child.clone()));
716                    }
717                }
718            }
719            Node::Extension(ext) => {
720                expected_count += 1;
721                stack.push((path.concat(&ext.prefix), ext.child.clone()));
722            }
723            Node::Leaf(_) => {}
724        }
725    }
726
727    if expected_count != 0 {
728        return Err(TrieError::Verify(format!(
729            "Node count mismatch in subtree, expected {expected_count} more"
730        )));
731    }
732    Ok(())
733}
734
735impl IntoIterator for Trie {
736    type Item = (Nibbles, Node);
737
738    type IntoIter = TrieIterator;
739
740    fn into_iter(self) -> Self::IntoIter {
741        TrieIterator::new(self)
742    }
743}
744
745pub struct ProofTrie(Trie);
746
747impl ProofTrie {
748    pub fn insert(
749        &mut self,
750        partial_path: Nibbles,
751        external_ref: NodeHash,
752    ) -> Result<(), TrieError> {
753        if self.0.root.is_valid() {
754            // If the trie is not empty, call the root node's insertion logic.
755            self.0
756                .root
757                .get_node_mut(self.0.db.as_ref(), Nibbles::default())?
758                .ok_or_else(|| {
759                    TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFoundNoHash))
760                })?
761                .insert(self.0.db.as_ref(), partial_path, external_ref)?;
762            self.0.root.clear_hash();
763        } else {
764            self.0.root = external_ref.into();
765        };
766
767        Ok(())
768    }
769
770    pub fn hash(&self, crypto: &dyn Crypto) -> H256 {
771        self.0.hash_no_commit(crypto)
772    }
773}
774
775impl From<Trie> for ProofTrie {
776    fn from(value: Trie) -> Self {
777        Self(value)
778    }
779}
ethrex_trie/trie.rs

ethrex_trie/
trie.rs
