DO NOT MERGE: always store branch nodes in database

src/hash_builder/mod.rs

@@ -132,7 +132,7 @@ impl HashBuilder {
     }
 
     /// Adds a new branch element and its hash to the trie hash builder.
-    pub fn add_branch(&mut self, key: Nibbles, value: B256, stored_in_database: bool) {
+    pub fn add_branch(&mut self, key: Nibbles, value: RlpNode, stored_in_database: bool) {
         assert!(
             key > self.key || (self.key.is_empty() && key.is_empty()),
             "add_branch key {:?} self.key {:?}",
@@ -142,9 +142,9 @@ impl HashBuilder {
         if !self.key.is_empty() {
             self.update(&key);
         } else if key.is_empty() {
-            self.stack.push(RlpNode::word_rlp(&value));
+            self.stack.push(value.clone());
         }
-        self.set_key_value(key, HashBuilderValueRef::Hash(&value));
+        self.set_key_value(key, HashBuilderValueRef::Hash(value));
         self.stored_in_database = stored_in_database;
     }
 
@@ -225,18 +225,39 @@ impl HashBuilder {
                 "prefix lengths after comparing keys"
             );
 
-            // Adjust the state masks for branch calculation
-            let extra_digit = current[len];
+            // Resize masks if needed
             if self.state_masks.len() <= len {
                 let new_len = len + 1;
                 trace!(target: "trie::hash_builder", new_len, old_len = self.state_masks.len(), "scaling state masks to fit");
                 self.state_masks.resize(new_len, TrieMask::default());
             }
+            if self.stored_in_database && self.tree_masks.len() <= len {
+                let new_len = len + 1;
+                trace!(target: "trie::hash_builder", new_len, old_len = self.tree_masks.len(), "scaling tree masks to fit");
+                self.tree_masks.resize(new_len, TrieMask::default());
+            }
+            if self.hash_masks.len() <= len {
+                let new_len = len + 1;
+                trace!(target: "trie::hash_builder", new_len, old_len = self.hash_masks.len(), "scaling hash masks to fit");
+                self.hash_masks.resize(new_len, TrieMask::default());
+            }
+
+            // Adjust the masks for branch calculation
+            let extra_digit = current[len];
             self.state_masks[len] |= TrieMask::from_nibble(extra_digit);
+            if self.stored_in_database {
+                self.tree_masks[len] |= TrieMask::from_nibble(extra_digit);
+            }
+            self.hash_masks[len] |= TrieMask::from_nibble(extra_digit);
+
             trace!(
                 target: "trie::hash_builder",
-                ?extra_digit,
-                state_masks = ?self.state_masks,
+                len,
+                extra_digit,
+                state_mask = ?self.state_masks[len],
+                tree_mask = ?self.stored_in_database.then(|| self.tree_masks[len]),
+                hash_mask = ?self.hash_masks[len],
+                "adjusted branch node masks"
             );
 
             // Adjust the tree masks for exporting to the DB
@@ -272,14 +293,7 @@ impl HashBuilder {
                     }
                     HashBuilderValueRef::Hash(hash) => {
                         trace!(target: "trie::hash_builder", ?hash, "pushing branch node hash");
-                        self.stack.push(RlpNode::word_rlp(hash));
-
-                        if self.stored_in_database {
-                            self.tree_masks[current.len() - 1] |=
-                                TrieMask::from_nibble(current.last().unwrap());
-                        }
-                        self.hash_masks[current.len() - 1] |=
-                            TrieMask::from_nibble(current.last().unwrap());
+                        self.stack.push(hash.clone());
 
                         build_extensions = true;
                     }
@@ -345,7 +359,7 @@ impl HashBuilder {
     ///
     /// Returns the hashes of the children of the branch node, only if `updated_branch_nodes` is
     /// enabled.
-    fn push_branch_node(&mut self, current: &Nibbles, len: usize) -> Vec<B256> {
+    fn push_branch_node(&mut self, current: &Nibbles, len: usize) -> Vec<RlpNode> {
         let state_mask = self.state_masks[len];
         let hash_mask = self.hash_masks[len];
         let branch_node = BranchNodeRef::new(&self.stack, state_mask);
@@ -370,7 +384,14 @@ impl HashBuilder {
         );
         self.stack.resize_with(first_child_idx, Default::default);
 
-        trace!(target: "trie::hash_builder", ?rlp, "pushing branch node with {state_mask:?} mask from stack");
+        trace!(
+            target: "trie::hash_builder",
+            ?rlp,
+            ?state_mask,
+            ?hash_mask,
+            ?children,
+            "pushing branch node to stack"
+        );
         self.stack.push(rlp);
         children
     }
@@ -379,31 +400,24 @@ impl HashBuilder {
     /// to update the masks for the next level and store the branch node and the
     /// masks in the database. We will use that when consuming the intermediate nodes
     /// from the database to efficiently build the trie.
-    fn store_branch_node(&mut self, current: &Nibbles, len: usize, children: Vec<B256>) {
+    fn store_branch_node(&mut self, current: &Nibbles, len: usize, children: Vec<RlpNode>) {
         if len > 0 {
             let parent_index = len - 1;
             self.hash_masks[parent_index] |= TrieMask::from_nibble(current[parent_index]);
+            self.tree_masks[parent_index] |= TrieMask::from_nibble(current[parent_index]);
         }
 
-        let store_in_db_trie = !self.tree_masks[len].is_empty() || !self.hash_masks[len].is_empty();
-        if store_in_db_trie {
-            if len > 0 {
-                let parent_index = len - 1;
-                self.tree_masks[parent_index] |= TrieMask::from_nibble(current[parent_index]);
-            }
-
-            if self.updated_branch_nodes.is_some() {
-                let common_prefix = current.slice(..len);
-                let node = BranchNodeCompact::new(
-                    self.state_masks[len],
-                    self.tree_masks[len],
-                    self.hash_masks[len],
-                    children,
-                    (len == 0).then(|| self.current_root()),
-                );
-                trace!(target: "trie::hash_builder", ?node, "intermediate node");
-                self.updated_branch_nodes.as_mut().unwrap().insert(common_prefix, node);
-            }
+        if self.updated_branch_nodes.is_some() {
+            let common_prefix = current.slice(..len);
+            let node = BranchNodeCompact::new(
+                self.state_masks[len],
+                self.tree_masks[len],
+                self.hash_masks[len],
+                children,
+                (len == 0).then(|| self.current_root()),
+            );
+            trace!(target: "trie::hash_builder", ?node, "intermediate node");
+            self.updated_branch_nodes.as_mut().unwrap().insert(common_prefix, node);
         }
     }
 
@@ -417,7 +431,9 @@ impl HashBuilder {
         if len_from > 0 {
             let flag = TrieMask::from_nibble(current[len_from - 1]);
 
-            self.hash_masks[len_from - 1] &= !flag;
+            if !self.hash_masks[current.len() - 1].is_empty() {
+                self.hash_masks[len_from - 1] |= flag;
+            }
 
             if !self.tree_masks[current.len() - 1].is_empty() {
                 self.tree_masks[len_from - 1] |= flag;
@@ -563,11 +579,11 @@ mod tests {
         // Nibbles 0, 1, 2, 3 have children
         assert_eq!(update.state_mask, TrieMask::new(0b1111));
         // None of the children are stored in the database
-        assert_eq!(update.tree_mask, TrieMask::new(0b0000));
+        assert_eq!(update.tree_mask, TrieMask::new(0b0110));
         // Children under nibbles `1` and `2` are branch nodes with `hashes`
-        assert_eq!(update.hash_mask, TrieMask::new(0b0110));
+        assert_eq!(update.hash_mask, TrieMask::new(0b1111));
         // Calculated when running the hash builder
-        assert_eq!(update.hashes.len(), 2);
+        assert_eq!(update.hashes.len(), 4);
 
         assert_eq!(_root, triehash_trie_root(data));
     }
@@ -596,7 +612,7 @@ mod tests {
     fn test_root_known_hash() {
         let root_hash = b256!("45596e474b536a6b4d64764e4f75514d544577646c414e684271706871446456");
         let mut hb = HashBuilder::default();
-        hb.add_branch(Nibbles::default(), root_hash, false);
+        hb.add_branch(Nibbles::default(), RlpNode::word_rlp(&root_hash), false);
         assert_eq!(hb.root(), root_hash);
     }
 
@@ -630,7 +646,11 @@ mod tests {
 
         let mut hb2 = HashBuilder::default();
         // Insert the branch with the `0x6` shared prefix.
-        hb2.add_branch(Nibbles::from_nibbles_unchecked([0x6]), branch_node_hash, false);
+        hb2.add_branch(
+            Nibbles::from_nibbles_unchecked([0x6]),
+            RlpNode::word_rlp(&branch_node_hash),
+            false,
+        );
 
         assert_eq!(hb.root(), expected);
         assert_eq!(hb2.root(), expected);

src/hash_builder/value.rs

@@ -1,7 +1,9 @@
 use alloc::vec::Vec;
-use alloy_primitives::{hex, B256};
+use alloy_primitives::hex;
 use core::fmt;
 
+use crate::nodes::RlpNode;
+
 /// Hash builder value.
 ///
 /// Stores [`HashBuilderValueRef`] efficiently by reusing resources.
@@ -33,7 +35,7 @@ impl<'u> arbitrary::Arbitrary<'u> for HashBuilderValue {
         let kind = HashBuilderValueKind::arbitrary(g)?;
         let buf = match kind {
             HashBuilderValueKind::Bytes => Vec::arbitrary(g)?,
-            HashBuilderValueKind::Hash => B256::arbitrary(g)?.to_vec(),
+            HashBuilderValueKind::Hash => alloy_primitives::B256::arbitrary(g)?.to_vec(),
         };
         Ok(Self { buf, kind })
     }
@@ -71,10 +73,9 @@ impl HashBuilderValue {
     pub fn as_ref(&self) -> HashBuilderValueRef<'_> {
         match self.kind {
             HashBuilderValueKind::Bytes => HashBuilderValueRef::Bytes(&self.buf),
-            HashBuilderValueKind::Hash => {
-                debug_assert_eq!(self.buf.len(), 32);
-                HashBuilderValueRef::Hash(unsafe { self.buf[..].try_into().unwrap_unchecked() })
-            }
+            HashBuilderValueKind::Hash => HashBuilderValueRef::Hash(unsafe {
+                RlpNode::from_raw_rlp(&self.buf[..]).unwrap_unchecked()
+            }),
         }
     }
 
@@ -122,13 +123,13 @@ pub enum HashBuilderValueRef<'a> {
     /// Value of the leaf node.
     Bytes(&'a [u8]),
     /// Hash of adjacent nodes.
-    Hash(&'a B256),
+    Hash(RlpNode),
 }
 
-impl<'a> HashBuilderValueRef<'a> {
+impl HashBuilderValueRef<'_> {
     /// Returns the value as a slice.
-    pub const fn as_slice(&self) -> &'a [u8] {
-        match *self {
+    pub fn as_slice(&self) -> &[u8] {
+        match self {
             HashBuilderValueRef::Bytes(bytes) => bytes,
             HashBuilderValueRef::Hash(hash) => hash.as_slice(),
         }

src/nodes/branch.rs

@@ -169,11 +169,11 @@ impl<'a> BranchNodeRef<'a> {
     /// Given the hash mask of children, return an iterator over stack items
     /// that match the mask.
     #[inline]
-    pub fn child_hashes(&self, hash_mask: TrieMask) -> impl Iterator<Item = B256> + '_ {
+    pub fn child_hashes(&self, hash_mask: TrieMask) -> impl Iterator<Item = RlpNode> + '_ {
         self.children()
             .filter_map(|(i, c)| c.map(|c| (i, c)))
             .filter(move |(index, _)| hash_mask.is_bit_set(*index))
-            .map(|(_, child)| B256::from_slice(&child[1..]))
+            .map(|(_, child)| child.clone())
     }
 
     /// RLP-encodes the node and returns either `rlp(node)` or `rlp(keccak(rlp(node)))`.
@@ -279,7 +279,7 @@ pub struct BranchNodeCompact {
     pub hash_mask: TrieMask,
     /// Collection of hashes associated with the children of the branch node.
     /// Each child hash is calculated by hashing two consecutive sub-branch roots.
-    pub hashes: Arc<Vec<B256>>,
+    pub hashes: Arc<Vec<RlpNode>>,
     /// An optional root hash of the subtree rooted at this branch node.
     pub root_hash: Option<B256>,
 }
@@ -290,7 +290,7 @@ impl BranchNodeCompact {
         state_mask: impl Into<TrieMask>,
         tree_mask: impl Into<TrieMask>,
         hash_mask: impl Into<TrieMask>,
-        hashes: Vec<B256>,
+        hashes: Vec<RlpNode>,
         root_hash: Option<B256>,
     ) -> Self {
         let (state_mask, tree_mask, hash_mask) =
@@ -308,10 +308,10 @@ impl BranchNodeCompact {
     }
 
     /// Returns the hash associated with the given nibble.
-    pub fn hash_for_nibble(&self, nibble: u8) -> B256 {
+    pub fn hash_for_nibble(&self, nibble: u8) -> RlpNode {
         let mask = *TrieMask::from_nibble(nibble) - 1;
         let index = (*self.hash_mask & mask).count_ones();
-        self.hashes[index as usize]
+        self.hashes[index as usize].clone()
     }
 }
 

src/nodes/rlp.rs

@@ -6,7 +6,7 @@ use core::fmt;
 const MAX: usize = 33;
 
 /// An RLP-encoded node.
-#[derive(Clone, Default, PartialEq, Eq)]
+#[derive(Clone, Default, PartialEq, Eq, PartialOrd, Ord)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub struct RlpNode(ArrayVec<u8, MAX>);