diff --git a/Cargo.toml b/Cargo.toml index a994d73..263d26e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,13 +11,13 @@ exclude = ["/fixtures", "/proptest-regressions"] # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [features] -default = ["std", "blake2b"] -std = [] +default = ["blake2b"] blake2b = ["blake2b-rs"] [dependencies] cfg-if = "0.1" blake2b-rs = { version = "0.1", optional = true } +numext-fixed-hash = "0.1.6" [dev-dependencies] proptest = "0.9" diff --git a/benches/smt_benchmark.rs b/benches/smt_benchmark.rs index d067942..800bda8 100644 --- a/benches/smt_benchmark.rs +++ b/benches/smt_benchmark.rs @@ -2,11 +2,14 @@ extern crate criterion; use criterion::Criterion; +use numext_fixed_hash; use rand::{thread_rng, Rng}; use sparse_merkle_tree::{ - blake2b::Blake2bHasher, default_store::DefaultStore, tree::SparseMerkleTree, H256, + blake2b::Blake2bHasher, default_store::DefaultStore, tree::SparseMerkleTree, }; +// Because the direct 'use numext_fixed_hash::H256' may cause part of the IDE errors. (may be about macro expansion) +type H256 = numext_fixed_hash::H256; const TARGET_LEAVES_COUNT: usize = 20; type SMT = SparseMerkleTree>; @@ -23,7 +26,7 @@ fn random_smt(update_count: usize, rng: &mut impl Rng) -> (SMT, Vec) { for _ in 0..update_count { let key = random_h256(rng); let value = random_h256(rng); - smt.update(key, value).unwrap(); + smt.update(key.clone(), value.clone()).unwrap(); keys.push(key); } (smt, keys) @@ -71,7 +74,7 @@ fn bench(c: &mut Criterion) { let leaves: Vec<_> = keys .iter() .take(TARGET_LEAVES_COUNT) - .map(|k| (*k, smt.get(k).unwrap())) + .map(|k| (k.clone(), smt.get(&k).unwrap())) .collect(); let proof = smt .merkle_proof(keys.into_iter().take(TARGET_LEAVES_COUNT).collect()) diff --git a/c/rust-tests/Cargo.toml b/c/rust-tests/Cargo.toml index a31add7..0dbd47c 100644 --- a/c/rust-tests/Cargo.toml +++ b/c/rust-tests/Cargo.toml @@ -17,6 +17,7 @@ hex = "0.4" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" anyhow = "1.0" +numext-fixed-hash = "0.1.6" [build-dependencies] cc = "1.0" \ No newline at end of file diff --git a/c/rust-tests/src/tests/smt.rs b/c/rust-tests/src/tests/smt.rs index e0419a5..a9e0f82 100644 --- a/c/rust-tests/src/tests/smt.rs +++ b/c/rust-tests/src/tests/smt.rs @@ -6,9 +6,10 @@ use proptest::prelude::*; use rand::prelude::Rng; use serde::{Deserialize, Serialize}; use sparse_merkle_tree::traits::Hasher; -use sparse_merkle_tree::{default_store::DefaultStore, SparseMerkleTree, H256}; +use sparse_merkle_tree::{default_store::DefaultStore, SparseMerkleTree}; use std::collections::HashMap; -use std::fs; +//use std::fs; +use numext_fixed_hash::H256; #[link(name = "dl-c-impl", kind = "static")] extern "C" { @@ -145,7 +146,7 @@ impl Hasher for CkbBlake2bHasher { self.0.update(&[b][..]); } fn write_h256(&mut self, h: &H256) { - self.0.update(h.as_slice()); + self.0.update(h.as_bytes()); } fn finish(self) -> H256 { let mut hash = [0u8; 32]; @@ -273,7 +274,7 @@ fn run_test_case(case: Case) -> AnyResult<()> { assert_eq!(smt_state.len(), leaves.len() as u32); smt_state - .verify(ckb_smt.root().as_slice(), &ckb_actual_compiled_proof_bin) + .verify(ckb_smt.root().as_bytes(), &ckb_actual_compiled_proof_bin) .unwrap(); } Ok(()) @@ -300,27 +301,27 @@ proptest! { const EXPECTED_PROOF_SIZE: usize = 16; let mut tree = CkbSMT::default(); - tree.update(key, value).expect("update"); + tree.update(key.clone(), value.clone()).expect("update"); if !tree.is_empty() { - let proof = tree.merkle_proof(vec![key]).expect("proof"); + let proof = tree.merkle_proof(vec![key.clone()]).expect("proof"); let compiled_proof = proof .clone() - .compile(vec![(key, value)]) + .compile(vec![(key.clone(), value.clone())]) .expect("compile proof"); assert!(proof.merkle_path().len() < EXPECTED_PROOF_SIZE); assert!(proof - .verify::(tree.root(), vec![(key, value)]) + .verify::(tree.root(), vec![(key.clone(), value.clone())]) .expect("verify")); assert!(compiled_proof - .verify::(tree.root(), vec![(key, value)]) + .verify::(tree.root(), vec![(key.clone(), value.clone())]) .expect("compiled verify")); let compiled_proof_bin: Vec = compiled_proof.into(); let mut smt_state = SmtCImpl::new(8); - smt_state.insert(key.as_slice(), value.as_slice()).unwrap(); + smt_state.insert(key.as_bytes(), value.as_bytes()).unwrap(); smt_state.normalize(); smt_state - .verify(tree.root().as_slice(), &compiled_proof_bin) + .verify(tree.root().as_bytes(), &compiled_proof_bin) .expect("verify with c"); } } diff --git a/src/blake2b.rs b/src/blake2b.rs index 47f50e7..a264619 100644 --- a/src/blake2b.rs +++ b/src/blake2b.rs @@ -19,7 +19,7 @@ impl Default for Blake2bHasher { impl Hasher for Blake2bHasher { fn write_h256(&mut self, h: &H256) { - self.0.update(h.as_slice()); + self.0.update(h.as_bytes()); } fn write_byte(&mut self, b: u8) { self.0.update(&[b][..]); diff --git a/src/default_store.rs b/src/default_store.rs index fda98b1..70bc8dd 100644 --- a/src/default_store.rs +++ b/src/default_store.rs @@ -50,12 +50,5 @@ impl Store for DefaultStore { } } -cfg_if::cfg_if! { - if #[cfg(feature = "std")] { - pub type Map = collections::HashMap; - pub type Entry<'a, K, V> = collections::hash_map::Entry<'a, K, V>; - } else { - pub type Map = collections::BTreeMap; - pub type Entry<'a, K, V> = collections::btree_map::Entry<'a, K, V>; - } -} +pub type Map = collections::BTreeMap; +pub type Entry<'a, K, V> = collections::btree_map::Entry<'a, K, V>; diff --git a/src/error.rs b/src/error.rs index 3f4f79f..f18af7e 100644 --- a/src/error.rs +++ b/src/error.rs @@ -66,5 +66,4 @@ impl core::fmt::Display for Error { } } -#[cfg(feature = "std")] impl std::error::Error for Error {} diff --git a/src/h256.rs b/src/h256.rs index b76143b..f620747 100644 --- a/src/h256.rs +++ b/src/h256.rs @@ -1,112 +1,48 @@ -use core::cmp::Ordering; +use numext_fixed_hash; -/// Represent 256 bits -#[derive(Eq, PartialEq, Debug, Default, Hash, Clone, Copy)] -pub struct H256([u8; 32]); +pub(crate) type H256 = numext_fixed_hash::H256; -const ZERO: H256 = H256([0u8; 32]); const BYTE_SIZE: u8 = 8; -impl H256 { - pub const fn zero() -> Self { - ZERO - } - - pub fn is_zero(&self) -> bool { - self == &ZERO - } - - #[inline] - pub fn get_bit(&self, i: u8) -> bool { - let byte_pos = i / BYTE_SIZE; - let bit_pos = i % BYTE_SIZE; - let bit = self.0[byte_pos as usize] >> bit_pos & 1; - bit != 0 - } - - #[inline] - pub fn set_bit(&mut self, i: u8) { - let byte_pos = i / BYTE_SIZE; - let bit_pos = i % BYTE_SIZE; - self.0[byte_pos as usize] |= 1 << bit_pos as u8; - } - - #[inline] - pub fn clear_bit(&mut self, i: u8) { - let byte_pos = i / BYTE_SIZE; - let bit_pos = i % BYTE_SIZE; - self.0[byte_pos as usize] &= !((1 << bit_pos) as u8); - } - - #[inline] - pub fn is_right(&self, height: u8) -> bool { - self.get_bit(height) - } - - pub fn as_slice(&self) -> &[u8] { - &self.0[..] - } - - /// Treat H256 as a path in a tree - /// fork height is the number of common bits(from heigher to lower: 255..=0) of two H256 - pub fn fork_height(&self, key: &H256) -> u8 { - for h in (0..=core::u8::MAX).rev() { - if self.get_bit(h) != key.get_bit(h) { - return h; - } +pub fn fork_height(data: &H256, key: &H256) -> u8 { + for h in (0..=core::u8::MAX).rev() { + if data.bit(h.into()).unwrap_or(false) != key.bit(h.into()).unwrap_or(false) { + return h; } - 0 } + 0 +} - /// Treat H256 as a path in a tree - /// return parent_path of self - pub fn parent_path(&self, height: u8) -> Self { - if height == core::u8::MAX { - H256::zero() - } else { - self.copy_bits(height + 1) - } +pub fn parent_path(data: &H256, height: u8) -> H256 { + if height == core::u8::MAX { + H256::empty() + } else { + copy_bits(data, height + 1) } +} - /// Copy bits and return a new H256 - pub fn copy_bits(&self, start: u8) -> Self { - let mut target = H256::zero(); - - let start_byte = (start / BYTE_SIZE) as usize; - // copy bytes - target.0[start_byte..].copy_from_slice(&self.0[start_byte..]); +/// Copy bits and return a new H256 +pub fn copy_bits(data: &H256, start: u8) -> H256 { + // It can also be implemented with And, but the performance is not as good as this + let mut target = H256::empty(); - // reset remain bytes - let remain = start % BYTE_SIZE; - if remain > 0 { - target.0[start_byte] &= 0b11111111 << remain - } + let start_byte = (start / BYTE_SIZE) as usize; + // copy bytes + target.0[start_byte..].copy_from_slice(&data.0[start_byte..]); - target + // reset remain bytes + let remain = start % BYTE_SIZE; + if remain > 0 { + target.0[start_byte] &= 0b11111111 << remain } -} -impl PartialOrd for H256 { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for H256 { - fn cmp(&self, other: &Self) -> Ordering { - // Compare bits from heigher to lower (255..0) - self.0.iter().rev().cmp(other.0.iter().rev()) - } + target } -impl From<[u8; 32]> for H256 { - fn from(v: [u8; 32]) -> H256 { - H256(v) - } +pub(crate) fn smt_sort_unstable(v: &mut Vec) { + (*v).sort_unstable_by(|k1, k2| k1.0.iter().rev().cmp(k2.0.iter().rev())); } -impl Into<[u8; 32]> for H256 { - fn into(self: H256) -> [u8; 32] { - self.0 - } +pub(crate) fn smt_sort_unstable_kv(v: &mut Vec<(H256, H256)>) { + v.sort_unstable_by(|(k1, _v1), (k2, _v2)| k1.0.iter().rev().cmp(k2.0.iter().rev())); } diff --git a/src/lib.rs b/src/lib.rs index b33f77e..bec4d98 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,9 +6,10 @@ //! use sparse_merkle_tree::{ //! blake2b::Blake2bHasher, default_store::DefaultStore, //! error::Error, MerkleProof, -//! SparseMerkleTree, traits::Value, H256 +//! SparseMerkleTree, traits::Value //! }; //! use blake2b_rs::{Blake2b, Blake2bBuilder}; +//! use numext_fixed_hash::H256; //! //! // define SMT //! type SMT = SparseMerkleTree>; @@ -19,7 +20,7 @@ //! impl Value for Word { //! fn to_h256(&self) -> H256 { //! if self.0.is_empty() { -//! return H256::zero(); +//! return H256::empty(); //! } //! let mut buf = [0u8; 32]; //! let mut hasher = new_blake2b(); @@ -59,8 +60,6 @@ //! } //! ``` -#![cfg_attr(not(feature = "std"), no_std)] - #[cfg(feature = "blake2b")] pub mod blake2b; pub mod default_store; @@ -73,7 +72,7 @@ mod tests; pub mod traits; pub mod tree; -pub use h256::H256; +pub(crate) use h256::H256; pub use merkle_proof::{CompiledMerkleProof, MerkleProof}; pub use tree::SparseMerkleTree; @@ -82,15 +81,6 @@ pub const EXPECTED_PATH_SIZE: usize = 16; // Max stack size can be used when verify compiled proof pub(crate) const MAX_STACK_SIZE: usize = 257; -cfg_if::cfg_if! { - if #[cfg(feature = "std")] { - use std::collections; - use std::vec; - use std::string; - } else { - extern crate alloc; - use alloc::collections; - use alloc::vec; - use alloc::string; - } -} +extern crate alloc; +use alloc::collections; +use alloc::string; diff --git a/src/merge.rs b/src/merge.rs index c42cedb..4eeb866 100644 --- a/src/merge.rs +++ b/src/merge.rs @@ -20,19 +20,19 @@ impl MergeValue { } pub fn zero() -> Self { - MergeValue::Value(H256::zero()) + MergeValue::Value(H256::empty()) } pub fn is_zero(&self) -> bool { if let MergeValue::Value(v) = self { - return v.is_zero(); + return v.is_empty(); } false } pub fn hash(&self) -> H256 { match self { - MergeValue::Value(v) => *v, + MergeValue::Value(v) => v.clone(), MergeValue::MergeWithZero { base_node, zero_bits, @@ -97,9 +97,9 @@ fn merge_with_zero( ) -> MergeValue { match value { MergeValue::Value(v) => { - let mut zero_bits = H256::zero(); + let mut zero_bits = H256::empty(); if set_bit { - zero_bits.set_bit(height); + zero_bits.set_bit(height.into(), true); } let base_node = hash_base_node::(height, node_key, v); MergeValue::MergeWithZero { @@ -113,12 +113,12 @@ fn merge_with_zero( zero_bits, zero_count, } => { - let mut zero_bits = *zero_bits; + let mut zero_bits = zero_bits.clone(); if set_bit { - zero_bits.set_bit(height); + zero_bits.set_bit(height.into(), true); } MergeValue::MergeWithZero { - base_node: *base_node, + base_node: base_node.clone(), zero_bits, zero_count: zero_count.wrapping_add(1), } diff --git a/src/merkle_proof.rs b/src/merkle_proof.rs index ba102f7..710301d 100644 --- a/src/merkle_proof.rs +++ b/src/merkle_proof.rs @@ -1,8 +1,8 @@ use crate::{ error::{Error, Result}, + h256, merge::{merge, MergeValue}, traits::Hasher, - vec::Vec, H256, MAX_STACK_SIZE, }; @@ -49,7 +49,8 @@ impl MerkleProof { &self.merkle_path } - pub fn compile(self, mut leaves: Vec<(H256, H256)>) -> Result { + pub fn compile(self, leaves: Vec<(H256, H256)>) -> Result { + let mut leaves = leaves; if leaves.is_empty() { return Err(Error::EmptyKeys); } else if leaves.len() != self.leaves_count() { @@ -59,7 +60,7 @@ impl MerkleProof { }); } // sort leaves - leaves.sort_unstable_by_key(|(k, _v)| *k); + h256::smt_sort_unstable_kv(&mut leaves); let (leaves_bitmap, merkle_path) = self.take(); @@ -69,9 +70,9 @@ impl MerkleProof { let mut leaf_index = 0; let mut merkle_path_index = 0; while leaf_index < leaves.len() { - let (leaf_key, _value) = leaves[leaf_index]; + let (leaf_key, _value) = leaves[leaf_index].clone(); let fork_height = if leaf_index + 1 < leaves.len() { - leaf_key.fork_height(&leaves[leaf_index + 1].0) + h256::fork_height(&leaf_key, &leaves[leaf_index + 1].0) } else { core::u8::MAX }; @@ -86,14 +87,17 @@ impl MerkleProof { if stack_top > 0 && stack_fork_height[stack_top - 1] == height { stack_top -= 1; (Some(0x48), None) - } else if leaves_bitmap[leaf_index].get_bit(height) { + } else if leaves_bitmap[leaf_index] + .bit(height.into()) + .unwrap_or(false) + { if merkle_path_index >= merkle_path.len() { return Err(Error::CorruptedProof); } let node = &merkle_path[merkle_path_index]; merkle_path_index += 1; match node { - MergeValue::Value(v) => (Some(0x50), Some(v.as_slice().to_vec())), + MergeValue::Value(v) => (Some(0x50), Some(v.as_bytes().to_vec())), MergeValue::MergeWithZero { base_node, zero_bits, @@ -101,8 +105,8 @@ impl MerkleProof { } => { let mut buffer = Vec::new(); buffer.push(*zero_count); - buffer.extend_from_slice(base_node.as_slice()); - buffer.extend_from_slice(zero_bits.as_slice()); + buffer.extend_from_slice(base_node.as_bytes()); + buffer.extend_from_slice(zero_bits.as_bytes()); (Some(0x51), Some(buffer)) } } @@ -183,8 +187,9 @@ impl MerkleProof { pub struct CompiledMerkleProof(pub Vec); impl CompiledMerkleProof { - pub fn compute_root(&self, mut leaves: Vec<(H256, H256)>) -> Result { - leaves.sort_unstable_by_key(|(k, _v)| *k); + pub fn compute_root(&self, leaves: Vec<(H256, H256)>) -> Result { + let mut leaves = leaves; + h256::smt_sort_unstable_kv(&mut leaves); let mut program_index = 0; let mut leaf_index = 0; let mut stack: Vec<(u16, H256, MergeValue)> = Vec::new(); @@ -197,8 +202,8 @@ impl CompiledMerkleProof { if leaf_index >= leaves.len() { return Err(Error::CorruptedStack); } - let (k, v) = leaves[leaf_index]; - stack.push((0, k, MergeValue::from_h256(v))); + let (k, v) = leaves[leaf_index].clone(); + stack.push((0, k.clone(), MergeValue::from_h256(v))); leaf_index += 1; } // P : hash stack top item with sibling node in proof @@ -218,8 +223,8 @@ impl CompiledMerkleProof { return Err(Error::CorruptedProof); } let height = height_u16 as u8; - let parent_key = key.parent_path(height); - let parent = if key.get_bit(height) { + let parent_key = h256::parent_path(&key, height); + let parent = if key.bit(height.into()).unwrap_or(false) { merge::(height, &parent_key, &sibling_node, &value) } else { merge::(height, &parent_key, &value, &sibling_node) @@ -258,8 +263,8 @@ impl CompiledMerkleProof { return Err(Error::CorruptedProof); } let height = height_u16 as u8; - let parent_key = key.parent_path(height); - let parent = if key.get_bit(height) { + let parent_key = h256::parent_path(&key, height); + let parent = if key.bit(height.into()).unwrap_or(false) { merge::(height, &parent_key, &sibling_node, &value) } else { merge::(height, &parent_key, &value, &sibling_node) @@ -281,12 +286,12 @@ impl CompiledMerkleProof { } let height_u16 = height_a; let height = height_u16 as u8; - let parent_key_a = key_a.parent_path(height); - let parent_key_b = key_b.parent_path(height); + let parent_key_a = h256::parent_path(&key_a, height); + let parent_key_b = h256::parent_path(&key_b, height); if parent_key_a != parent_key_b { return Err(Error::CorruptedProof); } - let parent = if key_a.get_bit(height) { + let parent = if key_a.bit(height.into()).unwrap_or(false) { merge::(height, &parent_key_a, &value_b, &value_a) } else { merge::(height, &parent_key_a, &value_a, &value_b) @@ -308,7 +313,7 @@ impl CompiledMerkleProof { if base_height > 255 { return Err(Error::CorruptedProof); } - let mut parent_key = key; + let mut parent_key = key.clone(); let mut height_u16 = base_height; for idx in 0..zero_count { if base_height + idx > 255 { @@ -316,8 +321,8 @@ impl CompiledMerkleProof { } height_u16 = base_height + idx; let height = height_u16 as u8; - parent_key = key.parent_path(height); - value = if key.get_bit(height) { + parent_key = h256::parent_path(&key, height); + value = if key.bit(height.into()).unwrap_or(false) { merge::(height, &parent_key, &MergeValue::zero(), &value) } else { merge::(height, &parent_key, &value, &MergeValue::zero()) diff --git a/src/tests/tree.rs b/src/tests/tree.rs index ce72fc5..bed485c 100644 --- a/src/tests/tree.rs +++ b/src/tests/tree.rs @@ -13,33 +13,37 @@ fn test_default_root() { let mut tree = SMT::default(); assert_eq!(tree.store().branches_map().len(), 0); assert_eq!(tree.store().leaves_map().len(), 0); - assert_eq!(tree.root(), &H256::zero()); + assert_eq!(tree.root(), &H256::empty()); // insert a key-value - tree.update(H256::zero(), [42u8; 32].into()) + tree.update(H256::empty(), [42u8; 32].into()) .expect("update"); - assert_ne!(tree.root(), &H256::zero()); + assert_ne!(tree.root(), &H256::empty()); assert_ne!(tree.store().branches_map().len(), 0); assert_ne!(tree.store().leaves_map().len(), 0); - assert_eq!(tree.get(&H256::zero()).expect("get"), [42u8; 32].into()); + assert_eq!(tree.get(&H256::empty()).expect("get"), [42u8; 32].into()); // update zero is to delete the key - tree.update(H256::zero(), H256::zero()).expect("update"); - assert_eq!(tree.root(), &H256::zero()); - assert_eq!(tree.get(&H256::zero()).expect("get"), H256::zero()); + tree.update(H256::empty(), H256::empty()).expect("update"); + assert_eq!(tree.root(), &H256::empty()); + assert_eq!(tree.get(&H256::empty()).expect("get"), H256::empty()); } #[test] fn test_default_tree() { let tree = SMT::default(); - assert_eq!(tree.get(&H256::zero()).expect("get"), H256::zero()); - let proof = tree.merkle_proof(vec![H256::zero()]).expect("merkle proof"); + assert_eq!(tree.get(&H256::empty()).expect("get"), H256::empty()); + let proof = tree + .merkle_proof(vec![H256::empty()]) + .expect("merkle proof"); let root = proof - .compute_root::(vec![(H256::zero(), H256::zero())]) + .compute_root::(vec![(H256::empty(), H256::empty())]) .expect("root"); assert_eq!(&root, tree.root()); - let proof = tree.merkle_proof(vec![H256::zero()]).expect("merkle proof"); + let proof = tree + .merkle_proof(vec![H256::empty()]) + .expect("merkle proof"); let root2 = proof - .compute_root::(vec![(H256::zero(), [42u8; 32].into())]) + .compute_root::(vec![(H256::empty(), [42u8; 32].into())]) .expect("root"); assert_ne!(&root2, tree.root()); } @@ -61,7 +65,7 @@ fn test_default_merkle_proof() { // let root = proof // .compute_root::(vec![([42u8; 32].into(), [42u8; 32].into())]) // .expect("compute root"); - // assert_ne!(root, H256::zero()); + // assert_ne!(root, H256::empty()); } #[test] @@ -109,9 +113,9 @@ fn test_zero_value_donot_change_root() { 0, 1, ] .into(); - let value = H256::zero(); + let value = H256::empty(); tree.update(key, value).unwrap(); - assert_eq!(tree.root(), &H256::zero()); + assert_eq!(tree.root(), &H256::empty()); assert_eq!(tree.store().leaves_map().len(), 0); assert_eq!(tree.store().branches_map().len(), 0); } @@ -130,8 +134,8 @@ fn test_zero_value_donot_change_store() { ] .into(); tree.update(key, value).unwrap(); - assert_ne!(tree.root(), &H256::zero()); - let root = *tree.root(); + assert_ne!(tree.root(), &H256::empty()); + let root = tree.root().clone(); let store = tree.store().clone(); // insert a zero value leaf @@ -165,12 +169,12 @@ fn test_delete_a_leaf() { ] .into(); tree.update(key, value).unwrap(); - assert_ne!(tree.root(), &H256::zero()); - let root = *tree.root(); + assert_ne!(tree.root(), &H256::empty()); + let root = tree.root().clone(); let store = tree.store().clone(); // insert a leaf - let key = [ + let key: H256 = [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, ] @@ -180,11 +184,11 @@ fn test_delete_a_leaf() { 0, 1, ] .into(); - tree.update(key, value).unwrap(); + tree.update(key.clone(), value).unwrap(); assert_ne!(tree.root(), &root); // delete a leaf - tree.update(key, H256::zero()).unwrap(); + tree.update(key, H256::empty()).unwrap(); assert_eq!(tree.root(), &root); assert_eq!(tree.store().leaves_map(), store.leaves_map()); assert_eq!(tree.store().branches_map(), store.branches_map()); @@ -206,7 +210,7 @@ fn test_sibling_key_get() { 0, 0, 0, ]); // get non exists sibling key should return zero value; - assert_eq!(H256::zero(), tree.get(&sibling_key).unwrap()); + assert_eq!(H256::empty(), tree.get(&sibling_key).unwrap()); } { @@ -216,14 +220,15 @@ fn test_sibling_key_get() { 0, 0, 0, ]); let value = H256::from([1u8; 32]); - tree.update(key, value).expect("update"); + tree.update(key.clone(), value.clone()).expect("update"); let sibling_key = H256::from([ 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ]); let sibling_value = H256::from([2u8; 32]); - tree.update(sibling_key, sibling_value).expect("update"); + tree.update(sibling_key.clone(), sibling_value.clone()) + .expect("update"); // get sibling key should return corresponding value assert_eq!(value, tree.get(&key).unwrap()); assert_eq!(sibling_value, tree.get(&sibling_key).unwrap()); @@ -234,13 +239,13 @@ fn test_construct(key: H256, value: H256) { // insert same value to sibling key will construct a different root let mut tree = SMT::default(); - tree.update(key, value).expect("update"); + tree.update(key.clone(), value.clone()).expect("update"); let mut sibling_key = key; - if sibling_key.get_bit(0) { - sibling_key.clear_bit(0); + if sibling_key.bit(0).unwrap_or(false) { + sibling_key.set_bit(0, false); } else { - sibling_key.set_bit(0); + sibling_key.set_bit(0, true); } let mut tree2 = SMT::default(); tree2.update(sibling_key, value).expect("update"); @@ -249,7 +254,7 @@ fn test_construct(key: H256, value: H256) { fn test_update(key: H256, value: H256) { let mut tree = SMT::default(); - tree.update(key, value).expect("update"); + tree.update(key.clone(), value.clone()).expect("update"); assert_eq!(tree.get(&key), Ok(value)); } @@ -257,10 +262,10 @@ fn test_update_tree_store(key: H256, value: H256, value2: H256) { const EXPECTED_LEAVES_LEN: usize = 1; let mut tree = SMT::default(); - tree.update(key, value).expect("update"); + tree.update(key.clone(), value).expect("update"); assert_eq!(tree.store().branches_map().len(), 256); assert_eq!(tree.store().leaves_map().len(), EXPECTED_LEAVES_LEN); - tree.update(key, value2).expect("update"); + tree.update(key.clone(), value2.clone()).expect("update"); assert_eq!(tree.store().branches_map().len(), 256); assert_eq!(tree.store().leaves_map().len(), EXPECTED_LEAVES_LEN); assert_eq!(tree.get(&key), Ok(value2)); @@ -270,16 +275,16 @@ fn test_merkle_proof(key: H256, value: H256) { const EXPECTED_MERKLE_PATH_SIZE: usize = 1; let mut tree = SMT::default(); - tree.update(key, value).expect("update"); + tree.update(key.clone(), value.clone()).expect("update"); if !tree.is_empty() { - let proof = tree.merkle_proof(vec![key]).expect("proof"); + let proof = tree.merkle_proof(vec![key.clone()]).expect("proof"); let compiled_proof = proof .clone() - .compile(vec![(key, value)]) + .compile(vec![(key.clone(), value.clone())]) .expect("compile proof"); assert!(proof.merkle_path().len() < EXPECTED_MERKLE_PATH_SIZE); assert!(proof - .verify::(tree.root(), vec![(key, value)]) + .verify::(tree.root(), vec![(key.clone(), value.clone())]) .expect("verify")); assert!(compiled_proof .verify::(tree.root(), vec![(key, value)]) @@ -338,15 +343,16 @@ fn merkle_proof(max_proof: usize) -> impl Strategy> { proptest! { #[test] fn test_h256(key: [u8; 32], key2: [u8; 32]) { - let mut list1: Vec = vec![key.into() , key2.into()]; + let mut list1: Vec = vec![H256::from(key) , H256::from(key2)]; let mut list2 = list1.clone(); // sort H256 - list1.sort_unstable_by_key(|k| *k); + list1.sort_unstable_by_key(|k| k.clone()); + h256::smt_sort_unstable(&mut list1); // sort by high bits to lower bits list2.sort_unstable_by(|k1, k2| { for i in (0u8..=255).rev() { - let b1 = if k1.get_bit(i) { 1 } else { 0 }; - let b2 = if k2.get_bit(i) { 1 } else { 0 }; + let b1 = if (*k1).bit(i.into()).unwrap_or(false) { 1 } else { 0 }; + let b2 = if (*k2).bit(i.into()).unwrap_or(false) { 1 } else { 0 }; let o = b1.cmp(&b2); if o != std::cmp::Ordering::Equal { return o; @@ -360,12 +366,12 @@ proptest! { #[test] fn test_h256_copy_bits(start: u8) { let one: H256 = [255u8; 32].into(); - let target = one.copy_bits(start); + let target = h256::copy_bits(&one, start); for i in start..=core::u8::MAX { - assert_eq!(one.get_bit(i), target.get_bit(i)); + assert_eq!(one.bit(i.into()).unwrap_or(false), target.bit(i.into()).unwrap_or(false)); } for i in 0..start { - assert!(!target.get_bit(i)); + assert!(!target.bit(i.into()).unwrap_or(false)); } } @@ -393,9 +399,9 @@ proptest! { fn test_smt_single_leaf_small((pairs, _n) in leaves(1, 50)){ let smt = new_smt(pairs.clone()); for (k, v) in pairs { - let proof = smt.merkle_proof(vec![k]).expect("gen proof"); - let compiled_proof = proof.clone().compile(vec![(k, v)]).expect("compile proof"); - assert!(proof.verify::(smt.root(), vec![(k, v)]).expect("verify proof")); + let proof = smt.merkle_proof(vec![k.clone()]).expect("gen proof"); + let compiled_proof = proof.clone().compile(vec![(k.clone(), v.clone())]).expect("compile proof"); + assert!(proof.verify::(smt.root(), vec![(k.clone(), v.clone())]).expect("verify proof")); assert!(compiled_proof.verify::(smt.root(), vec![(k, v)]).expect("verify compiled proof")); } } @@ -404,18 +410,18 @@ proptest! { fn test_smt_single_leaf_large((pairs, _n) in leaves(50, 100)){ let smt = new_smt(pairs.clone()); for (k, v) in pairs { - let proof = smt.merkle_proof(vec![k]).expect("gen proof"); - let compiled_proof = proof.clone().compile(vec![(k, v)]).expect("compile proof"); - assert!(proof.verify::(smt.root(), vec![(k, v)]).expect("verify proof")); + let proof = smt.merkle_proof(vec![k.clone()]).expect("gen proof"); + let compiled_proof = proof.clone().compile(vec![(k.clone(), v.clone())]).expect("compile proof"); + assert!(proof.verify::(smt.root(), vec![(k.clone(), v.clone())]).expect("verify proof")); assert!(compiled_proof.verify::(smt.root(), vec![(k, v)]).expect("verify compiled proof")); } } #[test] - fn test_smt_multi_leaves_small((pairs, n) in leaves(1, 50)){ + fn test_smt_multi_leaves_small((pairs, _n) in leaves(1, 50)){ let smt = new_smt(pairs.clone()); - let proof = smt.merkle_proof(pairs.iter().take(n).map(|(k, _v)| *k).collect()).expect("gen proof"); - let data: Vec<(H256, H256)> = pairs.into_iter().take(n).collect(); + let proof = smt.merkle_proof(pairs.iter().map(|(k, _v)| k.clone()).collect()).expect("gen proof"); + let data: Vec<(H256, H256)> = pairs.into_iter().collect(); let compiled_proof = proof.clone().compile(data.clone()).expect("compile proof"); assert!(proof.verify::(smt.root(), data.clone()).expect("verify proof")); assert!(compiled_proof.verify::(smt.root(), data).expect("verify compiled proof")); @@ -425,7 +431,7 @@ proptest! { fn test_smt_multi_leaves_large((pairs, _n) in leaves(50, 100)){ let n = 20; let smt = new_smt(pairs.clone()); - let proof = smt.merkle_proof(pairs.iter().take(n).map(|(k, _v)| *k).collect()).expect("gen proof"); + let proof = smt.merkle_proof(pairs.iter().take(n).map(|(k, _v)| k.clone()).collect()).expect("gen proof"); let data: Vec<(H256, H256)> = pairs.into_iter().take(n).collect(); let compiled_proof = proof.clone().compile(data.clone()).expect("compile proof"); assert!(proof.verify::(smt.root(), data.clone()).expect("verify proof")); @@ -437,7 +443,7 @@ proptest! { let smt = new_smt(pairs); let non_exists_keys: Vec<_> = pairs2.into_iter().map(|(k, _v)|k).collect(); let proof = smt.merkle_proof(non_exists_keys.clone()).expect("gen proof"); - let data: Vec<(H256, H256)> = non_exists_keys.into_iter().map(|k|(k, H256::zero())).collect(); + let data: Vec<(H256, H256)> = non_exists_keys.into_iter().map(|k|(k, H256::empty())).collect(); let compiled_proof = proof.clone().compile(data.clone()).expect("compile proof"); assert!(proof.verify::(smt.root(), data.clone()).expect("verify proof")); assert!(compiled_proof.verify::(smt.root(), data).expect("verify compiled proof")); @@ -453,7 +459,7 @@ proptest! { let mut keys: Vec<_> = exists_keys.into_iter().take(exists_keys_len).chain(non_exists_keys.into_iter().take(non_exists_keys_len)).collect(); keys.dedup(); let proof = smt.merkle_proof(keys.clone()).expect("gen proof"); - let data: Vec<(H256, H256)> = keys.into_iter().map(|k|(k, smt.get(&k).expect("get"))).collect(); + let data: Vec<(H256, H256)> = keys.into_iter().map(|k|(k.clone(), smt.get(&k).expect("get"))).collect(); let compiled_proof = proof.clone().compile(data.clone()).expect("compile proof"); assert!(proof.verify::(smt.root(), data.clone()).expect("verify proof")); assert!(compiled_proof.verify::(smt.root(), data).expect("verify compiled proof")); @@ -470,7 +476,7 @@ proptest! { #[test] fn test_smt_random_insert_order((pairs, _n) in leaves(5, 50)){ let smt = new_smt(pairs.clone()); - let root = *smt.root(); + let root = smt.root().clone(); let mut pairs = pairs; let mut rng = rand::thread_rng(); @@ -484,13 +490,13 @@ proptest! { // check leaves for (k, v) in &pairs { - assert_eq!(&smt2.get(k).unwrap(), v, "key value must be consisted"); + assert_eq!(&smt2.get(&k).unwrap(), v, "key value must be consisted"); - let origin_proof = smt.merkle_proof(vec![*k]).unwrap(); - let proof = smt2.merkle_proof(vec![*k]).unwrap(); + let origin_proof = smt.merkle_proof(vec![k.clone()]).unwrap(); + let proof = smt2.merkle_proof(vec![k.clone()]).unwrap(); assert_eq!(origin_proof, proof, "merkle proof must be consisted"); - let calculated_root = proof.compute_root::(vec![(*k, *v)]).unwrap(); + let calculated_root = proof.compute_root::(vec![(k.clone(), v.clone())]).unwrap(); assert_eq!(root, calculated_root, "root must be consisted"); } } @@ -501,18 +507,18 @@ proptest! { let mut rng = rand::thread_rng(); let len = rng.gen_range(0, pairs.len()); let mut smt = new_smt(pairs[..len].to_vec()); - let root = *smt.root(); + let root = smt.root().clone(); // insert zero values for (k, _v) in pairs[len..].iter() { - smt.update(*k, H256::zero()).unwrap(); + smt.update(k.clone(), H256::empty()).unwrap(); } // check root - let current_root = *smt.root(); + let current_root = smt.root().clone(); assert_eq!(root, current_root); // check inserted pairs for (k, v) in pairs[..len].iter() { - let value = smt.get(k).unwrap(); + let value = smt.get(&k).unwrap(); assert_eq!(v, &value); } } @@ -596,16 +602,21 @@ fn test_v0_2_broken_sample() { .into_iter() .map(parse_h256) .collect::>(); - let mut pairs = keys.into_iter().zip(values.into_iter()).collect::>(); + let mut pairs = keys + .into_iter() + .map(|k| k.clone()) + .into_iter() + .zip(values.into_iter()) + .collect::>(); let smt = new_smt(pairs.clone()); - let base_root = *smt.root(); + let base_root = smt.root().clone(); // insert in random order let mut rng = rand::thread_rng(); for _i in 0..10 { pairs.shuffle(&mut rng); let smt = new_smt(pairs.clone()); - let current_root = *smt.root(); + let current_root = smt.root().clone(); assert_eq!(base_root, current_root); } } @@ -687,14 +698,14 @@ fn test_replay_to_pass_proof() { ] .into(); let pairs = vec![ - (key1, existing), - (key2, non_existing), - (key3, non_existing), - (key4, non_existing), + (key1.clone(), existing), + (key2.clone(), non_existing.clone()), + (key3.clone(), non_existing.clone()), + (key4.clone(), non_existing.clone()), ]; let smt = new_smt(pairs); - let leaf_a_bl = vec![(key1, H256::zero())]; - let leaf_c = vec![(key3, non_existing)]; + let leaf_a_bl = vec![(key1, H256::empty())]; + let leaf_c = vec![(key3.clone(), non_existing)]; let leaf_other = vec![(key3, other_value)]; let proofc = smt .merkle_proof(leaf_c.clone().into_iter().map(|(k, _)| k).collect()) @@ -734,13 +745,16 @@ fn test_sibling_leaf() { H256::from(rand_data) } let rand_key = gen_rand_h256(); - let mut sibling_key = rand_key; - if rand_key.is_right(0) { - sibling_key.clear_bit(0); + let mut sibling_key = rand_key.clone(); + if rand_key.bit(0).unwrap_or(false) { + sibling_key.set_bit(0, false); } else { - sibling_key.set_bit(0); + sibling_key.set_bit(0, true); } - let pairs = vec![(rand_key, gen_rand_h256()), (sibling_key, gen_rand_h256())]; + let pairs = vec![ + (rand_key.clone(), gen_rand_h256()), + (sibling_key.clone(), gen_rand_h256()), + ]; let keys = vec![rand_key, sibling_key]; let smt = new_smt(pairs.clone()); let proof = smt.merkle_proof(keys).expect("gen proof"); @@ -753,9 +767,9 @@ fn test_sibling_leaf() { fn test_max_stack_size() { fn gen_h256(height: u8) -> H256 { // The key path is first go right `256 - height` times then go left `height` times. - let mut key = H256::zero(); + let mut key = H256::empty(); for h in height..=255 { - key.set_bit(h); + key.set_bit(h.into(), true); } key } @@ -763,15 +777,15 @@ fn test_max_stack_size() { .map(|height| (gen_h256(height), gen_h256(1))) .collect(); // Most left key - pairs.push((H256::zero(), gen_h256(1))); + pairs.push((H256::empty(), gen_h256(1))); { // A pair of sibling keys in between - let mut left_key = H256::zero(); + let mut left_key = H256::empty(); for h in 12..56 { - left_key.set_bit(h); + left_key.set_bit(h, true); } let mut right_key = left_key.clone(); - right_key.set_bit(0); + right_key.set_bit(0, true); pairs.push((left_key, gen_h256(1))); pairs.push((right_key, gen_h256(1))); } diff --git a/src/traits.rs b/src/traits.rs index a7b617e..73cd6c4 100644 --- a/src/traits.rs +++ b/src/traits.rs @@ -19,10 +19,10 @@ pub trait Value { impl Value for H256 { fn to_h256(&self) -> H256 { - *self + self.clone() } fn zero() -> Self { - H256::zero() + H256::empty() } } diff --git a/src/tree.rs b/src/tree.rs index 7304d25..6709433 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -1,9 +1,9 @@ use crate::{ error::{Error, Result}, + h256, merge::{merge, MergeValue}, merkle_proof::MerkleProof, traits::{Hasher, Store, Value}, - vec::Vec, H256, MAX_STACK_SIZE, }; use core::cmp::Ordering; @@ -68,7 +68,7 @@ impl> SparseMerkleTree { /// Check empty of the tree pub fn is_empty(&self) -> bool { - self.root.is_zero() + self.root.is_empty() } /// Destroy current tree and retake store @@ -93,7 +93,7 @@ impl> SparseMerkleTree { let node = MergeValue::from_h256(value.to_h256()); // notice when value is zero the leaf is deleted, so we do not need to store it if !node.is_zero() { - self.store.insert_leaf(key, value)?; + self.store.insert_leaf(key.clone(), value)?; } else { self.store.remove_leaf(&key)?; } @@ -102,16 +102,16 @@ impl> SparseMerkleTree { let mut current_key = key; let mut current_node = node; for height in 0..=core::u8::MAX { - let parent_key = current_key.parent_path(height); - let parent_branch_key = BranchKey::new(height, parent_key); + let parent_key = h256::parent_path(¤t_key, height); + let parent_branch_key = BranchKey::new(height, parent_key.clone()); let (left, right) = if let Some(parent_branch) = self.store.get_branch(&parent_branch_key)? { - if current_key.is_right(height) { + if current_key.bit(height.into()).unwrap_or(false) { (parent_branch.left, current_node) } else { (current_node, parent_branch.right) } - } else if current_key.is_right(height) { + } else if current_key.bit(height.into()).unwrap_or(false) { (MergeValue::zero(), current_node) } else { (current_node, MergeValue::zero()) @@ -131,7 +131,7 @@ impl> SparseMerkleTree { self.store.remove_branch(&parent_branch_key)?; } // prepare for next round - current_key = parent_key; + current_key = parent_key.clone(); current_node = merge::(height, &parent_key, &left, &right); } @@ -149,29 +149,30 @@ impl> SparseMerkleTree { } /// Generate merkle proof - pub fn merkle_proof(&self, mut keys: Vec) -> Result { + pub fn merkle_proof(&self, keys: Vec) -> Result { + let mut keys = keys; if keys.is_empty() { return Err(Error::EmptyKeys); } // sort keys - keys.sort_unstable(); + h256::smt_sort_unstable(&mut keys); // Collect leaf bitmaps let mut leaves_bitmap: Vec = Default::default(); for current_key in &keys { - let mut bitmap = H256::zero(); + let mut bitmap = H256::empty(); for height in 0..=core::u8::MAX { - let parent_key = current_key.parent_path(height); + let parent_key = h256::parent_path(¤t_key, height); let parent_branch_key = BranchKey::new(height, parent_key); if let Some(parent_branch) = self.store.get_branch(&parent_branch_key)? { - let sibling = if current_key.is_right(height) { + let sibling = if current_key.bit(height.into()).unwrap_or(false) { parent_branch.left } else { parent_branch.right }; if !sibling.is_zero() { - bitmap.set_bit(height); + bitmap.set_bit(height.into(), true); } } else { // The key is not in the tree (support non-inclusion proof) @@ -185,9 +186,9 @@ impl> SparseMerkleTree { let mut stack_top = 0; let mut leaf_index = 0; while leaf_index < keys.len() { - let leaf_key = keys[leaf_index]; + let leaf_key = keys[leaf_index].clone(); let fork_height = if leaf_index + 1 < keys.len() { - leaf_key.fork_height(&keys[leaf_index + 1]) + h256::fork_height(&leaf_key, &keys[leaf_index + 1]) } else { core::u8::MAX }; @@ -196,13 +197,16 @@ impl> SparseMerkleTree { // If it's not final round, we don't need to merge to root (height=255) break; } - let parent_key = leaf_key.parent_path(height); - let is_right = leaf_key.is_right(height); + let parent_key = h256::parent_path(&leaf_key, height); + let is_right = leaf_key.bit(height.into()).unwrap_or(false); // has non-zero sibling if stack_top > 0 && stack_fork_height[stack_top - 1] == height { stack_top -= 1; - } else if leaves_bitmap[leaf_index].get_bit(height) { + } else if leaves_bitmap[leaf_index] + .bit(height.into()) + .unwrap_or(false) + { let parent_branch_key = BranchKey::new(height, parent_key); if let Some(parent_branch) = self.store.get_branch(&parent_branch_key)? { let sibling = if is_right {