Skip to content

Commit 2641f59

Browse files
quakequake
authored
Merge pull request #34 from code-monad/master
TrieTree implementation
2 parents 671eed3 + da46840 commit 2641f59

File tree

10 files changed

+778
-10
lines changed

10 files changed

+778
-10
lines changed

‎Cargo.toml‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@ default = ["std"]
2121
std = []
2222
# SMT implemented in C
2323
smtc = []
24+
# A storage optimized SMT implemented in trie (https://ouvrard-pierre-alain.medium.com/sparse-merkle-tree-86e6e2fc26da)
25+
trie = []
2426

2527
[dependencies]
2628
cfg-if = "0.1"
@@ -33,11 +35,15 @@ rand = "0.8"
3335
hex = "0.4.3"
3436
serde = { version = "1.0", features = ["derive"] }
3537
serde_json = "1.0"
36-
anyhow = "1.0"
38+
anyhow = "1.0.65"
3739

3840
[[bench]]
3941
name = "smt_benchmark"
4042
harness = false
4143

44+
[[bench]]
45+
name = "store_counter_benchmark"
46+
harness = false
47+
4248
[build-dependencies]
4349
cc = "1.0"

‎Makefile‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,21 @@
1-
default: fmt clippy test bench-test check test-c-impl test-cxx-build
1+
default: fmt clippy clippy-trie test test-trie bench-test bench-test-trie check test-c-impl test-cxx-build
22

33
test:
4+
cargo test --all --features std,smtc
5+
6+
test-trie:
47
cargo test --all --all-features
58

69
bench-test:
710
cargo bench -- --test
811

12+
bench-test-trie:
13+
cargo bench --features trie -- --test
14+
915
clippy:
16+
cargo clippy --all --features std,smtc --all-targets
17+
18+
clippy-trie:
1019
cargo clippy --all --all-features --all-targets
1120

1221
fmt:

‎benches/smt_benchmark.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ extern crate criterion;
44
use criterion::Criterion;
55
use rand::{thread_rng, Rng};
66
use sparse_merkle_tree::{
7-
blake2b::Blake2bHasher, default_store::DefaultStore, tree::SparseMerkleTree, H256,
7+
blake2b::Blake2bHasher, default_store::DefaultStore, SparseMerkleTree, H256,
88
};
99

1010
const TARGET_LEAVES_COUNT: usize = 20;

