fix root updates bug in hash_builder

.github/workflows/ci.yml

@@ -16,7 +16,7 @@ jobs:
         strategy:
             fail-fast: false
             matrix:
-                rust: ["stable", "beta", "nightly", "1.66"] # MSRV
+                rust: ["stable", "beta", "nightly", "1.81"] # MSRV
                 flags: ["--no-default-features", "", "--all-features"]
         steps:
             - uses: actions/checkout@v3
@@ -29,7 +29,7 @@ jobs:
             - name: build
               run: cargo build --workspace ${{ matrix.flags }}
             - name: test
-              if: ${{ matrix.rust != '1.66' }} # MSRV
+              if: ${{ matrix.rust != '1.81' }} # MSRV
               run: cargo test --workspace ${{ matrix.flags }}
 
     miri:

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "alloy-trie"
-version = "0.4.1"
+version = "0.6.0"
 authors = [
     "rkrasiuk <rokrassyuk@gmail.com>",
     "gakonst <me@gakonst.com>",
@@ -11,7 +11,7 @@ Fast Merkle-Patricia Trie (MPT) state root calculator
 and proof generator for prefix-sorted nibbles
 """
 edition = "2021"
-rust-version = "1.66"
+rust-version = "1.81"
 license = "MIT OR Apache-2.0"
 categories = ["data-structures", "no-std"]
 keywords = ["nibbles", "trie", "mpt", "merkle", "ethereum"]
@@ -20,8 +20,8 @@ repository = "https://github.com/alloy-rs/trie"
 exclude = [".github/", "deny.toml", "release.toml", "rustfmt.toml"]
 
 [dependencies]
-alloy-primitives = { version = "0.7", default-features = false, features = ["rlp"] }
-alloy-rlp = { version = "0.3", default-features = false, features = ["derive"] }
+alloy-primitives = { version = "0.8.5", default-features = false, features = ["rlp"] }
+alloy-rlp = { version = "0.3.8", default-features = false, features = ["derive"] }
 derive_more = "0.99"
 hashbrown = { version = "0.14", features = ["ahash", "inline-more"] }
 nybbles = { version = "0.2", default-features = false }

src/hash_builder/mod.rs

@@ -8,14 +8,16 @@ use super::{
 use crate::HashMap;
 use alloy_primitives::{hex, keccak256, Bytes, B256};
 use core::cmp;
-use tracing::trace;
+use tracing::{error, trace};
 
 #[allow(unused_imports)]
 use alloc::{collections::BTreeMap, vec::Vec};
 
 mod value;
 pub use value::HashBuilderValue;
 
+/// HashBuilder
+///
 /// A component used to construct the root hash of the trie. The primary purpose of a Hash Builder
 /// is to build the Merkle proof that is essential for verifying the integrity and authenticity of
 /// the trie's contents. It achieves this by constructing the root hash from the hashes of child
@@ -111,7 +113,14 @@ impl HashBuilder {
 
     /// Adds a new leaf element and its value to the trie hash builder.
     pub fn add_leaf(&mut self, key: Nibbles, value: &[u8]) {
-        assert!(key > self.key);
+        if key <= self.key {
+            error!(
+                "add_leaf invalid call, just ignore it. key = {:?}, self.key = {:?}",
+                key, self.key
+            );
+            return;
+        }
+        // assert!(key > self.key, "key: {:?}, self.key: {:?}", key, self.key);
         if !self.key.is_empty() {
             self.update(&key);
         }
@@ -120,7 +129,14 @@ 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) {
-        assert!(key > self.key || (self.key.is_empty() && key.is_empty()));
+        if !(key > self.key || (self.key.is_empty() && key.is_empty())) {
+            error!(
+                "add_branch invalid call, just ignore it. key = {:?}, self.key = {:?}",
+                key, self.key
+            );
+            return;
+        }
+        // assert!(key > self.key || (self.key.is_empty() && key.is_empty()));
         if !self.key.is_empty() {
             self.update(&key);
         } else if key.is_empty() {
@@ -343,25 +359,24 @@ impl HashBuilder {
             self.hash_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();
+        let store_in_db_trie = len == 0
+            || !self.tree_masks[len.saturating_sub(1)].is_empty()
+            || !self.hash_masks[len.saturating_sub(1)].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]);
             }
 
-            let mut n = BranchNodeCompact::new(
+            let n = BranchNodeCompact::new(
                 self.groups[len],
                 self.tree_masks[len],
                 self.hash_masks[len],
                 children,
-                None,
+                Some(self.current_root()),
             );
 
-            if len == 0 {
-                n.root_hash = Some(self.current_root());
-            }
-
             // Send it over to the provided channel which will handle it on the
             // other side of the HashBuilder
             trace!(target: "trie::hash_builder", node = ?n, "intermediate node");
@@ -413,40 +428,42 @@ mod tests {
     // Hashes the keys, RLP encodes the values, compares the trie builder with the upstream root.
     fn assert_hashed_trie_root<'a, I, K>(iter: I)
     where
-        I: Iterator<Item = (K, &'a U256)>,
-        K: AsRef<[u8]> + Ord,
+        I: Iterator<Item = (K, &'a U256)> + Clone,
+        K: AsRef<[u8]> + Ord + Clone,
     {
-        let hashed = iter
-            .map(|(k, v)| (keccak256(k.as_ref()), alloy_rlp::encode(v).to_vec()))
-            // Collect into a btree map to sort the data
-            .collect::<BTreeMap<_, _>>();
-
-        let mut hb = HashBuilder::default();
-
-        hashed.iter().for_each(|(key, val)| {
-            let nibbles = Nibbles::unpack(key);
-            hb.add_leaf(nibbles, val);
-        });
-
-        assert_eq!(hb.root(), triehash_trie_root(&hashed));
+        let iter = iter.map(|(k, v)| (keccak256(k.as_ref()), alloy_rlp::encode(v).to_vec()));
+        let mut hb = build_hash_builder(iter.clone(), None);
+        assert_eq!(hb.root(), triehash_trie_root(iter));
     }
 
     // No hashing involved
     fn assert_trie_root<I, K, V>(iter: I)
+    where
+        I: Iterator<Item = (K, V)> + Clone,
+        K: AsRef<[u8]> + Ord + Clone,
+        V: AsRef<[u8]> + Clone,
+    {
+        let mut hb = build_hash_builder(iter.clone(), None);
+        assert_eq!(hb.root(), triehash_trie_root(iter));
+    }
+
+    fn build_hash_builder<I, K, V>(iter: I, proof_retainer: Option<ProofRetainer>) -> HashBuilder
     where
         I: Iterator<Item = (K, V)>,
         K: AsRef<[u8]> + Ord,
         V: AsRef<[u8]>,
     {
-        let mut hb = HashBuilder::default();
+        let mut hb = HashBuilder::default().with_updates(true);
+        if let Some(retainer) = proof_retainer {
+            hb = hb.with_proof_retainer(retainer)
+        }
 
         let data = iter.collect::<BTreeMap<_, _>>();
         data.iter().for_each(|(key, val)| {
             let nibbles = Nibbles::unpack(key);
             hb.add_leaf(nibbles, val.as_ref());
         });
-
-        assert_eq!(hb.root(), triehash_trie_root(data));
+        hb
     }
 
     #[test]
@@ -518,7 +535,7 @@ mod tests {
 
         let update = updates.get(&Nibbles::from_nibbles_unchecked(hex!("01"))).unwrap();
         assert_eq!(update.state_mask, TrieMask::new(0b1111)); // 1st nibble: 0, 1, 2, 3
-        assert_eq!(update.tree_mask, TrieMask::new(0));
+        assert_eq!(update.tree_mask, TrieMask::new(6)); // in the 1st nibble, the ones with 1 and 2 are branches. value:0000000000000110
         assert_eq!(update.hash_mask, TrieMask::new(6)); // in the 1st nibble, the ones with 1 and 2 are branches with `hashes`
         assert_eq!(update.hashes.len(), 2); // calculated while the builder is running
 
@@ -588,4 +605,69 @@ mod tests {
         assert_eq!(hb.root(), expected);
         assert_eq!(hb2.root(), expected);
     }
+
+    #[test]
+    fn test_updates_root() {
+        let mut hb = HashBuilder::default().with_updates(true);
+        let account = Vec::new();
+
+        let mut key = Nibbles::unpack(hex!(
+            "a711355ec1c8f7e26bb3ccbcb0b75d870d15846c0b98e5cc452db46c37faea40"
+        ));
+        hb.add_leaf(key, account.as_ref());
+
+        key = Nibbles::unpack(hex!(
+            "a77d337781e762f3577784bab7491fcc43e291ce5a356b9bc517ac52eed3a37a"
+        ));
+        hb.add_leaf(key, account.as_ref());
+
+        key = Nibbles::unpack(hex!(
+            "a77d397a32b8ab5eb4b043c65b1f00c93f517bc8883c5cd31baf8e8a279475e3"
+        ));
+        hb.add_leaf(key, account.as_ref());
+
+        key = Nibbles::unpack(hex!(
+            "a7f936599f93b769acf90c7178fd2ddcac1b5b4bc9949ee5a04b7e0823c2446e"
+        ));
+        hb.add_leaf(key, account.as_ref());
+
+        let _root = hb.root();
+        let (_, updates) = hb.split();
+        assert!(!updates.is_empty());
+    }
+
+    /// Test the tree handling top branch edge case.
+    #[test]
+    fn test_top_branch_logic() {
+        let default_leaf = "hello".as_bytes();
+        // mpt tree like(B = branch node, E = ext node, L = leaf node):
+        // 0[B] -> 0[E] -> 0[B] -> 0[L]
+        //                 2[B] -> 0[L]
+        // 1[B] -> 000[L]
+        // 2[B] -> 0[B] -> 00[L]
+        //         1[B] -> 1[B] -> 1[L]
+        //                 2[B] -> 2[B] -> ()[L]
+        //                         3[B] -> ()[]
+        // 3[B] -> 00[E] -> 0[B] -> ()[L]
+        //               -> 1[B] -> ()[L]
+        let data = vec![
+            (hex!("0000").to_vec(), default_leaf.to_vec()),
+            (hex!("0020").to_vec(), default_leaf.to_vec()),
+            (hex!("1000").to_vec(), default_leaf.to_vec()),
+            (hex!("2000").to_vec(), default_leaf.to_vec()),
+            (hex!("2111").to_vec(), default_leaf.to_vec()),
+            (hex!("2122").to_vec(), default_leaf.to_vec()),
+            (hex!("2123").to_vec(), default_leaf.to_vec()),
+            (hex!("3000").to_vec(), default_leaf.to_vec()),
+            (hex!("3001").to_vec(), default_leaf.to_vec()),
+        ];
+
+        let mut hb = build_hash_builder(data.into_iter(), None);
+
+        // add empty succeeding as ending.
+        hb.root();
+        let (_, updates) = hb.split();
+        // according to the data graph, there's should be 6 branch nodes to be updated.
+        assert_eq!(updates.len(), 6);
+    }
 }

src/nodes/branch.rs

@@ -7,6 +7,8 @@ use nybbles::Nibbles;
 #[allow(unused_imports)]
 use alloc::vec::Vec;
 
+/// BranchNode.
+///
 /// A branch node in an Merkle Patricia Trie is a 17-element array consisting of 16 slots that
 /// correspond to each hexadecimal character and an additional slot for a value. We do exclude
 /// the node value since all paths have a fixed size.

src/nodes/extension.rs

@@ -6,6 +6,8 @@ use core::fmt;
 #[allow(unused_imports)]
 use alloc::vec::Vec;
 
+/// ExtensionNode.
+///
 /// An intermediate node that exists solely to compress the trie's paths. It contains a path segment
 /// (a shared prefix of keys) and a single child pointer. Essentially, an extension node can be
 /// thought of as a shortcut within the trie to reduce its overall depth.

src/nodes/leaf.rs

@@ -6,6 +6,8 @@ use core::fmt;
 #[allow(unused_imports)]
 use alloc::vec::Vec;
 
+/// LeafNode.
+///
 /// A leaf node represents the endpoint or terminal node in the trie. In other words, a leaf node is
 /// where actual values are stored.
 ///

src/nodes/mod.rs

@@ -174,6 +174,8 @@ pub(crate) fn unpack_path_to_nibbles(first: Option<u8>, rest: &[u8]) -> Nibbles
     Nibbles::from_vec_unchecked(nibbles)
 }
 
+/// encode_path_leaf
+///
 /// Encodes a given path leaf as a compact array of bytes, where each byte represents two
 /// "nibbles" (half-bytes or 4 bits) of the original hex data, along with additional information
 /// about the leaf itself.

src/proof/verify.rs

@@ -10,6 +10,8 @@ use alloy_primitives::{Bytes, B256};
 use alloy_rlp::Decodable;
 use nybbles::Nibbles;
 
+/// verify_proof
+///
 /// Verify the proof for given key value pair against the provided state root.
 /// The expected node value can be either [Some] if it's expected to be present
 /// in the tree or [None] if this is an exclusion proof.