TrieTree implementation

Cargo.toml

@@ -21,6 +21,8 @@ default = ["std"]
 std = []
 # SMT implemented in C
 smtc = []
+# A storage optimized SMT implemented in trie (https://ouvrard-pierre-alain.medium.com/sparse-merkle-tree-86e6e2fc26da)
+trie = []
 
 [dependencies]
 cfg-if = "0.1"
@@ -33,11 +35,15 @@ 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"
 harness = false
 
+[[bench]]
+name = "store_counter_benchmark"
+harness = false
+
 [build-dependencies]
 cc = "1.0"

Makefile

@@ -1,12 +1,21 @@
-default: fmt clippy test bench-test check test-c-impl test-cxx-build
+default: fmt clippy clippy-trie test test-trie bench-test bench-test-trie check test-c-impl test-cxx-build
 
 test:
+	cargo test --all --features std,smtc
+
+test-trie:
 	cargo test --all --all-features
 
 bench-test:
 	cargo bench -- --test
 
+bench-test-trie:
+	cargo bench --features trie -- --test
+
 clippy:
+	cargo clippy  --all --features std,smtc --all-targets
+
+clippy-trie:
 	cargo clippy  --all --all-features --all-targets
 
 fmt:

benches/smt_benchmark.rs

@@ -4,7 +4,7 @@ extern crate criterion;
 use criterion::Criterion;
 use rand::{thread_rng, Rng};
 use sparse_merkle_tree::{
-    blake2b::Blake2bHasher, default_store::DefaultStore, tree::SparseMerkleTree, H256,
+    blake2b::Blake2bHasher, default_store::DefaultStore, SparseMerkleTree, H256,
 };
 
 const TARGET_LEAVES_COUNT: usize = 20;

benches/store_counter_benchmark.rs

@@ -0,0 +1,120 @@
+#[macro_use]
+extern crate criterion;
+
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use rand::{thread_rng, Rng};
+use sparse_merkle_tree::{
+    blake2b::Blake2bHasher,
+    default_store::DefaultStore,
+    error::Error,
+    traits::{StoreReadOps, StoreWriteOps},
+    BranchKey, BranchNode, SparseMerkleTree, H256,
+};
+
+#[derive(Debug, Default)]
+struct DefaultStoreWithCounters<V> {
+    store: DefaultStore<V>,
+    counters: Counters,
+}
+
+#[derive(Debug, Default)]
+struct Counters {
+    get_branch_counter: AtomicUsize,
+    get_leaf_counter: AtomicUsize,
+    insert_branch_counter: AtomicUsize,
+    insert_leaf_counter: AtomicUsize,
+    remove_branch_counter: AtomicUsize,
+    remove_leaf_counter: AtomicUsize,
+}
+
+impl<V: Clone> StoreReadOps<V> for DefaultStoreWithCounters<V> {
+    fn get_branch(&self, branch_key: &BranchKey) -> Result<Option<BranchNode>, Error> {
+        self.counters
+            .get_branch_counter
+            .fetch_add(1, Ordering::SeqCst);
+        self.store.get_branch(branch_key)
+    }
+    fn get_leaf(&self, leaf_key: &H256) -> Result<Option<V>, Error> {
+        self.counters
+            .get_leaf_counter
+            .fetch_add(1, Ordering::SeqCst);
+        self.store.get_leaf(leaf_key)
+    }
+}
+
+impl<V> StoreWriteOps<V> for DefaultStoreWithCounters<V> {
+    fn insert_branch(&mut self, branch_key: BranchKey, branch: BranchNode) -> Result<(), Error> {
+        self.counters
+            .insert_branch_counter
+            .fetch_add(1, Ordering::SeqCst);
+        self.store.insert_branch(branch_key, branch)
+    }
+    fn insert_leaf(&mut self, leaf_key: H256, leaf: V) -> Result<(), Error> {
+        self.counters
+            .insert_leaf_counter
+            .fetch_add(1, Ordering::SeqCst);
+        self.store.insert_leaf(leaf_key, leaf)
+    }
+    fn remove_branch(&mut self, branch_key: &BranchKey) -> Result<(), Error> {
+        self.counters
+            .remove_branch_counter
+            .fetch_add(1, Ordering::SeqCst);
+        self.store.remove_branch(branch_key)
+    }
+    fn remove_leaf(&mut self, leaf_key: &H256) -> Result<(), Error> {
+        self.counters
+            .remove_leaf_counter
+            .fetch_add(1, Ordering::SeqCst);
+        self.store.remove_leaf(leaf_key)
+    }
+}
+
+#[allow(clippy::upper_case_acronyms)]
+type SMT = SparseMerkleTree<Blake2bHasher, H256, DefaultStoreWithCounters<H256>>;
+
+fn random_h256(rng: &mut impl Rng) -> H256 {
+    let mut buf = [0u8; 32];
+    rng.fill(&mut buf);
+    buf.into()
+}
+
+fn random_smt(update_count: usize, rng: &mut impl Rng) {
+    let mut smt = SMT::default();
+    let mut keys = Vec::with_capacity(update_count);
+    for _ in 0..update_count {
+        let key = random_h256(rng);
+        let value = random_h256(rng);
+        smt.update(key, value).unwrap();
+        keys.push(key);
+    }
+    println!(
+        "random update {} keys, store counters: {:?}",
+        update_count,
+        smt.store().counters
+    );
+}
+
+fn random_smt_update_all(update_count: usize, rng: &mut impl Rng) {
+    let mut smt = SMT::default();
+    let mut kvs = Vec::with_capacity(update_count);
+    for _ in 0..update_count {
+        let key = random_h256(rng);
+        let value = random_h256(rng);
+        kvs.push((key, value));
+    }
+    smt.update_all(kvs).unwrap();
+    println!(
+        "random update_all {} keys, store counters: {:?}",
+        update_count,
+        smt.store().counters
+    );
+}
+
+fn main() {
+    let mut rng = thread_rng();
+    random_smt(100, &mut rng);
+    random_smt(10000, &mut rng);
+    random_smt_update_all(100, &mut rng);
+    random_smt_update_all(10000, &mut rng);
+}