‎benches/store_counter_benchmark.rs‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
#[macro_use]
2+
extern crate criterion;
3+
4+
use std::sync::atomic::{AtomicUsize, Ordering};
5+
6+
use rand::{thread_rng, Rng};
7+
use sparse_merkle_tree::{
8+
blake2b::Blake2bHasher,
9+
default_store::DefaultStore,
10+
error::Error,
11+
traits::{StoreReadOps, StoreWriteOps},
12+
BranchKey, BranchNode, SparseMerkleTree, H256,
13+
};
14+
15+
#[derive(Debug, Default)]
16+
struct DefaultStoreWithCounters<V> {
17+
store: DefaultStore<V>,
18+
counters: Counters,
19+
}
20+
21+
#[derive(Debug, Default)]
22+
struct Counters {
23+
get_branch_counter: AtomicUsize,
24+
get_leaf_counter: AtomicUsize,
25+
insert_branch_counter: AtomicUsize,
26+
insert_leaf_counter: AtomicUsize,
27+
remove_branch_counter: AtomicUsize,
28+
remove_leaf_counter: AtomicUsize,
29+
}
30+
31+
impl<V: Clone> StoreReadOps<V> for DefaultStoreWithCounters<V> {
32+
fn get_branch(&self, branch_key: &BranchKey) -> Result<Option<BranchNode>, Error> {
33+
self.counters
34+
.get_branch_counter
35+
.fetch_add(1, Ordering::SeqCst);
36+
self.store.get_branch(branch_key)
37+
}
38+
fn get_leaf(&self, leaf_key: &H256) -> Result<Option<V>, Error> {
39+
self.counters
40+
.get_leaf_counter
41+
.fetch_add(1, Ordering::SeqCst);
42+
self.store.get_leaf(leaf_key)
43+
}
44+
}
45+
46+
impl<V> StoreWriteOps<V> for DefaultStoreWithCounters<V> {
47+
fn insert_branch(&mut self, branch_key: BranchKey, branch: BranchNode) -> Result<(), Error> {
48+
self.counters
49+
.insert_branch_counter
50+
.fetch_add(1, Ordering::SeqCst);
51+
self.store.insert_branch(branch_key, branch)
52+
}
53+
fn insert_leaf(&mut self, leaf_key: H256, leaf: V) -> Result<(), Error> {
54+
self.counters
55+
.insert_leaf_counter
56+
.fetch_add(1, Ordering::SeqCst);
57+
self.store.insert_leaf(leaf_key, leaf)
58+
}
59+
fn remove_branch(&mut self, branch_key: &BranchKey) -> Result<(), Error> {
60+
self.counters
61+
.remove_branch_counter
62+
.fetch_add(1, Ordering::SeqCst);
63+
self.store.remove_branch(branch_key)
64+
}
65+
fn remove_leaf(&mut self, leaf_key: &H256) -> Result<(), Error> {
66+
self.counters
67+
.remove_leaf_counter
68+
.fetch_add(1, Ordering::SeqCst);
69+
self.store.remove_leaf(leaf_key)
70+
}
71+
}
72+
73+
#[allow(clippy::upper_case_acronyms)]
74+
type SMT = SparseMerkleTree<Blake2bHasher, H256, DefaultStoreWithCounters<H256>>;
75+
76+
fn random_h256(rng: &mut impl Rng) -> H256 {
77+
let mut buf = [0u8; 32];
78+
rng.fill(&mut buf);
79+
buf.into()
80+
}
81+
82+
fn random_smt(update_count: usize, rng: &mut impl Rng) {
83+
let mut smt = SMT::default();
84+
let mut keys = Vec::with_capacity(update_count);
85+
for _ in 0..update_count {
86+
let key = random_h256(rng);
87+
let value = random_h256(rng);
88+
smt.update(key, value).unwrap();
89+
keys.push(key);
90+
}
91+
println!(
92+
"random update {} keys, store counters: {:?}",
93+
update_count,
94+
smt.store().counters
95+
);
96+
}
97+
98+
fn random_smt_update_all(update_count: usize, rng: &mut impl Rng) {
99+
let mut smt = SMT::default();
100+
let mut kvs = Vec::with_capacity(update_count);
101+
for _ in 0..update_count {
102+
let key = random_h256(rng);
103+
let value = random_h256(rng);
104+
kvs.push((key, value));
105+
}
106+
smt.update_all(kvs).unwrap();
107+
println!(
108+
"random update_all {} keys, store counters: {:?}",
109+
update_count,
110+
smt.store().counters
111+
);
112+
}
113+
114+
fn main() {
115+
let mut rng = thread_rng();
116+
random_smt(100, &mut rng);
117+
random_smt(10000, &mut rng);
118+
random_smt_update_all(100, &mut rng);
119+
random_smt_update_all(10000, &mut rng);
120+
}

‎src/lib.rs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,13 +72,19 @@ pub mod merkle_proof;
7272
#[cfg(test)]
7373
mod tests;
7474
pub mod traits;
75-
pub mod tree;
75+
mod tree;
76+
#[cfg(feature = "trie")]
77+
mod trie_tree;
7678

7779
#[cfg(feature = "smtc")]
7880
pub use ckb_smt::{SMTBuilder, SMT};
7981
pub use h256::H256;
8082
pub use merkle_proof::{CompiledMerkleProof, MerkleProof};
83+
#[cfg(not(feature = "trie"))]
8184
pub use tree::SparseMerkleTree;
85+
pub use tree::{BranchKey, BranchNode};
86+
#[cfg(feature = "trie")]
87+
pub use trie_tree::SparseMerkleTree;
8288

8389
/// Expected path size: log2(256) * 2, used for hint vector capacity
8490
pub const EXPECTED_PATH_SIZE: usize = 16;

‎src/merge.rs‎

Lines changed: 72 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,12 @@ pub enum MergeValue {
1212
zero_bits: H256,
1313
zero_count: u8,
1414
},
15+
#[cfg(feature = "trie")]
16+
ShortCut {
17+
key: H256,
18+
value: H256,
19+
height: u8,
20+
},
1521
}
1622

