feat: add test case

Cargo.toml

@@ -42,7 +42,6 @@ hash-db = "0.15"
 plain_hasher = "0.2"
 triehash = "0.8.4"
 criterion = "0.5"
-tracing-subscriber = "0.3"
 
 [features]
 default = ["std"]

src/hash_builder/mod.rs

@@ -343,8 +343,8 @@ 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.saturating_sub(1)].is_empty();
+        let store_in_db_trie = !self.tree_masks[len.saturating_sub(1)].is_empty()
+            || !self.hash_masks[len.saturating_sub(1)].is_empty();
         if store_in_db_trie || len == 0 {
             if len > 0 {
                 let parent_index = len - 1;
@@ -414,40 +414,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]
@@ -592,10 +594,6 @@ mod tests {
 
     #[test]
     fn test_updates_root() {
-        // let subscriber = tracing_subscriber::FmtSubscriber::builder()
-        //     .with_max_level(tracing::Level::TRACE)
-        //     .finish();
-        // tracing::subscriber::set_global_default(subscriber).unwrap();
         let mut hb = HashBuilder::default().with_updates(true);
         let account = Vec::new();
 
@@ -623,4 +621,39 @@ mod tests {
         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);
+    }
 }