src/lib.rs

@@ -72,13 +72,19 @@ pub mod merkle_proof;
 #[cfg(test)]
 mod tests;
 pub mod traits;
-pub mod tree;
+mod tree;
+#[cfg(feature = "trie")]
+mod trie_tree;
 
 #[cfg(feature = "smtc")]
 pub use ckb_smt::{SMTBuilder, SMT};
 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;
 
 /// Expected path size: log2(256) * 2, used for hint vector capacity
 pub const EXPECTED_PATH_SIZE: usize = 16;

src/merge.rs

@@ -12,6 +12,12 @@ pub enum MergeValue {
         zero_bits: H256,
         zero_count: u8,
     },
+    #[cfg(feature = "trie")]
+    ShortCut {
+        key: H256,
+        value: H256,
+        height: u8,
+    },
 }
 
 impl MergeValue {
@@ -24,10 +30,26 @@ impl MergeValue {
     }
 
     pub fn is_zero(&self) -> bool {
-        if let MergeValue::Value(v) = self {
-            return v.is_zero();
+        match self {
+            MergeValue::Value(v) => v.is_zero(),
+            MergeValue::MergeWithZero { .. } => false,
+            #[cfg(feature = "trie")]
+            MergeValue::ShortCut { .. } => false,
+        }
+    }
+
+    #[cfg(feature = "trie")]
+    pub fn shortcut_or_value(key: H256, value: H256, height: u8) -> Self {
+        if height == 0 || value.is_zero() {
+            MergeValue::Value(value)
+        } else {
+            MergeValue::ShortCut { key, value, height }
         }
-        false
+    }
+
+    #[cfg(feature = "trie")]
+    pub fn is_shortcut(&self) -> bool {
+        matches!(self, MergeValue::ShortCut { .. })
     }
 
     pub fn hash<H: Hasher + Default>(&self) -> H256 {
@@ -45,6 +67,34 @@ impl MergeValue {
                 hasher.write_byte(*zero_count);
                 hasher.finish()
             }
+            #[cfg(feature = "trie")]
+            MergeValue::ShortCut { key, value, height } => {
+                into_merge_value::<H>(*key, *value, *height).hash::<H>()
+            }
+        }
+    }
+}
+
+/// Helper function for Shortcut node
+/// Transform it into a MergeValue or MergeWithZero node
+#[cfg(feature = "trie")]
+pub fn into_merge_value<H: Hasher + Default>(key: H256, value: H256, height: u8) -> MergeValue {
+    // try keep hash same with MergeWithZero
+    if value.is_zero() {
+        MergeValue::from_h256(H256::zero())
+    } else {
+        let base_key = key.parent_path(0);
+        let base_node = hash_base_node::<H>(0, &base_key, &value);
+        let mut zero_bits = key;
+        for i in height..=core::u8::MAX {
+            if key.get_bit(i) {
+                zero_bits.clear_bit(i);
+            }
+        }
+        MergeValue::MergeWithZero {
+            base_node,
+            zero_bits,
+            zero_count: height,
         }
     }
 }
@@ -89,7 +139,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,
@@ -123,5 +173,23 @@ fn merge_with_zero<H: Hasher + Default>(
                 zero_count: zero_count.wrapping_add(1),
             }
         }
+        #[cfg(feature = "trie")]
+        MergeValue::ShortCut { key, value, .. } => {
+            if height == core::u8::MAX {
+                let base_key = key.parent_path(0);
+                let base_node = hash_base_node::<H>(0, &base_key, value);
+                MergeValue::MergeWithZero {
+                    base_node,
+                    zero_bits: *key,
+                    zero_count: 0,
+                }
+            } else {
+                MergeValue::ShortCut {
+                    key: *key,
+                    value: *value,
+                    height: height + 1,
+                }
+            }
+        }
     }
 }

src/merkle_proof.rs

@@ -104,6 +104,8 @@ impl MerkleProof {
                                 buffer.extend_from_slice(zero_bits.as_slice());
                                 (Some(0x51), Some(buffer))
                             }
+                            #[cfg(feature = "trie")]
+                            _ => unreachable!(),
                         }
                     } else {
                         zero_count += 1;
@@ -501,6 +503,8 @@ impl CompiledMerkleProof {
                                 sub_proof.extend(zero_bits.as_slice());
                                 is_last_merge_zero = false;
                             }
+                            #[cfg(feature = "trie")]
+                            _ => {}
                         };
                     }
                 }