1723
impl MergeValue {
@@ -24,10 +30,26 @@ impl MergeValue {
2430
}
2531

2632
pub fn is_zero(&self) -> bool {
27-
if let MergeValue::Value(v) = self {
28-
return v.is_zero();
33+
match self {
34+
MergeValue::Value(v) => v.is_zero(),
35+
MergeValue::MergeWithZero { .. } => false,
36+
#[cfg(feature = "trie")]
37+
MergeValue::ShortCut { .. } => false,
38+
}
39+
}
40+
41+
#[cfg(feature = "trie")]
42+
pub fn shortcut_or_value(key: H256, value: H256, height: u8) -> Self {
43+
if height == 0 || value.is_zero() {
44+
MergeValue::Value(value)
45+
} else {
46+
MergeValue::ShortCut { key, value, height }
2947
}
30-
false
48+
}
49+
50+
#[cfg(feature = "trie")]
51+
pub fn is_shortcut(&self) -> bool {
52+
matches!(self, MergeValue::ShortCut { .. })
3153
}
3254

3355
pub fn hash<H: Hasher + Default>(&self) -> H256 {
@@ -45,6 +67,34 @@ impl MergeValue {
4567
hasher.write_byte(*zero_count);
4668
hasher.finish()
4769
}
70+
#[cfg(feature = "trie")]
71+
MergeValue::ShortCut { key, value, height } => {
72+
into_merge_value::<H>(*key, *value, *height).hash::<H>()
73+
}
74+
}
75+
}
76+
}
77+
78+
/// Helper function for Shortcut node
79+
/// Transform it into a MergeValue or MergeWithZero node
80+
#[cfg(feature = "trie")]
81+
pub fn into_merge_value<H: Hasher + Default>(key: H256, value: H256, height: u8) -> MergeValue {
82+
// try keep hash same with MergeWithZero
83+
if value.is_zero() {
84+
MergeValue::from_h256(H256::zero())
85+
} else {
86+
let base_key = key.parent_path(0);
87+
let base_node = hash_base_node::<H>(0, &base_key, &value);
88+
let mut zero_bits = key;
89+
for i in height..=core::u8::MAX {
90+
if key.get_bit(i) {
91+
zero_bits.clear_bit(i);
92+
}
93+
}
94+
MergeValue::MergeWithZero {
95+
base_node,
96+
zero_bits,
97+
zero_count: height,
4898
}
4999
}
50100
}
@@ -89,7 +139,7 @@ pub fn merge<H: Hasher + Default>(
89139
MergeValue::Value(hasher.finish())
90140
}
91141

92-
fn merge_with_zero<H: Hasher + Default>(
142+
pub fn merge_with_zero<H: Hasher + Default>(
93143
height: u8,
94144
node_key: &H256,
95145
value: &MergeValue,
@@ -123,5 +173,23 @@ fn merge_with_zero<H: Hasher + Default>(
123173
zero_count: zero_count.wrapping_add(1),
124174
}
125175
}
176+
#[cfg(feature = "trie")]
177+
MergeValue::ShortCut { key, value, .. } => {
178+
if height == core::u8::MAX {
179+
let base_key = key.parent_path(0);
180+
let base_node = hash_base_node::<H>(0, &base_key, value);
181+
MergeValue::MergeWithZero {
182+
base_node,
183+
zero_bits: *key,
184+
zero_count: 0,
185+
}
186+
} else {
187+
MergeValue::ShortCut {
188+
key: *key,
189+
value: *value,
190+
height: height + 1,
191+
}
192+
}
193+
}
126194
}
127195
}

‎src/merkle_proof.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,8 @@ impl MerkleProof {
104104
buffer.extend_from_slice(zero_bits.as_slice());
105105
(Some(0x51), Some(buffer))
106106
}
107+
#[cfg(feature = "trie")]
108+
_ => unreachable!(),
107109
}
108110
} else {
109111
zero_count += 1;
@@ -501,6 +503,8 @@ impl CompiledMerkleProof {
501503
sub_proof.extend(zero_bits.as_slice());
502504
is_last_merge_zero = false;
503505
}
506+
#[cfg(feature = "trie")]
507+
_ => {}
504508
};
505509
}
506510
}

0 commit comments

Comments
 (0)