Merge branch (#8)

* Removed unused code, fix warnings, code path tweaking * Merkle Proof implementation(inclusion&exclusion), compatible with original one
nervosnetwork · Nov 10, 2022 · c4e8b90 · c4e8b90
1 parent 6113d1c
commit c4e8b90
Show file tree

Hide file tree

Showing 11 changed files with 301 additions and 119 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -35,7 +35,7 @@ rand = "0.8"
 hex = "0.4.3"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
-anyhow = "1.0"
+anyhow = "1.0.65"
 
 [[bench]]
 name = "smt_benchmark"

diff --git a/benches/store_counter_benchmark.rs b/benches/store_counter_benchmark.rs
@@ -6,7 +6,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use rand::{thread_rng, Rng};
 use sparse_merkle_tree::{
     blake2b::Blake2bHasher,
-    branch::{BranchKey, BranchNode},
+    tree::{BranchKey, BranchNode},
     default_store::DefaultStore,
     error::Error,
     traits::{StoreReadOps, StoreWriteOps},

diff --git a/src/branch.rs b/src/branch.rs
diff --git a/src/default_store.rs b/src/default_store.rs
@@ -1,8 +1,8 @@
 use crate::{
-    branch::{BranchKey, BranchNode},
     collections,
     error::Error,
     traits::{StoreReadOps, StoreWriteOps},
+    tree::{BranchKey, BranchNode},
     H256,
 };
 

diff --git a/src/lib.rs b/src/lib.rs
@@ -62,7 +62,6 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
 pub mod blake2b;
-pub mod branch;
 #[cfg(feature = "smtc")]
 pub mod ckb_smt;
 pub mod default_store;
@@ -84,6 +83,7 @@ pub use h256::H256;
 pub use merkle_proof::{CompiledMerkleProof, MerkleProof};
 #[cfg(not(feature = "trie"))]
 pub use tree::SparseMerkleTree;
+pub use tree::{BranchKey, BranchNode};
 #[cfg(feature = "trie")]
 pub use trie_tree::SparseMerkleTree;
 

diff --git a/src/merge.rs b/src/merge.rs
@@ -72,16 +72,27 @@ impl MergeValue {
                 hasher.finish()
             }
             #[cfg(feature = "trie")]
-            MergeValue::ShortCut { key, value, height } => {
+            MergeValue::ShortCut {
+                key: _,
+                value,
+                height,
+            } => {
                 // try keep hash same with MergeWithZero
                 if value.is_zero() {
                     return H256::zero();
                 }
-                self.into_merge_with_zero::<H>().hash::<H>()
+                if *height == 0 {
+                    *value
+                } else {
+                    self.into_merge_with_zero::<H>().hash::<H>()
+                }
             }
         }
     }
 
+
+    /// Helper function for get base_node
+    /// When call with MergeValue::Value(v), it would be v
     pub fn base_node<H: Hasher + Default>(&self) -> H256 {
         match self {
             MergeValue::ShortCut {
@@ -101,7 +112,10 @@ impl MergeValue {
         }
     }
 
-    fn into_merge_with_zero<H: Hasher + Default>(&self) -> MergeValue {
+
+    /// Helper function for Shortcut node
+    /// Transform it into a MergeWithZero node
+    pub fn into_merge_with_zero<H: Hasher + Default>(&self) -> MergeValue {
         match self {
             MergeValue::ShortCut { key, value, height } => {
                 let base_key = key.parent_path(0);
@@ -117,13 +131,8 @@ impl MergeValue {
                     zero_bits,
                     zero_count: *height,
                 }
-            }
-            MergeValue::MergeWithZero {
-                base_node: _,
-                zero_bits: _,
-                zero_count: _,
-            } => self.clone(),
-            MergeValue::Value(_) => {
+            },
+            _ => {
                 unreachable!();
             }
         }
@@ -170,7 +179,7 @@ pub fn merge<H: Hasher + Default>(
     MergeValue::Value(hasher.finish())
 }
 
-fn merge_with_zero<H: Hasher + Default>(
+pub fn merge_with_zero<H: Hasher + Default>(
     height: u8,
     node_key: &H256,
     value: &MergeValue,

diff --git a/src/merkle_proof.rs b/src/merkle_proof.rs
@@ -104,12 +104,7 @@ impl MerkleProof {
                                 buffer.extend_from_slice(zero_bits.as_slice());
                                 (Some(0x51), Some(buffer))
                             }
-                            #[cfg(feature = "trie")]
-                            MergeValue::ShortCut {
-                                key: _,
-                                value: val,
-                                height: _,
-                            } => (Some(0x52), Some(val.as_slice().to_vec())),
+                            _ => unreachable!(),
                         }
                     } else {
                         zero_count += 1;
@@ -507,12 +502,7 @@ impl CompiledMerkleProof {
                                 sub_proof.extend(zero_bits.as_slice());
                                 is_last_merge_zero = false;
                             }
-                            #[cfg(feature = "trie")]
-                            MergeValue::ShortCut {
-                                key: _,
-                                value: _,
-                                height: _,
-                            } => {}
+                            _ => {}
                         };
                     }
                 }

diff --git a/src/tests/tree.rs b/src/tests/tree.rs
@@ -264,6 +264,7 @@ fn test_update(key: H256, value: H256) {
     assert_eq!(tree.get(&key), Ok(value));
 }
 
+#[cfg(not(feature = "trie"))]
 fn test_update_tree_store(key: H256, value: H256, value2: H256) {
     const EXPECTED_LEAVES_LEN: usize = 1;
 
@@ -389,6 +390,7 @@ proptest! {
         test_update(key.into(), value.into());
     }
 
+    #[cfg(not(feature = "trie"))]
     #[test]
     fn test_random_update_tree_store(key: [u8;32], value: [u8;32], value2: [u8;32]) {
         test_update_tree_store(key.into(), value.into(), value2.into());
@@ -721,16 +723,10 @@ fn test_trie_broken_sample() {
 
     let mut pairs = keys.zip(values).collect::<Vec<_>>();
     let smt = new_smt(pairs.clone());
-    let base_root = *smt.root();
     let base_branches = smt.store().branches_map();
-
     pairs.reverse();
     let smt = new_smt(pairs.clone());
-    let current_root = *smt.root();
-
     let current_branches = smt.store().branches_map();
-
-    //assert_eq!(base_root, current_root);
     assert_eq!(base_branches, current_branches);
 }
 

diff --git a/src/traits.rs b/src/traits.rs
@@ -1,6 +1,6 @@
 use crate::{
-    branch::{BranchKey, BranchNode},
     error::Error,
+    tree::{BranchKey, BranchNode},
     H256,
 };
 

diff --git a/src/tree.rs b/src/tree.rs
@@ -1,5 +1,4 @@
 use crate::{
-    branch::*,
     collections::VecDeque,
     error::{Error, Result},
     merge::{merge, MergeValue},
@@ -8,7 +7,56 @@ use crate::{
     vec::Vec,
     H256, MAX_STACK_SIZE,
 };
+use core::cmp::Ordering;
 use core::marker::PhantomData;
+/// The branch key
+#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+pub struct BranchKey {
+    pub height: u8,
+    pub node_key: H256,
+}
+
+impl BranchKey {
+    pub fn new(height: u8, node_key: H256) -> BranchKey {
+        BranchKey { height, node_key }
+    }
+}
+
+impl PartialOrd for BranchKey {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+impl Ord for BranchKey {
+    fn cmp(&self, other: &Self) -> Ordering {
+        match self.height.cmp(&other.height) {
+            Ordering::Equal => self.node_key.cmp(&other.node_key),
+            ordering => ordering,
+        }
+    }
+}
+
+/// A branch in the SMT
+#[derive(Debug, Eq, PartialEq, Clone)]
+pub struct BranchNode {
+    pub left: MergeValue,
+    pub right: MergeValue,
+}
+
+impl BranchNode {
+    /// Create a new empty branch
+    pub fn new_empty() -> BranchNode {
+        BranchNode {
+            left: MergeValue::zero(),
+            right: MergeValue::zero(),
+        }
+    }
+
+    /// Determine whether a node did not store any value
+    pub fn is_empty(&self) -> bool {
+        self.left.is_zero() && self.right.is_zero()
+    }
+}
 
 /// Sparse merkle tree
 #[derive(Default, Debug)]