From d37f3f5e841ffb6eb45832ab307dc8238ded88ae Mon Sep 17 00:00:00 2001 From: Victor Lopez Date: Wed, 29 Mar 2023 08:31:07 +0200 Subject: [PATCH 01/10] feat: Add `depth` as store SMT argument Prior to this commit, MerkleStore allowed the creation of Sparse Merkle tree only with the maximum depth of 63. However, this doesn't fit the Tiered Sparse Merkle tree requirements, as it will contain trees of depth 16. This commit adds the `depth` argument to the MerkleStore methods that will create Sparse Merkle trees. --- README.md | 2 +- benches/store.rs | 28 ++++++++++++++++------------ src/merkle/simple_smt/mod.rs | 6 +++--- src/merkle/store/mod.rs | 30 ++++++++++++++++++++---------- src/merkle/store/tests.rs | 12 ++++++++---- 5 files changed, 48 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 365c7c7..3583271 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ For performance benchmarks of these hash functions and their comparison to other [Merkle module](./src/merkle/) provides a set of data structures related to Merkle trees. All these data structures are implemented using the RPO hash function described above. The data structures are: * `MerkleTree`: a regular fully-balanced binary Merkle tree. The depth of this tree can be at most 64. -* `SimpleSmt`: a Sparse Merkle Tree, mapping 63-bit keys to 4-element leaf values. +* `SimpleSmt`: a Sparse Merkle Tree, mapping 64-bit keys to 4-element leaf values. * `MerklePathSet`: a collection of Merkle authentication paths all resolving to the same root. The length of the paths can be at most 64. * `MerkleStore`: a collection of Merkle trees of different heights designed to efficiently store trees with common subtrees. * `Mmr`: a Merkle mountain range structure designed to function as an append-only log. diff --git a/benches/store.rs b/benches/store.rs index 793118a..b59a537 100644 --- a/benches/store.rs +++ b/benches/store.rs @@ -27,7 +27,7 @@ fn random_index(range: u64) -> u64 { fn get_empty_leaf_simplesmt(c: &mut Criterion) { let mut group = c.benchmark_group("get_empty_leaf_simplesmt"); - let depth = 63u8; + let depth = SimpleSmt::MAX_DEPTH; let size = 2u64.pow(depth as u32); // both SMT and the store are pre-populated with empty hashes, accessing these values is what is @@ -103,12 +103,12 @@ fn get_leaf_simplesmt(c: &mut Criterion) { .enumerate() .map(|(c, v)| (c.try_into().unwrap(), v.into())) .collect::>(); - let smt = SimpleSmt::new(63) + let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) .unwrap() .with_leaves(smt_leaves.clone()) .unwrap(); let store = MerkleStore::new() - .with_sparse_merkle_tree(smt_leaves) + .with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, smt_leaves) .unwrap(); let depth = smt.depth(); let root = smt.root(); @@ -136,7 +136,7 @@ fn get_leaf_simplesmt(c: &mut Criterion) { fn get_node_of_empty_simplesmt(c: &mut Criterion) { let mut group = c.benchmark_group("get_node_of_empty_simplesmt"); - let depth = 63u8; + let depth = SimpleSmt::MAX_DEPTH; let size = 2u64.pow(depth as u32); // both SMT and the store are pre-populated with the empty hashes, accessing the internal nodes @@ -216,12 +216,12 @@ fn get_node_simplesmt(c: &mut Criterion) { .enumerate() .map(|(c, v)| (c.try_into().unwrap(), v.into())) .collect::>(); - let smt = SimpleSmt::new(63) + let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) .unwrap() .with_leaves(smt_leaves.clone()) .unwrap(); let store = MerkleStore::new() - .with_sparse_merkle_tree(smt_leaves) + .with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, smt_leaves) .unwrap(); let root = smt.root(); let size_u64 = size as u64; @@ -295,12 +295,12 @@ fn get_leaf_path_simplesmt(c: &mut Criterion) { .enumerate() .map(|(c, v)| (c.try_into().unwrap(), v.into())) .collect::>(); - let smt = SimpleSmt::new(63) + let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) .unwrap() .with_leaves(smt_leaves.clone()) .unwrap(); let store = MerkleStore::new() - .with_sparse_merkle_tree(smt_leaves) + .with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, smt_leaves) .unwrap(); let depth = smt.depth(); let root = smt.root(); @@ -366,7 +366,7 @@ fn new(c: &mut Criterion) { .map(|(c, v)| (c.try_into().unwrap(), v.into())) .collect::>() }, - |l| black_box(SimpleSmt::new(63).unwrap().with_leaves(l)), + |l| black_box(SimpleSmt::new(SimpleSmt::MAX_DEPTH).unwrap().with_leaves(l)), BatchSize::SmallInput, ) }); @@ -382,7 +382,11 @@ fn new(c: &mut Criterion) { .map(|(c, v)| (c.try_into().unwrap(), v.into())) .collect::>() }, - |l| black_box(MerkleStore::new().with_sparse_merkle_tree(l)), + |l| { + black_box( + MerkleStore::new().with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, l), + ) + }, BatchSize::SmallInput, ) }, @@ -450,12 +454,12 @@ fn update_leaf_simplesmt(c: &mut Criterion) { .enumerate() .map(|(c, v)| (c.try_into().unwrap(), v.into())) .collect::>(); - let mut smt = SimpleSmt::new(63) + let mut smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) .unwrap() .with_leaves(smt_leaves.clone()) .unwrap(); let mut store = MerkleStore::new() - .with_sparse_merkle_tree(smt_leaves) + .with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, smt_leaves) .unwrap(); let depth = smt.depth(); let root = smt.root(); diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs index e330a97..841fa73 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/simple_smt/mod.rs @@ -8,7 +8,7 @@ mod tests; // SPARSE MERKLE TREE // ================================================================================================ -/// A sparse Merkle tree with 63-bit keys and 4-element leaf values, without compaction. +/// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction. /// Manipulation and retrieval of leaves and internal nodes is provided by its internal `Store`. /// The root of the tree is recomputed on each new leaf update. #[derive(Debug, Clone, PartialEq, Eq)] @@ -26,7 +26,7 @@ impl SimpleSmt { pub const MIN_DEPTH: u8 = 1; /// Maximum supported depth. - pub const MAX_DEPTH: u8 = 63; + pub const MAX_DEPTH: u8 = 64; // CONSTRUCTORS // -------------------------------------------------------------------------------------------- @@ -57,7 +57,7 @@ impl SimpleSmt { { // check if the leaves count will fit the depth setup let mut entries = entries.into_iter(); - let max = 1 << self.depth; + let max = 1 << self.depth.min(63); if entries.len() > max { return Err(MerkleError::InvalidEntriesCount(max, entries.len())); } diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index f98ae6f..2edc355 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -115,13 +115,19 @@ impl MerkleStore { Ok(self) } - /// Appends the provided sparse merkle tree represented by its `entries` to the set. - pub fn with_sparse_merkle_tree(mut self, entries: R) -> Result + /// Appends the provided Sparse Merkle tree represented by its `entries` to the set. + /// + /// For more information, check [MerkleStore::add_sparse_merkle_tree]. + pub fn with_sparse_merkle_tree( + mut self, + depth: u8, + entries: R, + ) -> Result where R: IntoIterator, I: Iterator + ExactSizeIterator, { - self.add_sparse_merkle_tree(entries)?; + self.add_sparse_merkle_tree(depth, entries)?; Ok(self) } @@ -272,20 +278,24 @@ impl MerkleStore { Ok(tree.nodes[1]) } - /// Adds all the nodes of a Sparse Merkle tree represented by `entries`. + /// Adds a Sparse Merkle tree defined by the specified `entries` to the store, and returns the + /// root of the added tree. /// - /// This will instantiate a Sparse Merkle tree using `entries` and include all the nodes into - /// the store. + /// The entries are expected to contain tuples of `(index, node)` describing nodes in the tree + /// at `depth`. /// /// # Errors - /// - /// This will return `InvalidEntriesCount` if the length of `entries` is not `63`. - pub fn add_sparse_merkle_tree(&mut self, entries: R) -> Result + /// Returns an error if the provided `depth` is greater than [SimpleSmt::MAX_DEPTH]. + pub fn add_sparse_merkle_tree( + &mut self, + depth: u8, + entries: R, + ) -> Result where R: IntoIterator, I: Iterator + ExactSizeIterator, { - let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)?.with_leaves(entries)?; + let smt = SimpleSmt::new(depth)?.with_leaves(entries)?; for branch in smt.store.branches.values() { let parent = Rpo256::merge(&[branch.left, branch.right]); self.nodes.insert( diff --git a/src/merkle/store/tests.rs b/src/merkle/store/tests.rs index 9ccd3ec..a14c557 100644 --- a/src/merkle/store/tests.rs +++ b/src/merkle/store/tests.rs @@ -205,7 +205,7 @@ fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> { let mut store = MerkleStore::default(); let keys2: [u64; 2] = [0, 1]; let leaves2: [Word; 2] = [int_to_node(1), int_to_node(2)]; - store.add_sparse_merkle_tree(keys2.into_iter().zip(leaves2.into_iter()))?; + store.add_sparse_merkle_tree(48, keys2.into_iter().zip(leaves2.into_iter()))?; let smt = SimpleSmt::new(1) .unwrap() .with_leaves(keys2.into_iter().zip(leaves2.into_iter())) @@ -231,7 +231,10 @@ fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> { #[test] fn test_sparse_merkle_tree() -> Result<(), MerkleError> { let mut store = MerkleStore::default(); - store.add_sparse_merkle_tree(KEYS4.into_iter().zip(LEAVES4.into_iter()))?; + store.add_sparse_merkle_tree( + SimpleSmt::MAX_DEPTH, + KEYS4.into_iter().zip(LEAVES4.into_iter()), + )?; let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) .unwrap() @@ -584,9 +587,10 @@ fn test_constructors() -> Result<(), MerkleError> { assert_eq!(mtree.get_path(index)?, value_path.path); } + let depth = 32; let store = MerkleStore::default() - .with_sparse_merkle_tree(KEYS4.into_iter().zip(LEAVES4.into_iter()))?; - let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH) + .with_sparse_merkle_tree(depth, KEYS4.into_iter().zip(LEAVES4.into_iter()))?; + let smt = SimpleSmt::new(depth) .unwrap() .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter())) .unwrap(); From f8a62dae76b696b7e8dc5eca859212c7d0e50b42 Mon Sep 17 00:00:00 2001 From: "Augusto F. Hack" Date: Thu, 30 Mar 2023 23:23:46 +0200 Subject: [PATCH 02/10] chore: remove simple_smt::Store --- src/merkle/simple_smt/mod.rs | 94 ++++++++++++------------------------ src/merkle/store/mod.rs | 2 +- 2 files changed, 33 insertions(+), 63 deletions(-) diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs index 841fa73..9699ea9 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/simple_smt/mod.rs @@ -9,13 +9,20 @@ mod tests; // ================================================================================================ /// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction. -/// Manipulation and retrieval of leaves and internal nodes is provided by its internal `Store`. /// The root of the tree is recomputed on each new leaf update. #[derive(Debug, Clone, PartialEq, Eq)] pub struct SimpleSmt { - root: Word, depth: u8, - pub(crate) store: Store, + root: Word, + leaves: BTreeMap, + pub(crate) branches: BTreeMap, + empty_hashes: Vec, +} + +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub(crate) struct BranchNode { + pub(crate) left: RpoDigest, + pub(crate) right: RpoDigest, } impl SimpleSmt { @@ -40,8 +47,16 @@ impl SimpleSmt { return Err(MerkleError::DepthTooBig(depth as u64)); } - let (store, root) = Store::new(depth); - Ok(Self { root, depth, store }) + let empty_hashes = EmptySubtreeRoots::empty_hashes(depth).to_vec(); + let root = empty_hashes[0].into(); + + Ok(Self { + root, + depth, + empty_hashes, + leaves: BTreeMap::new(), + branches: BTreeMap::new(), + }) } /// Appends the provided entries as leaves of the tree. @@ -72,8 +87,7 @@ impl SimpleSmt { where I: IntoIterator, { - self.store - .replace_empty_subtrees(hashes.into_iter().collect()); + self.replace_empty_subtrees(hashes.into_iter().collect()); self } @@ -95,7 +109,7 @@ impl SimpleSmt { /// Returns the set count of the keys of the leaves. pub fn leaves_count(&self) -> usize { - self.store.leaves_count() + self.leaves.len() } /// Returns a node at the specified key @@ -109,18 +123,16 @@ impl SimpleSmt { } else if index.depth() > self.depth() { Err(MerkleError::DepthTooBig(index.depth() as u64)) } else if index.depth() == self.depth() { - self.store - .get_leaf_node(index.value()) + self.get_leaf_node(index.value()) .or_else(|| { - self.store - .empty_hashes + self.empty_hashes .get(index.depth() as usize) .copied() .map(Word::from) }) .ok_or(MerkleError::InvalidIndex(*index)) } else { - let branch_node = self.store.get_branch_node(index); + let branch_node = self.get_branch_node(index); Ok(Rpo256::merge(&[branch_node.left, branch_node.right]).into()) } } @@ -142,7 +154,7 @@ impl SimpleSmt { for _ in 0..index.depth() { let is_right = index.is_value_odd(); index.move_up(); - let BranchNode { left, right } = self.store.get_branch_node(&index); + let BranchNode { left, right } = self.get_branch_node(&index); let value = if is_right { left } else { right }; path.push(*value); } @@ -167,7 +179,7 @@ impl SimpleSmt { /// # Errors /// Returns an error if the specified key is not a valid leaf index for this tree. pub fn update_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> { - if !self.store.check_leaf_node_exists(key) { + if !self.check_leaf_node_exists(key) { return Err(MerkleError::InvalidIndex(NodeIndex::new(self.depth(), key))); } self.insert_leaf(key, value)?; @@ -177,7 +189,7 @@ impl SimpleSmt { /// Inserts a leaf located at the specified key, and recomputes hashes by walking up the tree pub fn insert_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> { - self.store.insert_leaf_node(key, value); + self.insert_leaf_node(key, value); // TODO consider using a map `index |-> word` instead of `index |-> (word, word)` let mut index = NodeIndex::new(self.depth(), key); @@ -185,59 +197,21 @@ impl SimpleSmt { for _ in 0..index.depth() { let is_right = index.is_value_odd(); index.move_up(); - let BranchNode { left, right } = self.store.get_branch_node(&index); + let BranchNode { left, right } = self.get_branch_node(&index); let (left, right) = if is_right { (left, value) } else { (value, right) }; - self.store.insert_branch_node(index, left, right); + self.insert_branch_node(index, left, right); value = Rpo256::merge(&[left, right]); } self.root = value.into(); Ok(()) } -} - -// STORE -// ================================================================================================ - -/// A data store for sparse Merkle tree key-value pairs. -/// Leaves and branch nodes are stored separately in B-tree maps, indexed by key and (key, depth) -/// respectively. Hashes for blank subtrees at each layer are stored in `empty_hashes`, beginning -/// with the root hash of an empty tree, and ending with the zero value of a leaf node. -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct Store { - pub(crate) branches: BTreeMap, - leaves: BTreeMap, - pub(crate) empty_hashes: Vec, - depth: u8, -} - -#[derive(Debug, Default, Clone, PartialEq, Eq)] -pub(crate) struct BranchNode { - pub(crate) left: RpoDigest, - pub(crate) right: RpoDigest, -} - -impl Store { - fn new(depth: u8) -> (Self, Word) { - let branches = BTreeMap::new(); - let leaves = BTreeMap::new(); - // Construct empty node digests for each layer of the tree - let empty_hashes = EmptySubtreeRoots::empty_hashes(depth).to_vec(); - - let root = empty_hashes[0].into(); - let store = Self { - branches, - leaves, - empty_hashes, - depth, - }; - - (store, root) - } + // HELPER METHODS + // -------------------------------------------------------------------------------------------- fn replace_empty_subtrees(&mut self, hashes: Vec) { self.empty_hashes = hashes; @@ -269,8 +243,4 @@ impl Store { let branch = BranchNode { left, right }; self.branches.insert(index, branch); } - - fn leaves_count(&self) -> usize { - self.leaves.len() - } } diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index 2edc355..e21335f 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -296,7 +296,7 @@ impl MerkleStore { I: Iterator + ExactSizeIterator, { let smt = SimpleSmt::new(depth)?.with_leaves(entries)?; - for branch in smt.store.branches.values() { + for branch in smt.branches.values() { let parent = Rpo256::merge(&[branch.left, branch.right]); self.nodes.insert( parent, From 433b467953c15d760166fd4c9a82f2be62dd5481 Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Thu, 23 Feb 2023 17:06:19 -0800 Subject: [PATCH 03/10] feat: optimized hash_elements for blake3 hasher --- benches/README.md | 2 +- src/hash/blake/mod.rs | 26 +++++++++++++++++--------- src/hash/blake/tests.rs | 27 +++++++++++++++++++++++++++ 3 files changed, 45 insertions(+), 10 deletions(-) diff --git a/benches/README.md b/benches/README.md index 9e14f78..385e01e 100644 --- a/benches/README.md +++ b/benches/README.md @@ -28,7 +28,7 @@ The second scenario is that of sequential hashing where we take a sequence of le | Function | BLAKE3 | SHA3 | Poseidon | Rp64_256 | RPO_256 | | ------------------- | -------| ------- | --------- | --------- | ------- | -| Apple M1 Pro | 1.1 us | 1.5 us | 19.4 us | 118 us | 70 us | +| Apple M1 Pro | 1.0 us | 1.5 us | 19.4 us | 118 us | 70 us | | Apple M2 | 1.0 us | 1.5 us | 17.4 us | 103 us | 65 us | | Amazon Graviton 3 | 1.4 us | | | | 114 us | | AMD Ryzen 9 5950X | 0.8 us | 1.7 us | 15.7 us | 120 us | 72 us | diff --git a/src/hash/blake/mod.rs b/src/hash/blake/mod.rs index eb07ad6..a3bcfd0 100644 --- a/src/hash/blake/mod.rs +++ b/src/hash/blake/mod.rs @@ -1,7 +1,5 @@ use super::{Digest, ElementHasher, Felt, FieldElement, Hasher, StarkField}; -use crate::utils::{ - uninit_vector, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable, -}; +use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; use core::{ mem::{size_of, transmute, transmute_copy}, ops::Deref, @@ -290,15 +288,25 @@ where let digest = if Felt::IS_CANONICAL { blake3::hash(E::elements_as_bytes(elements)) } else { - let base_elements = E::slice_as_base_elements(elements); - let blen = base_elements.len() << 3; + let mut hasher = blake3::Hasher::new(); + + // BLAKE3 state is 64 bytes - so, we can absorb 64 bytes into the state in a single + // permutation. we move the elements into the hasher via the buffer to give the CPU + // a chance to process multiple element-to-byte conversions in parallel + let mut buf = [0_u8; 64]; + let mut chunk_iter = E::slice_as_base_elements(elements).chunks_exact(8); + for chunk in chunk_iter.by_ref() { + for i in 0..8 { + buf[i * 8..(i + 1) * 8].copy_from_slice(&chunk[i].as_int().to_le_bytes()); + } + hasher.update(&buf); + } - let mut bytes = unsafe { uninit_vector(blen) }; - for (idx, element) in base_elements.iter().enumerate() { - bytes[idx * 8..(idx + 1) * 8].copy_from_slice(&element.as_int().to_le_bytes()); + for element in chunk_iter.remainder() { + hasher.update(&element.as_int().to_le_bytes()); } - blake3::hash(&bytes) + hasher.finalize() }; *shrink_bytes(&digest.into()) } diff --git a/src/hash/blake/tests.rs b/src/hash/blake/tests.rs index 8897611..b03b06a 100644 --- a/src/hash/blake/tests.rs +++ b/src/hash/blake/tests.rs @@ -1,6 +1,22 @@ use super::*; use crate::utils::collections::Vec; use proptest::prelude::*; +use rand_utils::rand_vector; + +#[test] +fn blake3_hash_elements() { + // test multiple of 8 + let elements = rand_vector::(16); + let expected = compute_expected_element_hash(&elements); + let actual: [u8; 32] = hash_elements(&elements); + assert_eq!(&expected, &actual); + + // test not multiple of 8 + let elements = rand_vector::(17); + let expected = compute_expected_element_hash(&elements); + let actual: [u8; 32] = hash_elements(&elements); + assert_eq!(&expected, &actual); +} proptest! { #[test] @@ -18,3 +34,14 @@ proptest! { Blake3_256::hash(vec); } } + +// HELPER FUNCTIONS +// ================================================================================================ + +fn compute_expected_element_hash(elements: &[Felt]) -> blake3::Hash { + let mut bytes = Vec::new(); + for element in elements.iter() { + bytes.extend_from_slice(&element.as_int().to_le_bytes()); + } + blake3::hash(&bytes) +} From 1df4318399eeaf1e689a5959c364e0901f6174e6 Mon Sep 17 00:00:00 2001 From: "Augusto F. Hack" Date: Thu, 30 Mar 2023 23:51:44 +0200 Subject: [PATCH 04/10] feat: add node iterator to MerkleTree --- src/merkle/merkle_tree.rs | 83 +++++++++++++++++++++++++++++++++++++-- src/merkle/mod.rs | 3 ++ src/merkle/node.rs | 9 +++++ src/merkle/store/mod.rs | 33 +++++----------- 4 files changed, 101 insertions(+), 27 deletions(-) create mode 100644 src/merkle/node.rs diff --git a/src/merkle/merkle_tree.rs b/src/merkle/merkle_tree.rs index 434c8e7..6b767ae 100644 --- a/src/merkle/merkle_tree.rs +++ b/src/merkle/merkle_tree.rs @@ -1,4 +1,6 @@ -use super::{Felt, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word}; +use super::{ + Felt, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word, +}; use crate::{ utils::{string::String, uninit_vector, word_to_hex}, FieldElement, @@ -12,7 +14,7 @@ use winter_math::log2; /// A fully-balanced binary Merkle tree (i.e., a tree where the number of leaves is a power of two). #[derive(Debug, Clone, PartialEq, Eq)] pub struct MerkleTree { - pub(crate) nodes: Vec, + nodes: Vec, } impl MerkleTree { @@ -158,6 +160,47 @@ impl MerkleTree { Ok(()) } + + /// An iterator over every inner node in the tree. The iterator order is unspecified. + pub fn inner_nodes(&self) -> MerkleTreeNodes<'_> { + MerkleTreeNodes { + nodes: &self.nodes, + index: 1, // index 0 is just padding, start at 1 + } + } +} + +// ITERATORS +// ================================================================================================ + +/// An iterator over every inner node of the [MerkleTree]. +/// +/// Use this to extract the data of the tree, there is no guarantee on the order of the elements. +pub struct MerkleTreeNodes<'a> { + nodes: &'a Vec, + index: usize, +} + +impl<'a> Iterator for MerkleTreeNodes<'a> { + type Item = InnerNodeInfo; + + fn next(&mut self) -> Option { + if self.index < self.nodes.len() / 2 { + let value = self.index; + let left = self.index * 2; + let right = left + 1; + + self.index += 1; + + Some(InnerNodeInfo { + value: self.nodes[value], + left: self.nodes[left], + right: self.nodes[right], + }) + } else { + None + } + } } /// Utility to vizualize a [MerkleTree] in text. @@ -212,7 +255,7 @@ pub fn path_to_text(path: &MerklePath) -> Result { #[cfg(test)] mod tests { use super::*; - use crate::merkle::int_to_node; + use crate::merkle::{int_to_node, InnerNodeInfo}; use core::mem::size_of; use proptest::prelude::*; @@ -323,6 +366,40 @@ mod tests { assert_eq!(expected_tree.nodes, tree.nodes); } + #[test] + fn nodes() -> Result<(), MerkleError> { + let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap(); + let root = tree.root(); + let l1n0 = tree.get_node(NodeIndex::new(1, 0))?; + let l1n1 = tree.get_node(NodeIndex::new(1, 1))?; + let l2n0 = tree.get_node(NodeIndex::new(2, 0))?; + let l2n1 = tree.get_node(NodeIndex::new(2, 1))?; + let l2n2 = tree.get_node(NodeIndex::new(2, 2))?; + let l2n3 = tree.get_node(NodeIndex::new(2, 3))?; + + let nodes: Vec = tree.inner_nodes().collect(); + let expected = vec![ + InnerNodeInfo { + value: root, + left: l1n0, + right: l1n1, + }, + InnerNodeInfo { + value: l1n0, + left: l2n0, + right: l2n1, + }, + InnerNodeInfo { + value: l1n1, + left: l2n2, + right: l2n3, + }, + ]; + assert_eq!(nodes, expected); + + Ok(()) + } + proptest! { #[test] fn arbitrary_word_can_be_represented_as_digest( diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index b0c7b2d..c12ac7a 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -32,6 +32,9 @@ pub use mmr::{Mmr, MmrPeaks}; mod store; pub use store::MerkleStore; +mod node; +pub use node::InnerNodeInfo; + // ERRORS // ================================================================================================ diff --git a/src/merkle/node.rs b/src/merkle/node.rs new file mode 100644 index 0000000..681e306 --- /dev/null +++ b/src/merkle/node.rs @@ -0,0 +1,9 @@ +use super::Word; + +/// Representation of a node with two children used for iterating over containers. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InnerNodeInfo { + pub value: Word, + pub left: Word, + pub right: Word, +} diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index e21335f..557942c 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -249,33 +249,18 @@ impl MerkleStore { return Err(MerkleError::DepthTooSmall(leaves.len() as u8)); } - let layers = leaves.len().ilog2(); let tree = MerkleTree::new(leaves)?; - - let mut depth = 0; - let mut parent_offset = 1; - let mut child_offset = 2; - while depth < layers { - let layer_size = 1usize << depth; - for _ in 0..layer_size { - // merkle tree is using level form representation, so left and right siblings are - // next to each other - let left = tree.nodes[child_offset]; - let right = tree.nodes[child_offset + 1]; - self.nodes.insert( - tree.nodes[parent_offset].into(), - Node { - left: left.into(), - right: right.into(), - }, - ); - parent_offset += 1; - child_offset += 2; - } - depth += 1; + for node in tree.inner_nodes() { + self.nodes.insert( + node.value.into(), + Node { + left: node.left.into(), + right: node.right.into(), + }, + ); } - Ok(tree.nodes[1]) + Ok(tree.root()) } /// Adds a Sparse Merkle tree defined by the specified `entries` to the store, and returns the From 429d3bab6fb91541ad59660b5ba6a3a0064a30e9 Mon Sep 17 00:00:00 2001 From: "Augusto F. Hack" Date: Fri, 31 Mar 2023 04:47:12 +0200 Subject: [PATCH 05/10] feat: add support for MMR to the MerkleStore --- src/merkle/mmr/full.rs | 97 ++++++++++++++++++++++++++++++++++++++++- src/merkle/mmr/tests.rs | 40 ++++++++++++++++- src/merkle/store/mod.rs | 29 ++++++++++++ 3 files changed, 163 insertions(+), 3 deletions(-) diff --git a/src/merkle/mmr/full.rs b/src/merkle/mmr/full.rs index 92c350c..76f1e3d 100644 --- a/src/merkle/mmr/full.rs +++ b/src/merkle/mmr/full.rs @@ -11,8 +11,10 @@ //! merged, creating a new tree with depth d+1, this process is continued until the property is //! restabilished. use super::bit::TrueBitPositionIterator; -use super::{super::Vec, MmrPeaks, MmrProof, Rpo256, Word}; -use crate::merkle::MerklePath; +use super::{ + super::{InnerNodeInfo, MerklePath, Vec}, + MmrPeaks, MmrProof, Rpo256, Word, +}; use core::fmt::{Display, Formatter}; #[cfg(feature = "std")] @@ -190,6 +192,16 @@ impl Mmr { } } + /// An iterator over inner nodes in the [Mmm]. The order of iteration is unspecified. + pub fn inner_nodes(&self) -> MmrNodes { + MmrNodes { + mmr: self, + forest: 0, + last_right: 0, + index: 0, + } + } + // UTILITIES // ============================================================================================ @@ -246,6 +258,87 @@ where } } +// ITERATOR +// =============================================================================================== + +/// Yields inner nodes of the [Mmr]. +pub struct MmrNodes<'a> { + /// [Mmr] being yielded, when its `forest` value is matched, the iterations is finished. + mmr: &'a Mmr, + /// Keeps track of the left nodes yielded so far waiting for a right pair, this matches the + /// semantics of the [Mmr]'s forest attribute, since that too works as a buffer of left nodes + /// waiting for a pair to be hashed together. + forest: usize, + /// Keeps track of the last right node yielded, after this value is set, the next iteration + /// will be its parent with its corresponding left node that has been yield already. + last_right: usize, + /// The current index in the `nodes` vector. + index: usize, +} + +impl<'a> Iterator for MmrNodes<'a> { + type Item = InnerNodeInfo; + + fn next(&mut self) -> Option { + debug_assert!( + self.last_right.count_ones() <= 1, + "last_right tracks zero or one element" + ); + + // only parent nodes are emitted, remove the single node tree from the forest + let target = self.mmr.forest & (usize::MAX << 1); + + if self.forest < target { + if self.last_right == 0 { + // yield the left leaf + debug_assert!(self.last_right == 0, "left must be before right"); + self.forest |= 1; + self.index += 1; + + // yield the right leaf + debug_assert!((self.forest & 1) == 1, "right must be after left"); + self.last_right |= 1; + self.index += 1; + }; + + debug_assert!( + self.forest & self.last_right != 0, + "parent requires both a left and right", + ); + + // compute the number of nodes in the right tree, this is the offset to the + // previous left parent + let right_nodes = nodes_in_forest(self.last_right); + // the next parent position is one above the position of the pair + let parent = self.last_right << 1; + + // the left node has been paired and the current parent yielded, removed it from the forest + self.forest ^= self.last_right; + if self.forest & parent == 0 { + // this iteration yielded the left parent node + debug_assert!(self.forest & 1 == 0, "next iteration yields a left leaf"); + self.last_right = 0; + self.forest ^= parent; + } else { + // the left node of the parent level has been yielded already, this iteration + // was the right parent. Next iteration yields their parent. + self.last_right = parent; + } + + // yields a parent + let value = self.mmr.nodes[self.index]; + let right = self.mmr.nodes[self.index - 1]; + let left = self.mmr.nodes[self.index - 1 - right_nodes]; + self.index += 1; + let node = InnerNodeInfo { value, left, right }; + + Some(node) + } else { + None + } + } +} + // UTILITIES // =============================================================================================== diff --git a/src/merkle/mmr/tests.rs b/src/merkle/mmr/tests.rs index 577c6c4..1fb879e 100644 --- a/src/merkle/mmr/tests.rs +++ b/src/merkle/mmr/tests.rs @@ -1,6 +1,9 @@ use super::bit::TrueBitPositionIterator; use super::full::{high_bitmask, leaf_to_corresponding_tree, nodes_in_forest}; -use super::{super::Vec, Mmr, Rpo256, Word}; +use super::{ + super::{InnerNodeInfo, Vec}, + Mmr, Rpo256, Word, +}; use crate::merkle::{int_to_node, MerklePath}; #[test] @@ -410,6 +413,41 @@ fn test_bit_position_iterator() { ); } +#[test] +fn test_mmr_inner_nodes() { + let mmr: Mmr = LEAVES.into(); + let nodes: Vec = mmr.inner_nodes().collect(); + + let h01 = *Rpo256::hash_elements(&[LEAVES[0], LEAVES[1]].concat()); + let h23 = *Rpo256::hash_elements(&[LEAVES[2], LEAVES[3]].concat()); + let h0123 = *Rpo256::hash_elements(&[h01, h23].concat()); + let h45 = *Rpo256::hash_elements(&[LEAVES[4], LEAVES[5]].concat()); + let postorder = vec![ + InnerNodeInfo { + value: h01, + left: LEAVES[0], + right: LEAVES[1], + }, + InnerNodeInfo { + value: h23, + left: LEAVES[2], + right: LEAVES[3], + }, + InnerNodeInfo { + value: h0123, + left: h01, + right: h23, + }, + InnerNodeInfo { + value: h45, + left: LEAVES[4], + right: LEAVES[5], + }, + ]; + + assert_eq!(postorder, nodes); +} + mod property_tests { use super::leaf_to_corresponding_tree; use proptest::prelude::*; diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index 557942c..d1cffb6 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -1,3 +1,4 @@ +use super::mmr::{Mmr, MmrPeaks}; use super::{ BTreeMap, BTreeSet, EmptySubtreeRoots, MerkleError, MerklePath, MerklePathSet, MerkleTree, NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, ValuePath, Vec, Word, @@ -151,6 +152,15 @@ impl MerkleStore { Ok(self) } + /// Appends the provided [Mmr] represented by its `leaves` to the set. + pub fn with_mmr(mut self, leaves: I) -> Result + where + I: IntoIterator, + { + self.add_mmr(leaves)?; + Ok(self) + } + // PUBLIC ACCESSORS // -------------------------------------------------------------------------------------------- @@ -375,6 +385,25 @@ impl MerkleStore { Ok(root) } + /// Appends the provided [Mmr] into the store. + pub fn add_mmr(&mut self, leaves: I) -> Result + where + I: IntoIterator, + { + let mmr = Mmr::from(leaves); + for node in mmr.inner_nodes() { + self.nodes.insert( + node.value.into(), + Node { + left: node.left.into(), + right: node.right.into(), + }, + ); + } + + Ok(mmr.accumulator()) + } + /// Sets a node to `value`. /// /// # Errors From 9275dd00adf4f272d1a8c53f59c3cc6202f0150f Mon Sep 17 00:00:00 2001 From: "Augusto F. Hack" Date: Mon, 3 Apr 2023 19:08:42 +0200 Subject: [PATCH 06/10] feat: add parent node iterator for SimpleSMT --- src/merkle/simple_smt/mod.rs | 26 +++++++++++++++---- src/merkle/simple_smt/tests.rs | 47 +++++++++++++++++++++++++++++++++- src/merkle/store/mod.rs | 9 +++---- 3 files changed, 71 insertions(+), 11 deletions(-) diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs index 9699ea9..85fec0e 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/simple_smt/mod.rs @@ -1,5 +1,6 @@ use super::{ - BTreeMap, EmptySubtreeRoots, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, Vec, Word, + BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Rpo256, + RpoDigest, Vec, Word, }; #[cfg(test)] @@ -15,14 +16,20 @@ pub struct SimpleSmt { depth: u8, root: Word, leaves: BTreeMap, - pub(crate) branches: BTreeMap, + branches: BTreeMap, empty_hashes: Vec, } #[derive(Debug, Default, Clone, PartialEq, Eq)] -pub(crate) struct BranchNode { - pub(crate) left: RpoDigest, - pub(crate) right: RpoDigest, +struct BranchNode { + left: RpoDigest, + right: RpoDigest, +} + +impl BranchNode { + fn parent(&self) -> RpoDigest { + Rpo256::merge(&[self.left, self.right]) + } } impl SimpleSmt { @@ -171,6 +178,15 @@ impl SimpleSmt { self.get_path(NodeIndex::new(self.depth(), key)) } + /// Iterator over the inner nodes of the [SimpleSmt]. + pub fn inner_nodes(&self) -> impl Iterator + '_ { + self.branches.values().map(|e| InnerNodeInfo { + value: e.parent().into(), + left: e.left.into(), + right: e.right.into(), + }) + } + // STATE MUTATORS // -------------------------------------------------------------------------------------------- diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs index 6abd343..f7b8635 100644 --- a/src/merkle/simple_smt/tests.rs +++ b/src/merkle/simple_smt/tests.rs @@ -1,5 +1,5 @@ use super::{ - super::{int_to_node, MerkleTree, RpoDigest, SimpleSmt}, + super::{int_to_node, InnerNodeInfo, MerkleError, MerkleTree, RpoDigest, SimpleSmt}, NodeIndex, Rpo256, Vec, Word, }; use proptest::prelude::*; @@ -138,6 +138,51 @@ fn get_path() { assert_eq!(vec![node2], *tree.get_path(NodeIndex::new(1, 1)).unwrap()); } +#[test] +fn test_parent_node_iterator() -> Result<(), MerkleError> { + let tree = SimpleSmt::new(2) + .unwrap() + .with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter())) + .unwrap(); + + // check depth 2 + assert_eq!(VALUES4[0], tree.get_node(&NodeIndex::new(2, 0)).unwrap()); + assert_eq!(VALUES4[1], tree.get_node(&NodeIndex::new(2, 1)).unwrap()); + assert_eq!(VALUES4[2], tree.get_node(&NodeIndex::new(2, 2)).unwrap()); + assert_eq!(VALUES4[3], tree.get_node(&NodeIndex::new(2, 3)).unwrap()); + + // get parent nodes + let root = tree.root(); + let l1n0 = tree.get_node(&NodeIndex::new(1, 0))?; + let l1n1 = tree.get_node(&NodeIndex::new(1, 1))?; + let l2n0 = tree.get_node(&NodeIndex::new(2, 0))?; + let l2n1 = tree.get_node(&NodeIndex::new(2, 1))?; + let l2n2 = tree.get_node(&NodeIndex::new(2, 2))?; + let l2n3 = tree.get_node(&NodeIndex::new(2, 3))?; + + let nodes: Vec = tree.inner_nodes().collect(); + let expected = vec![ + InnerNodeInfo { + value: root.into(), + left: l1n0.into(), + right: l1n1.into(), + }, + InnerNodeInfo { + value: l1n0.into(), + left: l2n0.into(), + right: l2n1.into(), + }, + InnerNodeInfo { + value: l1n1.into(), + left: l2n2.into(), + right: l2n3.into(), + }, + ]; + assert_eq!(nodes, expected); + + Ok(()) +} + #[test] fn update_leaf() { let mut tree = SimpleSmt::new(3) diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index 557942c..7d47912 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -281,13 +281,12 @@ impl MerkleStore { I: Iterator + ExactSizeIterator, { let smt = SimpleSmt::new(depth)?.with_leaves(entries)?; - for branch in smt.branches.values() { - let parent = Rpo256::merge(&[branch.left, branch.right]); + for node in smt.inner_nodes() { self.nodes.insert( - parent, + node.value.into(), Node { - left: branch.left, - right: branch.right, + left: node.left.into(), + right: node.right.into(), }, ); } From bd557bc68c688fe5dc60352371bc28bd89ea2634 Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Wed, 5 Apr 2023 02:23:20 -0700 Subject: [PATCH 07/10] fix: add validation to NodeIndex constructor and remove BitIterator --- benches/store.rs | 99 ++++++++-------- src/bit.rs | 169 --------------------------- src/lib.rs | 1 - src/merkle/index.rs | 104 +++++++++++------ src/merkle/merkle_tree.rs | 57 ++++----- src/merkle/mod.rs | 8 +- src/merkle/path.rs | 16 ++- src/merkle/path_set.rs | 35 ++---- src/merkle/simple_smt/mod.rs | 19 +-- src/merkle/simple_smt/tests.rs | 56 ++++----- src/merkle/store/mod.rs | 70 +++++------ src/merkle/store/tests.rs | 206 ++++++++++++++++++--------------- 12 files changed, 338 insertions(+), 502 deletions(-) delete mode 100644 src/bit.rs diff --git a/benches/store.rs b/benches/store.rs index b59a537..804faa6 100644 --- a/benches/store.rs +++ b/benches/store.rs @@ -18,9 +18,10 @@ fn random_word() -> Word { rand_array::().into() } -/// Generates a u64 in `0..range`. -fn random_index(range: u64) -> u64 { - rand_value::() % range +/// Generates an index at the specified depth in `0..range`. +fn random_index(range: u64, depth: u8) -> NodeIndex { + let value = rand_value::() % range; + NodeIndex::new(depth, value).unwrap() } /// Benchmarks getting an empty leaf from the SMT and MerkleStore backends. @@ -28,7 +29,7 @@ fn get_empty_leaf_simplesmt(c: &mut Criterion) { let mut group = c.benchmark_group("get_empty_leaf_simplesmt"); let depth = SimpleSmt::MAX_DEPTH; - let size = 2u64.pow(depth as u32); + let size = u64::MAX; // both SMT and the store are pre-populated with empty hashes, accessing these values is what is // being benchmarked here, so no values are inserted into the backends @@ -38,16 +39,16 @@ fn get_empty_leaf_simplesmt(c: &mut Criterion) { group.bench_function(BenchmarkId::new("SimpleSmt", depth), |b| { b.iter_batched( - || random_index(size), - |value| black_box(smt.get_node(&NodeIndex::new(depth, value))), + || random_index(size, depth), + |index| black_box(smt.get_node(index)), BatchSize::SmallInput, ) }); group.bench_function(BenchmarkId::new("MerkleStore", depth), |b| { b.iter_batched( - || random_index(size), - |value| black_box(store.get_node(root, NodeIndex::new(depth, value))), + || random_index(size, depth), + |index| black_box(store.get_node(root, index)), BatchSize::SmallInput, ) }); @@ -72,16 +73,16 @@ fn get_leaf_merkletree(c: &mut Criterion) { group.bench_function(BenchmarkId::new("MerkleTree", size), |b| { b.iter_batched( - || random_index(size_u64), - |value| black_box(mtree.get_node(NodeIndex::new(depth, value))), + || random_index(size_u64, depth), + |index| black_box(mtree.get_node(index)), BatchSize::SmallInput, ) }); group.bench_function(BenchmarkId::new("MerkleStore", size), |b| { b.iter_batched( - || random_index(size_u64), - |value| black_box(store.get_node(root, NodeIndex::new(depth, value))), + || random_index(size_u64, depth), + |index| black_box(store.get_node(root, index)), BatchSize::SmallInput, ) }); @@ -116,16 +117,16 @@ fn get_leaf_simplesmt(c: &mut Criterion) { group.bench_function(BenchmarkId::new("SimpleSmt", size), |b| { b.iter_batched( - || random_index(size_u64), - |value| black_box(smt.get_node(&NodeIndex::new(depth, value))), + || random_index(size_u64, depth), + |index| black_box(smt.get_node(index)), BatchSize::SmallInput, ) }); group.bench_function(BenchmarkId::new("MerkleStore", size), |b| { b.iter_batched( - || random_index(size_u64), - |value| black_box(store.get_node(root, NodeIndex::new(depth, value))), + || random_index(size_u64, depth), + |index| black_box(store.get_node(root, index)), BatchSize::SmallInput, ) }); @@ -137,7 +138,6 @@ fn get_node_of_empty_simplesmt(c: &mut Criterion) { let mut group = c.benchmark_group("get_node_of_empty_simplesmt"); let depth = SimpleSmt::MAX_DEPTH; - let size = 2u64.pow(depth as u32); // both SMT and the store are pre-populated with the empty hashes, accessing the internal nodes // of these values is what is being benchmarked here, so no values are inserted into the @@ -146,19 +146,20 @@ fn get_node_of_empty_simplesmt(c: &mut Criterion) { let store = MerkleStore::new(); let root = smt.root(); let half_depth = depth / 2; + let half_size = 2_u64.pow(half_depth as u32); group.bench_function(BenchmarkId::new("SimpleSmt", depth), |b| { b.iter_batched( - || random_index(size), - |value| black_box(smt.get_node(&NodeIndex::new(half_depth, value))), + || random_index(half_size, half_depth), + |index| black_box(smt.get_node(index)), BatchSize::SmallInput, ) }); group.bench_function(BenchmarkId::new("MerkleStore", depth), |b| { b.iter_batched( - || random_index(size), - |value| black_box(store.get_node(root, NodeIndex::new(half_depth, value))), + || random_index(half_size, half_depth), + |index| black_box(store.get_node(root, index)), BatchSize::SmallInput, ) }); @@ -178,22 +179,22 @@ fn get_node_merkletree(c: &mut Criterion) { let mtree_leaves: Vec = leaves.iter().map(|v| v.into()).collect(); let mtree = MerkleTree::new(mtree_leaves.clone()).unwrap(); let store = MerkleStore::new().with_merkle_tree(mtree_leaves).unwrap(); - let half_depth = mtree.depth() / 2; let root = mtree.root(); - let size_u64 = size as u64; + let half_depth = mtree.depth() / 2; + let half_size = 2_u64.pow(half_depth as u32); group.bench_function(BenchmarkId::new("MerkleTree", size), |b| { b.iter_batched( - || random_index(size_u64), - |value| black_box(mtree.get_node(NodeIndex::new(half_depth, value))), + || random_index(half_size, half_depth), + |index| black_box(mtree.get_node(index)), BatchSize::SmallInput, ) }); group.bench_function(BenchmarkId::new("MerkleStore", size), |b| { b.iter_batched( - || random_index(size_u64), - |value| black_box(store.get_node(root, NodeIndex::new(half_depth, value))), + || random_index(half_size, half_depth), + |index| black_box(store.get_node(root, index)), BatchSize::SmallInput, ) }); @@ -224,21 +225,21 @@ fn get_node_simplesmt(c: &mut Criterion) { .with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, smt_leaves) .unwrap(); let root = smt.root(); - let size_u64 = size as u64; let half_depth = smt.depth() / 2; + let half_size = 2_u64.pow(half_depth as u32); group.bench_function(BenchmarkId::new("SimpleSmt", size), |b| { b.iter_batched( - || random_index(size_u64), - |value| black_box(smt.get_node(&NodeIndex::new(half_depth, value))), + || random_index(half_size, half_depth), + |index| black_box(smt.get_node(index)), BatchSize::SmallInput, ) }); group.bench_function(BenchmarkId::new("MerkleStore", size), |b| { b.iter_batched( - || random_index(size_u64), - |value| black_box(store.get_node(root, NodeIndex::new(half_depth, value))), + || random_index(half_size, half_depth), + |index| black_box(store.get_node(root, index)), BatchSize::SmallInput, ) }); @@ -264,16 +265,16 @@ fn get_leaf_path_merkletree(c: &mut Criterion) { group.bench_function(BenchmarkId::new("MerkleTree", size), |b| { b.iter_batched( - || random_index(size_u64), - |value| black_box(mtree.get_path(NodeIndex::new(depth, value))), + || random_index(size_u64, depth), + |index| black_box(mtree.get_path(index)), BatchSize::SmallInput, ) }); group.bench_function(BenchmarkId::new("MerkleStore", size), |b| { b.iter_batched( - || random_index(size_u64), - |value| black_box(store.get_path(root, NodeIndex::new(depth, value))), + || random_index(size_u64, depth), + |index| black_box(store.get_path(root, index)), BatchSize::SmallInput, ) }); @@ -308,16 +309,16 @@ fn get_leaf_path_simplesmt(c: &mut Criterion) { group.bench_function(BenchmarkId::new("SimpleSmt", size), |b| { b.iter_batched( - || random_index(size_u64), - |value| black_box(smt.get_path(NodeIndex::new(depth, value))), + || random_index(size_u64, depth), + |index| black_box(smt.get_path(index)), BatchSize::SmallInput, ) }); group.bench_function(BenchmarkId::new("MerkleStore", size), |b| { b.iter_batched( - || random_index(size_u64), - |value| black_box(store.get_path(root, NodeIndex::new(depth, value))), + || random_index(size_u64, depth), + |index| black_box(store.get_path(root, index)), BatchSize::SmallInput, ) }); @@ -413,7 +414,7 @@ fn update_leaf_merkletree(c: &mut Criterion) { group.bench_function(BenchmarkId::new("MerkleTree", size), |b| { b.iter_batched( - || (random_index(size_u64), random_word()), + || (rand_value::() % size_u64, random_word()), |(index, value)| black_box(mtree.update_leaf(index, value)), BatchSize::SmallInput, ) @@ -422,15 +423,12 @@ fn update_leaf_merkletree(c: &mut Criterion) { let mut store_root = root; group.bench_function(BenchmarkId::new("MerkleStore", size), |b| { b.iter_batched( - || (random_index(size_u64), random_word()), + || (random_index(size_u64, depth), random_word()), |(index, value)| { // The MerkleTree automatically updates its internal root, the Store maintains // the old root and adds the new one. Here we update the root to have a fair // comparison - store_root = store - .set_node(root, NodeIndex::new(depth, index), value) - .unwrap() - .root; + store_root = store.set_node(root, index, value).unwrap().root; black_box(store_root) }, BatchSize::SmallInput, @@ -467,7 +465,7 @@ fn update_leaf_simplesmt(c: &mut Criterion) { group.bench_function(BenchmarkId::new("SimpleSMT", size), |b| { b.iter_batched( - || (random_index(size_u64), random_word()), + || (rand_value::() % size_u64, random_word()), |(index, value)| black_box(smt.update_leaf(index, value)), BatchSize::SmallInput, ) @@ -476,15 +474,12 @@ fn update_leaf_simplesmt(c: &mut Criterion) { let mut store_root = root; group.bench_function(BenchmarkId::new("MerkleStore", size), |b| { b.iter_batched( - || (random_index(size_u64), random_word()), + || (random_index(size_u64, depth), random_word()), |(index, value)| { // The MerkleTree automatically updates its internal root, the Store maintains // the old root and adds the new one. Here we update the root to have a fair // comparison - store_root = store - .set_node(root, NodeIndex::new(depth, index), value) - .unwrap() - .root; + store_root = store.set_node(root, index, value).unwrap().root; black_box(store_root) }, BatchSize::SmallInput, diff --git a/src/bit.rs b/src/bit.rs deleted file mode 100644 index 5eb2577..0000000 --- a/src/bit.rs +++ /dev/null @@ -1,169 +0,0 @@ -/// Yields the bits of a `u64`. -pub struct BitIterator { - /// The value that is being iterated bit-wise - value: u64, - /// True bits in the `mask` are the bits that have been visited. - mask: u64, -} - -impl BitIterator { - pub fn new(value: u64) -> BitIterator { - BitIterator { value, mask: 0 } - } - - /// An efficient skip implementation. - /// - /// Note: The compiler is smart enough to translate a `skip(n)` into a single shift instruction - /// if the code is inlined, however inlining does not always happen. - pub fn skip_front(mut self, n: u32) -> Self { - let mask = bitmask(n); - let ones = self.mask.trailing_ones(); - let mask_position = ones; - self.mask ^= mask.checked_shl(mask_position).unwrap_or(0); - self - } - - /// An efficient skip from the back. - /// - /// Note: The compiler is smart enough to translate a `skip(n)` into a single shift instruction - /// if the code is inlined, however inlining does not always happen. - pub fn skip_back(mut self, n: u32) -> Self { - let mask = bitmask(n); - let ones = self.mask.leading_ones(); - let mask_position = u64::BITS - ones - n; - self.mask ^= mask.checked_shl(mask_position).unwrap_or(0); - self - } -} - -impl Iterator for BitIterator { - type Item = bool; - - fn next(&mut self) -> Option<::Item> { - // trailing_ones is implemented with trailing_zeros, and the zeros are computed with the - // intrinsic cttz. [Rust 1.67.0] x86 uses the `bsf` instruction. AArch64 uses the `rbit - // clz` instructions. - let ones = self.mask.trailing_ones(); - - if ones == u64::BITS { - None - } else { - let bit_position = ones; - let mask = 1 << bit_position; - self.mask ^= mask; - let bit = self.value & mask; - Some(bit != 0) - } - } -} - -impl DoubleEndedIterator for BitIterator { - fn next_back(&mut self) -> Option<::Item> { - // leading_ones is implemented with leading_zeros, and the zeros are computed with the - // intrinsic ctlz. [Rust 1.67.0] x86 uses the `bsr` instruction. AArch64 uses the `clz` - // instruction. - let ones = self.mask.leading_ones(); - - if ones == u64::BITS { - None - } else { - let bit_position = u64::BITS - ones - 1; - let mask = 1 << bit_position; - self.mask ^= mask; - let bit = self.value & mask; - Some(bit != 0) - } - } -} - -#[cfg(test)] -mod test { - use super::BitIterator; - - #[test] - fn test_bit_iterator() { - let v = 0b1; - let mut it = BitIterator::new(v); - assert!(it.next().unwrap(), "first bit is true"); - assert!(it.all(|v| v == false), "every other value is false"); - - let v = 0b10; - let mut it = BitIterator::new(v); - assert!(!it.next().unwrap(), "first bit is false"); - assert!(it.next().unwrap(), "first bit is true"); - assert!(it.all(|v| v == false), "every other value is false"); - - let v = 0b10; - let mut it = BitIterator::new(v); - assert!(!it.next_back().unwrap(), "last bit is false"); - assert!(!it.next().unwrap(), "first bit is false"); - assert!(it.next().unwrap(), "first bit is true"); - assert!(it.all(|v| v == false), "every other value is false"); - } - - #[test] - fn test_bit_iterator_skip() { - let v = 0b1; - let mut it = BitIterator::new(v).skip_front(1); - assert!(it.all(|v| v == false), "every other value is false"); - - let v = 0b10; - let mut it = BitIterator::new(v).skip_front(1); - assert!(it.next().unwrap(), "first bit is true"); - assert!(it.all(|v| v == false), "every other value is false"); - - let high_bit = 0b1 << (u64::BITS - 1); - let mut it = BitIterator::new(high_bit).skip_back(1); - assert!(it.all(|v| v == false), "every other value is false"); - - let v = 0b10; - let mut it = BitIterator::new(v).skip_back(1); - assert!(!it.next_back().unwrap(), "last bit is false"); - assert!(!it.next().unwrap(), "first bit is false"); - assert!(it.next().unwrap(), "first bit is true"); - assert!(it.all(|v| v == false), "every other value is false"); - } - - #[test] - fn test_skip_all() { - let v = 0b1; - let mut it = BitIterator::new(v).skip_front(u64::BITS); - assert!(it.next().is_none(), "iterator must be exhausted"); - - let v = 0b1; - let mut it = BitIterator::new(v).skip_back(u64::BITS); - assert!(it.next().is_none(), "iterator must be exhausted"); - } - - #[test] - fn test_bit_iterator_count_bits_after_skip() { - let any_value = 0b1; - for s in 0..u64::BITS { - let it = BitIterator::new(any_value).skip_front(s); - assert_eq!(it.count() as u32, u64::BITS - s) - } - - let any_value = 0b1; - for s in 1..u64::BITS { - let it = BitIterator::new(any_value).skip_back(s); - assert_eq!(it.count() as u32, u64::BITS - s) - } - } - - #[test] - fn test_bit_iterator_rev() { - let v = 0b1; - let mut it = BitIterator::new(v).rev(); - assert!(it.nth(63).unwrap(), "the last value is true"); - } -} - -// UTILITIES -// =============================================================================================== - -fn bitmask(s: u32) -> u64 { - match 1u64.checked_shl(s) { - Some(r) => r - 1, - None => u64::MAX, - } -} diff --git a/src/lib.rs b/src/lib.rs index 7701d36..7c7d753 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,7 +4,6 @@ #[cfg_attr(test, macro_use)] extern crate alloc; -mod bit; pub mod hash; pub mod merkle; pub mod utils; diff --git a/src/merkle/index.rs b/src/merkle/index.rs index 5729083..b4f3bdf 100644 --- a/src/merkle/index.rs +++ b/src/merkle/index.rs @@ -1,13 +1,23 @@ use super::{Felt, MerkleError, RpoDigest, StarkField}; -use crate::bit::BitIterator; // NODE INDEX // ================================================================================================ -/// A Merkle tree address to an arbitrary node. +/// Address to an arbitrary node in a binary tree using level order form. /// -/// The position is relative to a tree in level order, where for a given depth `d` elements are -/// numbered from $0..2^d$. +/// The position is represented by the pair `(depth, pos)`, where for a given depth `d` elements +/// are numbered from $0..(2^d)-1$. Example: +/// +/// ```ignore +/// depth +/// 0 0 +/// 1 0 1 +/// 2 0 1 2 3 +/// 3 0 1 2 3 4 5 6 7 +/// ``` +/// +/// The root is represented by the pair $(0, 0)$, its left child is $(1, 0)$ and its right child +/// $(1, 1)$. #[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)] pub struct NodeIndex { depth: u8, @@ -19,20 +29,37 @@ impl NodeIndex { // -------------------------------------------------------------------------------------------- /// Creates a new node index. - pub const fn new(depth: u8, value: u64) -> Self { - Self { depth, value } + /// + /// # Errors + /// Returns an error if the `value` is greater than or equal to 2^{depth}. + pub const fn new(depth: u8, value: u64) -> Result { + if (64 - value.leading_zeros()) > depth as u32 { + Err(MerkleError::InvalidIndex { depth, value }) + } else { + Ok(Self { depth, value }) + } + } + + /// Creates a new node index for testing purposes. + /// + /// # Panics + /// Panics if the `value` is greater than or equal to 2^{depth}. + #[cfg(test)] + pub fn make(depth: u8, value: u64) -> Self { + Self::new(depth, value).unwrap() } /// Creates a node index from a pair of field elements representing the depth and value. /// /// # Errors - /// - /// Will error if the `u64` representation of the depth doesn't fit a `u8`. + /// Returns an error if: + /// - `depth` doesn't fit in a `u8`. + /// - `value` is greater than or equal to 2^{depth}. pub fn from_elements(depth: &Felt, value: &Felt) -> Result { let depth = depth.as_int(); let depth = u8::try_from(depth).map_err(|_| MerkleError::DepthTooBig(depth))?; let value = value.as_int(); - Ok(Self::new(depth, value)) + Self::new(depth, value) } /// Creates a new node index pointing to the root of the tree. @@ -40,12 +67,6 @@ impl NodeIndex { Self { depth: 0, value: 0 } } - /// Mutates the instance and returns it, replacing the depth. - pub const fn with_depth(mut self, depth: u8) -> Self { - self.depth = depth; - self - } - /// Computes the value of the sibling of the current node. pub fn sibling(mut self) -> Self { self.value ^= 1; @@ -83,11 +104,6 @@ impl NodeIndex { self.value } - /// Returns true if the current value fits the current depth for a binary tree. - pub const fn is_valid(&self) -> bool { - self.value < (1 << self.depth as u64) - } - /// Returns true if the current instance points to a right sibling node. pub const fn is_value_odd(&self) -> bool { (self.value & 1) == 1 @@ -98,19 +114,6 @@ impl NodeIndex { self.depth == 0 } - /// Returns a bit iterator for the `value`. - /// - /// Bits read from left-to-right represent which internal node's child should be visited to - /// arrive at the leaf. From the right-to-left the bit represent the position the hash of the - /// current element should go. - /// - /// Additionally, the value that is not visited are the sibling values necessary for a Merkle - /// opening. - pub fn bit_iterator(&self) -> BitIterator { - let depth: u32 = self.depth.into(); - BitIterator::new(self.value).skip_back(u64::BITS - depth) - } - // STATE MUTATORS // -------------------------------------------------------------------------------------------- @@ -127,14 +130,43 @@ mod tests { use super::*; use proptest::prelude::*; + #[test] + fn test_node_index_value_too_high() { + assert_eq!( + NodeIndex::new(0, 0).unwrap(), + NodeIndex { depth: 0, value: 0 } + ); + match NodeIndex::new(0, 1) { + Err(MerkleError::InvalidIndex { depth, value }) => { + assert_eq!(depth, 0); + assert_eq!(value, 1); + } + _ => unreachable!(), + } + } + + #[test] + fn test_node_index_can_represent_depth_64() { + assert!(NodeIndex::new(64, u64::MAX).is_ok()); + } + + prop_compose! { + fn node_index()(value in 0..2u64.pow(u64::BITS - 1)) -> NodeIndex { + // unwrap never panics because the range of depth is 0..u64::BITS + let mut depth = value.ilog2() as u8; + if value > (1 << depth) { // round up + depth += 1; + } + NodeIndex::new(depth, value.into()).unwrap() + } + } + proptest! { #[test] fn arbitrary_index_wont_panic_on_move_up( - depth in prop::num::u8::ANY, - value in prop::num::u64::ANY, + mut index in node_index(), count in prop::num::u8::ANY, ) { - let mut index = NodeIndex::new(depth, value); for _ in 0..count { index.move_up(); } diff --git a/src/merkle/merkle_tree.rs b/src/merkle/merkle_tree.rs index 6b767ae..3587fc6 100644 --- a/src/merkle/merkle_tree.rs +++ b/src/merkle/merkle_tree.rs @@ -79,8 +79,6 @@ impl MerkleTree { return Err(MerkleError::DepthTooSmall(index.depth())); } else if index.depth() > self.depth() { return Err(MerkleError::DepthTooBig(index.depth() as u64)); - } else if !index.is_valid() { - return Err(MerkleError::InvalidIndex(index)); } let pos = index.to_scalar_index() as usize; @@ -99,8 +97,6 @@ impl MerkleTree { return Err(MerkleError::DepthTooSmall(index.depth())); } else if index.depth() > self.depth() { return Err(MerkleError::DepthTooBig(index.depth() as u64)); - } else if !index.is_valid() { - return Err(MerkleError::InvalidIndex(index)); } // TODO should we create a helper in `NodeIndex` that will encapsulate traversal to root so @@ -126,11 +122,7 @@ impl MerkleTree { /// # Errors /// Returns an error if the specified index value is not a valid leaf value for this tree. pub fn update_leaf<'a>(&'a mut self, index_value: u64, value: Word) -> Result<(), MerkleError> { - let depth = self.depth(); - let mut index = NodeIndex::new(depth, index_value); - if !index.is_valid() { - return Err(MerkleError::InvalidIndex(index)); - } + let mut index = NodeIndex::new(self.depth(), index_value)?; // we don't need to copy the pairs into a new address as we are logically guaranteed to not // overlap write instructions. however, it's important to bind the lifetime of pairs to @@ -203,7 +195,7 @@ impl<'a> Iterator for MerkleTreeNodes<'a> { } } -/// Utility to vizualize a [MerkleTree] in text. +/// Utility to visualize a [MerkleTree] in text. pub fn tree_to_text(tree: &MerkleTree) -> Result { let indent = " "; let mut s = String::new(); @@ -212,11 +204,8 @@ pub fn tree_to_text(tree: &MerkleTree) -> Result { for d in 1..=tree.depth() { let entries = 2u64.pow(d.into()); for i in 0..entries { - let index = NodeIndex::new(d, i); - - let node = tree - .get_node(index) - .expect("The index must always be valid"); + let index = NodeIndex::new(d, i).expect("The index must always be valid"); + let node = tree.get_node(index).expect("The node must always be found"); for _ in 0..d { s.push_str(indent); @@ -229,7 +218,7 @@ pub fn tree_to_text(tree: &MerkleTree) -> Result { Ok(s) } -/// Utility to vizualize a [MerklePath] in text. +/// Utility to visualize a [MerklePath] in text. pub fn path_to_text(path: &MerklePath) -> Result { let mut s = String::new(); s.push('['); @@ -301,16 +290,16 @@ mod tests { let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap(); // check depth 2 - assert_eq!(LEAVES4[0], tree.get_node(NodeIndex::new(2, 0)).unwrap()); - assert_eq!(LEAVES4[1], tree.get_node(NodeIndex::new(2, 1)).unwrap()); - assert_eq!(LEAVES4[2], tree.get_node(NodeIndex::new(2, 2)).unwrap()); - assert_eq!(LEAVES4[3], tree.get_node(NodeIndex::new(2, 3)).unwrap()); + assert_eq!(LEAVES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap()); + assert_eq!(LEAVES4[1], tree.get_node(NodeIndex::make(2, 1)).unwrap()); + assert_eq!(LEAVES4[2], tree.get_node(NodeIndex::make(2, 2)).unwrap()); + assert_eq!(LEAVES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap()); // check depth 1 let (_, node2, node3) = compute_internal_nodes(); - assert_eq!(node2, tree.get_node(NodeIndex::new(1, 0)).unwrap()); - assert_eq!(node3, tree.get_node(NodeIndex::new(1, 1)).unwrap()); + assert_eq!(node2, tree.get_node(NodeIndex::make(1, 0)).unwrap()); + assert_eq!(node3, tree.get_node(NodeIndex::make(1, 1)).unwrap()); } #[test] @@ -322,24 +311,24 @@ mod tests { // check depth 2 assert_eq!( vec![LEAVES4[1], node3], - *tree.get_path(NodeIndex::new(2, 0)).unwrap() + *tree.get_path(NodeIndex::make(2, 0)).unwrap() ); assert_eq!( vec![LEAVES4[0], node3], - *tree.get_path(NodeIndex::new(2, 1)).unwrap() + *tree.get_path(NodeIndex::make(2, 1)).unwrap() ); assert_eq!( vec![LEAVES4[3], node2], - *tree.get_path(NodeIndex::new(2, 2)).unwrap() + *tree.get_path(NodeIndex::make(2, 2)).unwrap() ); assert_eq!( vec![LEAVES4[2], node2], - *tree.get_path(NodeIndex::new(2, 3)).unwrap() + *tree.get_path(NodeIndex::make(2, 3)).unwrap() ); // check depth 1 - assert_eq!(vec![node3], *tree.get_path(NodeIndex::new(1, 0)).unwrap()); - assert_eq!(vec![node2], *tree.get_path(NodeIndex::new(1, 1)).unwrap()); + assert_eq!(vec![node3], *tree.get_path(NodeIndex::make(1, 0)).unwrap()); + assert_eq!(vec![node2], *tree.get_path(NodeIndex::make(1, 1)).unwrap()); } #[test] @@ -370,12 +359,12 @@ mod tests { fn nodes() -> Result<(), MerkleError> { let tree = super::MerkleTree::new(LEAVES4.to_vec()).unwrap(); let root = tree.root(); - let l1n0 = tree.get_node(NodeIndex::new(1, 0))?; - let l1n1 = tree.get_node(NodeIndex::new(1, 1))?; - let l2n0 = tree.get_node(NodeIndex::new(2, 0))?; - let l2n1 = tree.get_node(NodeIndex::new(2, 1))?; - let l2n2 = tree.get_node(NodeIndex::new(2, 2))?; - let l2n3 = tree.get_node(NodeIndex::new(2, 3))?; + let l1n0 = tree.get_node(NodeIndex::make(1, 0))?; + let l1n1 = tree.get_node(NodeIndex::make(1, 1))?; + let l2n0 = tree.get_node(NodeIndex::make(2, 0))?; + let l2n1 = tree.get_node(NodeIndex::make(2, 1))?; + let l2n2 = tree.get_node(NodeIndex::make(2, 2))?; + let l2n3 = tree.get_node(NodeIndex::make(2, 3))?; let nodes: Vec = tree.inner_nodes().collect(); let expected = vec![ diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index c12ac7a..4d36577 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -1,6 +1,6 @@ use super::{ hash::rpo::{Rpo256, RpoDigest}, - utils::collections::{vec, BTreeMap, BTreeSet, Vec}, + utils::collections::{vec, BTreeMap, Vec}, Felt, StarkField, Word, WORD_SIZE, ZERO, }; use core::fmt; @@ -45,7 +45,7 @@ pub enum MerkleError { DepthTooBig(u64), NodeNotInStore(Word, NodeIndex), NumLeavesNotPowerOfTwo(usize), - InvalidIndex(NodeIndex), + InvalidIndex { depth: u8, value: u64 }, InvalidDepth { expected: u8, provided: u8 }, InvalidPath(MerklePath), InvalidEntriesCount(usize, usize), @@ -63,9 +63,9 @@ impl fmt::Display for MerkleError { NumLeavesNotPowerOfTwo(leaves) => { write!(f, "the leaves count {leaves} is not a power of 2") } - InvalidIndex(index) => write!( + InvalidIndex{ depth, value} => write!( f, - "the index value {} is not valid for the depth {}", index.value(), index.depth() + "the index value {value} is not valid for the depth {depth}" ), InvalidDepth { expected, provided } => write!( f, diff --git a/src/merkle/path.rs b/src/merkle/path.rs index 9a4e46b..964f6d0 100644 --- a/src/merkle/path.rs +++ b/src/merkle/path.rs @@ -1,4 +1,4 @@ -use super::{vec, NodeIndex, Rpo256, Vec, Word}; +use super::{vec, MerkleError, NodeIndex, Rpo256, Vec, Word}; use core::ops::{Deref, DerefMut}; // MERKLE PATH @@ -23,14 +23,15 @@ impl MerklePath { // -------------------------------------------------------------------------------------------- /// Computes the merkle root for this opening. - pub fn compute_root(&self, index_value: u64, node: Word) -> Word { - let mut index = NodeIndex::new(self.depth(), index_value); - self.nodes.iter().copied().fold(node, |node, sibling| { + pub fn compute_root(&self, index: u64, node: Word) -> Result { + let mut index = NodeIndex::new(self.depth(), index)?; + let root = self.nodes.iter().copied().fold(node, |node, sibling| { // compute the node and move to the next iteration. let input = index.build_node(node.into(), sibling.into()); index.move_up(); Rpo256::merge(&input).into() - }) + }); + Ok(root) } /// Returns the depth in which this Merkle path proof is valid. @@ -42,7 +43,10 @@ impl MerklePath { /// /// Returns `true` if `node` exists at `index` in a Merkle tree with `root`. pub fn verify(&self, index: u64, node: Word, root: &Word) -> bool { - root == &self.compute_root(index, node) + match self.compute_root(index, node) { + Ok(computed_root) => root == &computed_root, + Err(_) => false, + } } } diff --git a/src/merkle/path_set.rs b/src/merkle/path_set.rs index 9a90687..ef3cc48 100644 --- a/src/merkle/path_set.rs +++ b/src/merkle/path_set.rs @@ -64,11 +64,6 @@ impl MerklePathSet { /// * The specified index is not valid for the depth of structure. /// * Requested node does not exist in the set. pub fn get_node(&self, index: NodeIndex) -> Result { - if !index.with_depth(self.total_depth).is_valid() { - return Err(MerkleError::InvalidIndex( - index.with_depth(self.total_depth), - )); - } if index.depth() != self.total_depth { return Err(MerkleError::InvalidDepth { expected: self.total_depth, @@ -90,7 +85,8 @@ impl MerklePathSet { /// * The specified index is not valid for the depth of the structure. /// * Leaf with the requested path does not exist in the set. pub fn get_leaf(&self, index: u64) -> Result { - self.get_node(NodeIndex::new(self.depth(), index)) + let index = NodeIndex::new(self.depth(), index)?; + self.get_node(index) } /// Returns a Merkle path to the node at the specified index. The node itself is @@ -101,9 +97,6 @@ impl MerklePathSet { /// * The specified index is not valid for the depth of structure. /// * Node of the requested path does not exist in the set. pub fn get_path(&self, index: NodeIndex) -> Result { - if !index.with_depth(self.total_depth).is_valid() { - return Err(MerkleError::InvalidIndex(index)); - } if index.depth() != self.total_depth { return Err(MerkleError::InvalidDepth { expected: self.total_depth, @@ -165,8 +158,7 @@ impl MerklePathSet { value: Word, mut path: MerklePath, ) -> Result<(), MerkleError> { - let depth = path.len() as u8; - let mut index = NodeIndex::new(depth, index_value); + let mut index = NodeIndex::new(path.len() as u8, index_value)?; if index.depth() != self.total_depth { return Err(MerkleError::InvalidDepth { expected: self.total_depth, @@ -190,7 +182,7 @@ impl MerklePathSet { if self.root == [ZERO; 4] { self.root = root; } else if self.root != root { - return Err(MerkleError::InvalidPath(path)); + return Err(MerkleError::ConflictingRoots([self.root, root].to_vec())); } // finish updating the path @@ -205,12 +197,7 @@ impl MerklePathSet { /// Returns an error if: /// * Requested node does not exist in the set. pub fn update_leaf(&mut self, base_index_value: u64, value: Word) -> Result<(), MerkleError> { - let depth = self.depth(); - let mut index = NodeIndex::new(depth, base_index_value); - if !index.is_valid() { - return Err(MerkleError::InvalidIndex(index)); - } - + let mut index = NodeIndex::new(self.depth(), base_index_value)?; let parity = index.value() & 1; let path_key = index.value() - parity; let path = match self.paths.get_mut(&path_key) { @@ -293,10 +280,9 @@ mod tests { let set = super::MerklePathSet::new(depth) .with_paths([(index, hash_6, path_6.clone().into())]) .unwrap(); - let stored_path_6 = set.get_path(NodeIndex::new(depth, index)).unwrap(); + let stored_path_6 = set.get_path(NodeIndex::make(depth, index)).unwrap(); assert_eq!(path_6, *stored_path_6); - assert!(set.get_path(NodeIndex::new(depth, 15_u64)).is_err()) } #[test] @@ -311,9 +297,8 @@ mod tests { assert_eq!( int_to_node(6u64), - set.get_node(NodeIndex::new(depth, index)).unwrap() + set.get_node(NodeIndex::make(depth, index)).unwrap() ); - assert!(set.get_node(NodeIndex::new(depth, 15_u64)).is_err()); } #[test] @@ -347,13 +332,13 @@ mod tests { let new_hash_5 = int_to_node(55); set.update_leaf(index_6, new_hash_6).unwrap(); - let new_path_4 = set.get_path(NodeIndex::new(depth, index_4)).unwrap(); + let new_path_4 = set.get_path(NodeIndex::make(depth, index_4)).unwrap(); let new_hash_67 = calculate_parent_hash(new_hash_6, 14_u64, hash_7); assert_eq!(new_hash_67, new_path_4[1]); set.update_leaf(index_5, new_hash_5).unwrap(); - let new_path_4 = set.get_path(NodeIndex::new(depth, index_4)).unwrap(); - let new_path_6 = set.get_path(NodeIndex::new(depth, index_6)).unwrap(); + let new_path_4 = set.get_path(NodeIndex::make(depth, index_4)).unwrap(); + let new_path_6 = set.get_path(NodeIndex::make(depth, index_6)).unwrap(); let new_hash_45 = calculate_parent_hash(new_hash_5, 13_u64, hash_4); assert_eq!(new_hash_45, new_path_6[1]); assert_eq!(new_hash_5, new_path_4[0]); diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs index 85fec0e..b9eda8a 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/simple_smt/mod.rs @@ -119,12 +119,12 @@ impl SimpleSmt { self.leaves.len() } - /// Returns a node at the specified key + /// Returns a node at the specified index. /// /// # Errors /// Returns an error if: /// * The specified depth is greater than the depth of the tree. - pub fn get_node(&self, index: &NodeIndex) -> Result { + pub fn get_node(&self, index: NodeIndex) -> Result { if index.is_root() { Err(MerkleError::DepthTooSmall(index.depth())) } else if index.depth() > self.depth() { @@ -137,9 +137,9 @@ impl SimpleSmt { .copied() .map(Word::from) }) - .ok_or(MerkleError::InvalidIndex(*index)) + .ok_or(MerkleError::NodeNotInSet(index.value())) } else { - let branch_node = self.get_branch_node(index); + let branch_node = self.get_branch_node(&index); Ok(Rpo256::merge(&[branch_node.left, branch_node.right]).into()) } } @@ -175,7 +175,8 @@ impl SimpleSmt { /// Returns an error if: /// * The specified key does not exist as a leaf node. pub fn get_leaf_path(&self, key: u64) -> Result { - self.get_path(NodeIndex::new(self.depth(), key)) + let index = NodeIndex::new(self.depth(), key)?; + self.get_path(index) } /// Iterator over the inner nodes of the [SimpleSmt]. @@ -190,13 +191,15 @@ impl SimpleSmt { // STATE MUTATORS // -------------------------------------------------------------------------------------------- - /// Replaces the leaf located at the specified key, and recomputes hashes by walking up the tree + /// Replaces the leaf located at the specified key, and recomputes hashes by walking up the + /// tree. /// /// # Errors /// Returns an error if the specified key is not a valid leaf index for this tree. pub fn update_leaf(&mut self, key: u64, value: Word) -> Result<(), MerkleError> { + let index = NodeIndex::new(self.depth(), key)?; if !self.check_leaf_node_exists(key) { - return Err(MerkleError::InvalidIndex(NodeIndex::new(self.depth(), key))); + return Err(MerkleError::NodeNotInSet(index.value())); } self.insert_leaf(key, value)?; @@ -208,7 +211,7 @@ impl SimpleSmt { self.insert_leaf_node(key, value); // TODO consider using a map `index |-> word` instead of `index |-> (word, word)` - let mut index = NodeIndex::new(self.depth(), key); + let mut index = NodeIndex::new(self.depth(), key)?; let mut value = RpoDigest::from(value); for _ in 0..index.depth() { let is_right = index.is_value_odd(); diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs index f7b8635..82938d3 100644 --- a/src/merkle/simple_smt/tests.rs +++ b/src/merkle/simple_smt/tests.rs @@ -61,8 +61,8 @@ fn build_sparse_tree() { let mt2 = MerkleTree::new(values.clone()).unwrap(); assert_eq!(mt2.root(), smt.root()); assert_eq!( - mt2.get_path(NodeIndex::new(3, 6)).unwrap(), - smt.get_path(NodeIndex::new(3, 6)).unwrap() + mt2.get_path(NodeIndex::make(3, 6)).unwrap(), + smt.get_path(NodeIndex::make(3, 6)).unwrap() ); // insert second value at distinct leaf branch @@ -74,8 +74,8 @@ fn build_sparse_tree() { let mt3 = MerkleTree::new(values).unwrap(); assert_eq!(mt3.root(), smt.root()); assert_eq!( - mt3.get_path(NodeIndex::new(3, 2)).unwrap(), - smt.get_path(NodeIndex::new(3, 2)).unwrap() + mt3.get_path(NodeIndex::make(3, 2)).unwrap(), + smt.get_path(NodeIndex::make(3, 2)).unwrap() ); } @@ -88,8 +88,8 @@ fn build_full_tree() { let (root, node2, node3) = compute_internal_nodes(); assert_eq!(root, tree.root()); - assert_eq!(node2, tree.get_node(&NodeIndex::new(1, 0)).unwrap()); - assert_eq!(node3, tree.get_node(&NodeIndex::new(1, 1)).unwrap()); + assert_eq!(node2, tree.get_node(NodeIndex::make(1, 0)).unwrap()); + assert_eq!(node3, tree.get_node(NodeIndex::make(1, 1)).unwrap()); } #[test] @@ -100,10 +100,10 @@ fn get_values() { .unwrap(); // check depth 2 - assert_eq!(VALUES4[0], tree.get_node(&NodeIndex::new(2, 0)).unwrap()); - assert_eq!(VALUES4[1], tree.get_node(&NodeIndex::new(2, 1)).unwrap()); - assert_eq!(VALUES4[2], tree.get_node(&NodeIndex::new(2, 2)).unwrap()); - assert_eq!(VALUES4[3], tree.get_node(&NodeIndex::new(2, 3)).unwrap()); + assert_eq!(VALUES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap()); + assert_eq!(VALUES4[1], tree.get_node(NodeIndex::make(2, 1)).unwrap()); + assert_eq!(VALUES4[2], tree.get_node(NodeIndex::make(2, 2)).unwrap()); + assert_eq!(VALUES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap()); } #[test] @@ -118,24 +118,24 @@ fn get_path() { // check depth 2 assert_eq!( vec![VALUES4[1], node3], - *tree.get_path(NodeIndex::new(2, 0)).unwrap() + *tree.get_path(NodeIndex::make(2, 0)).unwrap() ); assert_eq!( vec![VALUES4[0], node3], - *tree.get_path(NodeIndex::new(2, 1)).unwrap() + *tree.get_path(NodeIndex::make(2, 1)).unwrap() ); assert_eq!( vec![VALUES4[3], node2], - *tree.get_path(NodeIndex::new(2, 2)).unwrap() + *tree.get_path(NodeIndex::make(2, 2)).unwrap() ); assert_eq!( vec![VALUES4[2], node2], - *tree.get_path(NodeIndex::new(2, 3)).unwrap() + *tree.get_path(NodeIndex::make(2, 3)).unwrap() ); // check depth 1 - assert_eq!(vec![node3], *tree.get_path(NodeIndex::new(1, 0)).unwrap()); - assert_eq!(vec![node2], *tree.get_path(NodeIndex::new(1, 1)).unwrap()); + assert_eq!(vec![node3], *tree.get_path(NodeIndex::make(1, 0)).unwrap()); + assert_eq!(vec![node2], *tree.get_path(NodeIndex::make(1, 1)).unwrap()); } #[test] @@ -146,19 +146,19 @@ fn test_parent_node_iterator() -> Result<(), MerkleError> { .unwrap(); // check depth 2 - assert_eq!(VALUES4[0], tree.get_node(&NodeIndex::new(2, 0)).unwrap()); - assert_eq!(VALUES4[1], tree.get_node(&NodeIndex::new(2, 1)).unwrap()); - assert_eq!(VALUES4[2], tree.get_node(&NodeIndex::new(2, 2)).unwrap()); - assert_eq!(VALUES4[3], tree.get_node(&NodeIndex::new(2, 3)).unwrap()); + assert_eq!(VALUES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap()); + assert_eq!(VALUES4[1], tree.get_node(NodeIndex::make(2, 1)).unwrap()); + assert_eq!(VALUES4[2], tree.get_node(NodeIndex::make(2, 2)).unwrap()); + assert_eq!(VALUES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap()); // get parent nodes let root = tree.root(); - let l1n0 = tree.get_node(&NodeIndex::new(1, 0))?; - let l1n1 = tree.get_node(&NodeIndex::new(1, 1))?; - let l2n0 = tree.get_node(&NodeIndex::new(2, 0))?; - let l2n1 = tree.get_node(&NodeIndex::new(2, 1))?; - let l2n2 = tree.get_node(&NodeIndex::new(2, 2))?; - let l2n3 = tree.get_node(&NodeIndex::new(2, 3))?; + let l1n0 = tree.get_node(NodeIndex::make(1, 0))?; + let l1n1 = tree.get_node(NodeIndex::make(1, 1))?; + let l2n0 = tree.get_node(NodeIndex::make(2, 0))?; + let l2n1 = tree.get_node(NodeIndex::make(2, 1))?; + let l2n2 = tree.get_node(NodeIndex::make(2, 2))?; + let l2n3 = tree.get_node(NodeIndex::make(2, 3))?; let nodes: Vec = tree.inner_nodes().collect(); let expected = vec![ @@ -263,7 +263,7 @@ fn small_tree_opening_is_consistent() { ]; for (depth, key, path) in cases { - let opening = tree.get_path(NodeIndex::new(depth, key)).unwrap(); + let opening = tree.get_path(NodeIndex::make(depth, key)).unwrap(); assert_eq!(path, *opening); } @@ -287,7 +287,7 @@ proptest! { // traverse to root, fetching all paths for d in 1..depth { let k = key >> (depth - d); - tree.get_path(NodeIndex::new(d, k)).unwrap(); + tree.get_path(NodeIndex::make(d, k)).unwrap(); } } diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index 7a266af..3c52462 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -1,7 +1,7 @@ use super::mmr::{Mmr, MmrPeaks}; use super::{ - BTreeMap, BTreeSet, EmptySubtreeRoots, MerkleError, MerklePath, MerklePathSet, MerkleTree, - NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, ValuePath, Vec, Word, + BTreeMap, EmptySubtreeRoots, MerkleError, MerklePath, MerklePathSet, MerkleTree, NodeIndex, + RootPath, Rpo256, RpoDigest, SimpleSmt, ValuePath, Vec, Word, }; use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; @@ -53,15 +53,19 @@ pub struct Node { /// /// // every leaf except the last are the same /// for i in 0..7 { -/// let d0 = store.get_node(ROOT0, NodeIndex::new(3, i)).unwrap(); -/// let d1 = store.get_node(ROOT1, NodeIndex::new(3, i)).unwrap(); +/// let idx0 = NodeIndex::new(3, i).unwrap(); +/// let d0 = store.get_node(ROOT0, idx0).unwrap(); +/// let idx1 = NodeIndex::new(3, i).unwrap(); +/// let d1 = store.get_node(ROOT1, idx1).unwrap(); /// assert_eq!(d0, d1, "Both trees have the same leaf at pos {i}"); /// } /// /// // The leafs A-B-C-D are the same for both trees, so are their 2 immediate parents /// for i in 0..4 { -/// let d0 = store.get_path(ROOT0, NodeIndex::new(3, i)).unwrap(); -/// let d1 = store.get_path(ROOT1, NodeIndex::new(3, i)).unwrap(); +/// let idx0 = NodeIndex::new(3, i).unwrap(); +/// let d0 = store.get_path(ROOT0, idx0).unwrap(); +/// let idx1 = NodeIndex::new(3, i).unwrap(); +/// let d1 = store.get_path(ROOT1, idx1).unwrap(); /// assert_eq!(d0.path[0..2], d1.path[0..2], "Both sub-trees are equal up to two levels"); /// } /// @@ -184,12 +188,14 @@ impl MerkleStore { .get(&hash) .ok_or(MerkleError::RootNotInStore(hash.into()))?; - for bit in index.bit_iterator().rev() { + for i in (0..index.depth()).rev() { let node = self .nodes .get(&hash) .ok_or(MerkleError::NodeNotInStore(hash.into(), index))?; - hash = if bit { node.right } else { node.left } + + let bit = (index.value() >> i) & 1; + hash = if bit == 0 { node.left } else { node.right } } Ok(hash.into()) @@ -213,18 +219,19 @@ impl MerkleStore { .get(&hash) .ok_or(MerkleError::RootNotInStore(hash.into()))?; - for bit in index.bit_iterator().rev() { + for i in (0..index.depth()).rev() { let node = self .nodes .get(&hash) .ok_or(MerkleError::NodeNotInStore(hash.into(), index))?; - hash = if bit { - path.push(node.left.into()); - node.right - } else { + let bit = (index.value() >> i) & 1; + hash = if bit == 0 { path.push(node.right.into()); node.left + } else { + path.push(node.left.into()); + node.right } } @@ -315,7 +322,7 @@ impl MerkleStore { mut node: Word, path: MerklePath, ) -> Result { - let mut index = NodeIndex::new(self.nodes.len() as u8, index_value); + let mut index = NodeIndex::new(path.len() as u8, index_value)?; for sibling in path { let (left, right) = match index.is_value_odd() { @@ -344,33 +351,14 @@ impl MerkleStore { /// into the store. /// /// For further reference, check [MerkleStore::add_merkle_path]. - /// - /// # Errors - /// - /// Every path must resolve to the same root, otherwise this will return an `ConflictingRoots` - /// error. - pub fn add_merkle_paths(&mut self, paths: I) -> Result + pub fn add_merkle_paths(&mut self, paths: I) -> Result<(), MerkleError> where I: IntoIterator, { - let paths: Vec<(u64, Word, MerklePath)> = paths.into_iter().collect(); - - let roots: BTreeSet = paths - .iter() - .map(|(index, node, path)| path.compute_root(*index, *node).into()) - .collect(); - - if roots.len() != 1 { - return Err(MerkleError::ConflictingRoots( - roots.iter().map(|v| Word::from(*v)).collect(), - )); - } - - for (index_value, node, path) in paths { + for (index_value, node, path) in paths.into_iter() { self.add_merkle_path(index_value, node, path)?; } - - Ok(roots.iter().next().unwrap().into()) + Ok(()) } /// Appends the provided [MerklePathSet] into the store. @@ -432,15 +420,9 @@ impl MerkleStore { let root2: RpoDigest = root2.into(); if !self.nodes.contains_key(&root1) { - Err(MerkleError::NodeNotInStore( - root1.into(), - NodeIndex::new(0, 0), - )) + Err(MerkleError::NodeNotInStore(root1.into(), NodeIndex::root())) } else if !self.nodes.contains_key(&root1) { - Err(MerkleError::NodeNotInStore( - root2.into(), - NodeIndex::new(0, 0), - )) + Err(MerkleError::NodeNotInStore(root2.into(), NodeIndex::root())) } else { let parent: Word = Rpo256::merge(&[root1, root2]).into(); self.nodes.insert( diff --git a/src/merkle/store/tests.rs b/src/merkle/store/tests.rs index a14c557..e6648a9 100644 --- a/src/merkle/store/tests.rs +++ b/src/merkle/store/tests.rs @@ -22,12 +22,12 @@ fn test_root_not_in_store() -> Result<(), MerkleError> { let mtree = MerkleTree::new(LEAVES4.to_vec())?; let store = MerkleStore::default().with_merkle_tree(LEAVES4)?; assert_eq!( - store.get_node(LEAVES4[0], NodeIndex::new(mtree.depth(), 0)), + store.get_node(LEAVES4[0], NodeIndex::make(mtree.depth(), 0)), Err(MerkleError::RootNotInStore(LEAVES4[0])), "Leaf 0 is not a root" ); assert_eq!( - store.get_path(LEAVES4[0], NodeIndex::new(mtree.depth(), 0)), + store.get_path(LEAVES4[0], NodeIndex::make(mtree.depth(), 0)), Err(MerkleError::RootNotInStore(LEAVES4[0])), "Leaf 0 is not a root" ); @@ -45,22 +45,22 @@ fn test_merkle_tree() -> Result<(), MerkleError> { // STORE LEAVES ARE CORRECT ============================================================== // checks the leaves in the store corresponds to the expected values assert_eq!( - store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0)), + store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 0)), Ok(LEAVES4[0]), "node 0 must be in the tree" ); assert_eq!( - store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1)), + store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 1)), Ok(LEAVES4[1]), "node 1 must be in the tree" ); assert_eq!( - store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2)), + store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 2)), Ok(LEAVES4[2]), "node 2 must be in the tree" ); assert_eq!( - store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3)), + store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)), Ok(LEAVES4[3]), "node 3 must be in the tree" ); @@ -68,76 +68,76 @@ fn test_merkle_tree() -> Result<(), MerkleError> { // STORE LEAVES MATCH TREE =============================================================== // sanity check the values returned by the store and the tree assert_eq!( - mtree.get_node(NodeIndex::new(mtree.depth(), 0)), - store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0)), + mtree.get_node(NodeIndex::make(mtree.depth(), 0)), + store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 0)), "node 0 must be the same for both MerkleTree and MerkleStore" ); assert_eq!( - mtree.get_node(NodeIndex::new(mtree.depth(), 1)), - store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1)), + mtree.get_node(NodeIndex::make(mtree.depth(), 1)), + store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 1)), "node 1 must be the same for both MerkleTree and MerkleStore" ); assert_eq!( - mtree.get_node(NodeIndex::new(mtree.depth(), 2)), - store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2)), + mtree.get_node(NodeIndex::make(mtree.depth(), 2)), + store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 2)), "node 2 must be the same for both MerkleTree and MerkleStore" ); assert_eq!( - mtree.get_node(NodeIndex::new(mtree.depth(), 3)), - store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3)), + mtree.get_node(NodeIndex::make(mtree.depth(), 3)), + store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)), "node 3 must be the same for both MerkleTree and MerkleStore" ); // STORE MERKLE PATH MATCHS ============================================================== // assert the merkle path returned by the store is the same as the one in the tree let result = store - .get_path(mtree.root(), NodeIndex::new(mtree.depth(), 0)) + .get_path(mtree.root(), NodeIndex::make(mtree.depth(), 0)) .unwrap(); assert_eq!( LEAVES4[0], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( - mtree.get_path(NodeIndex::new(mtree.depth(), 0)), + mtree.get_path(NodeIndex::make(mtree.depth(), 0)), Ok(result.path), "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" ); let result = store - .get_path(mtree.root(), NodeIndex::new(mtree.depth(), 1)) + .get_path(mtree.root(), NodeIndex::make(mtree.depth(), 1)) .unwrap(); assert_eq!( LEAVES4[1], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( - mtree.get_path(NodeIndex::new(mtree.depth(), 1)), + mtree.get_path(NodeIndex::make(mtree.depth(), 1)), Ok(result.path), "merkle path for index 1 must be the same for the MerkleTree and MerkleStore" ); let result = store - .get_path(mtree.root(), NodeIndex::new(mtree.depth(), 2)) + .get_path(mtree.root(), NodeIndex::make(mtree.depth(), 2)) .unwrap(); assert_eq!( LEAVES4[2], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( - mtree.get_path(NodeIndex::new(mtree.depth(), 2)), + mtree.get_path(NodeIndex::make(mtree.depth(), 2)), Ok(result.path), "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" ); let result = store - .get_path(mtree.root(), NodeIndex::new(mtree.depth(), 3)) + .get_path(mtree.root(), NodeIndex::make(mtree.depth(), 3)) .unwrap(); assert_eq!( LEAVES4[3], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( - mtree.get_path(NodeIndex::new(mtree.depth(), 3)), + mtree.get_path(NodeIndex::make(mtree.depth(), 3)), Ok(result.path), "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" ); @@ -153,7 +153,7 @@ fn test_empty_roots() { for depth in 0..255 { root = Rpo256::merge(&[root; 2]); assert!( - store.get_node(root.into(), NodeIndex::new(0, 0)).is_ok(), + store.get_node(root.into(), NodeIndex::make(0, 0)).is_ok(), "The root of the empty tree of depth {depth} must be registered" ); } @@ -169,7 +169,7 @@ fn test_leaf_paths_for_empty_trees() -> Result<(), MerkleError> { for depth in 1..64 { let smt = SimpleSmt::new(depth)?; - let index = NodeIndex::new(depth, 0); + let index = NodeIndex::make(depth, 0); let store_path = store.get_path(smt.root(), index)?; let smt_path = smt.get_path(index)?; assert_eq!( @@ -181,7 +181,7 @@ fn test_leaf_paths_for_empty_trees() -> Result<(), MerkleError> { "the returned merkle path does not match the computed values" ); assert_eq!( - store_path.path.compute_root(depth.into(), EMPTY), + store_path.path.compute_root(depth.into(), EMPTY).unwrap(), smt.root(), "computed root from the path must match the empty tree root" ); @@ -197,7 +197,7 @@ fn test_get_invalid_node() { store .add_merkle_tree(LEAVES4.to_vec()) .expect("adding a merkle tree to the store must work"); - let _ = store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3)); + let _ = store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)); } #[test] @@ -211,18 +211,18 @@ fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> { .with_leaves(keys2.into_iter().zip(leaves2.into_iter())) .unwrap(); - let idx = NodeIndex::new(1, 0); - assert_eq!(smt.get_node(&idx).unwrap(), leaves2[0]); + let idx = NodeIndex::make(1, 0); + assert_eq!(smt.get_node(idx).unwrap(), leaves2[0]); assert_eq!( store.get_node(smt.root(), idx).unwrap(), - smt.get_node(&idx).unwrap() + smt.get_node(idx).unwrap() ); - let idx = NodeIndex::new(1, 1); - assert_eq!(smt.get_node(&idx).unwrap(), leaves2[1]); + let idx = NodeIndex::make(1, 1); + assert_eq!(smt.get_node(idx).unwrap(), leaves2[1]); assert_eq!( store.get_node(smt.root(), idx).unwrap(), - smt.get_node(&idx).unwrap() + smt.get_node(idx).unwrap() ); Ok(()) @@ -244,27 +244,27 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> { // STORE LEAVES ARE CORRECT ============================================================== // checks the leaves in the store corresponds to the expected values assert_eq!( - store.get_node(smt.root(), NodeIndex::new(smt.depth(), 0)), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 0)), Ok(LEAVES4[0]), "node 0 must be in the tree" ); assert_eq!( - store.get_node(smt.root(), NodeIndex::new(smt.depth(), 1)), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 1)), Ok(LEAVES4[1]), "node 1 must be in the tree" ); assert_eq!( - store.get_node(smt.root(), NodeIndex::new(smt.depth(), 2)), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 2)), Ok(LEAVES4[2]), "node 2 must be in the tree" ); assert_eq!( - store.get_node(smt.root(), NodeIndex::new(smt.depth(), 3)), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 3)), Ok(LEAVES4[3]), "node 3 must be in the tree" ); assert_eq!( - store.get_node(smt.root(), NodeIndex::new(smt.depth(), 4)), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 4)), Ok(EMPTY), "unmodified node 4 must be ZERO" ); @@ -272,94 +272,94 @@ fn test_sparse_merkle_tree() -> Result<(), MerkleError> { // STORE LEAVES MATCH TREE =============================================================== // sanity check the values returned by the store and the tree assert_eq!( - smt.get_node(&NodeIndex::new(smt.depth(), 0)), - store.get_node(smt.root(), NodeIndex::new(smt.depth(), 0)), + smt.get_node(NodeIndex::make(smt.depth(), 0)), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 0)), "node 0 must be the same for both SparseMerkleTree and MerkleStore" ); assert_eq!( - smt.get_node(&NodeIndex::new(smt.depth(), 1)), - store.get_node(smt.root(), NodeIndex::new(smt.depth(), 1)), + smt.get_node(NodeIndex::make(smt.depth(), 1)), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 1)), "node 1 must be the same for both SparseMerkleTree and MerkleStore" ); assert_eq!( - smt.get_node(&NodeIndex::new(smt.depth(), 2)), - store.get_node(smt.root(), NodeIndex::new(smt.depth(), 2)), + smt.get_node(NodeIndex::make(smt.depth(), 2)), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 2)), "node 2 must be the same for both SparseMerkleTree and MerkleStore" ); assert_eq!( - smt.get_node(&NodeIndex::new(smt.depth(), 3)), - store.get_node(smt.root(), NodeIndex::new(smt.depth(), 3)), + smt.get_node(NodeIndex::make(smt.depth(), 3)), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 3)), "node 3 must be the same for both SparseMerkleTree and MerkleStore" ); assert_eq!( - smt.get_node(&NodeIndex::new(smt.depth(), 4)), - store.get_node(smt.root(), NodeIndex::new(smt.depth(), 4)), + smt.get_node(NodeIndex::make(smt.depth(), 4)), + store.get_node(smt.root(), NodeIndex::make(smt.depth(), 4)), "node 4 must be the same for both SparseMerkleTree and MerkleStore" ); // STORE MERKLE PATH MATCHS ============================================================== // assert the merkle path returned by the store is the same as the one in the tree let result = store - .get_path(smt.root(), NodeIndex::new(smt.depth(), 0)) + .get_path(smt.root(), NodeIndex::make(smt.depth(), 0)) .unwrap(); assert_eq!( LEAVES4[0], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( - smt.get_path(NodeIndex::new(smt.depth(), 0)), + smt.get_path(NodeIndex::make(smt.depth(), 0)), Ok(result.path), "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" ); let result = store - .get_path(smt.root(), NodeIndex::new(smt.depth(), 1)) + .get_path(smt.root(), NodeIndex::make(smt.depth(), 1)) .unwrap(); assert_eq!( LEAVES4[1], result.value, "Value for merkle path at index 1 must match leaf value" ); assert_eq!( - smt.get_path(NodeIndex::new(smt.depth(), 1)), + smt.get_path(NodeIndex::make(smt.depth(), 1)), Ok(result.path), "merkle path for index 1 must be the same for the MerkleTree and MerkleStore" ); let result = store - .get_path(smt.root(), NodeIndex::new(smt.depth(), 2)) + .get_path(smt.root(), NodeIndex::make(smt.depth(), 2)) .unwrap(); assert_eq!( LEAVES4[2], result.value, "Value for merkle path at index 2 must match leaf value" ); assert_eq!( - smt.get_path(NodeIndex::new(smt.depth(), 2)), + smt.get_path(NodeIndex::make(smt.depth(), 2)), Ok(result.path), "merkle path for index 2 must be the same for the MerkleTree and MerkleStore" ); let result = store - .get_path(smt.root(), NodeIndex::new(smt.depth(), 3)) + .get_path(smt.root(), NodeIndex::make(smt.depth(), 3)) .unwrap(); assert_eq!( LEAVES4[3], result.value, "Value for merkle path at index 3 must match leaf value" ); assert_eq!( - smt.get_path(NodeIndex::new(smt.depth(), 3)), + smt.get_path(NodeIndex::make(smt.depth(), 3)), Ok(result.path), "merkle path for index 3 must be the same for the MerkleTree and MerkleStore" ); let result = store - .get_path(smt.root(), NodeIndex::new(smt.depth(), 4)) + .get_path(smt.root(), NodeIndex::make(smt.depth(), 4)) .unwrap(); assert_eq!( EMPTY, result.value, "Value for merkle path at index 4 must match leaf value" ); assert_eq!( - smt.get_path(NodeIndex::new(smt.depth(), 4)), + smt.get_path(NodeIndex::make(smt.depth(), 4)), Ok(result.path), "merkle path for index 4 must be the same for the MerkleTree and MerkleStore" ); @@ -372,16 +372,16 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { let mtree = MerkleTree::new(LEAVES4.to_vec())?; let i0 = 0; - let p0 = mtree.get_path(NodeIndex::new(2, i0)).unwrap(); + let p0 = mtree.get_path(NodeIndex::make(2, i0)).unwrap(); let i1 = 1; - let p1 = mtree.get_path(NodeIndex::new(2, i1)).unwrap(); + let p1 = mtree.get_path(NodeIndex::make(2, i1)).unwrap(); let i2 = 2; - let p2 = mtree.get_path(NodeIndex::new(2, i2)).unwrap(); + let p2 = mtree.get_path(NodeIndex::make(2, i2)).unwrap(); let i3 = 3; - let p3 = mtree.get_path(NodeIndex::new(2, i3)).unwrap(); + let p3 = mtree.get_path(NodeIndex::make(2, i3)).unwrap(); let paths = [ (i0, LEAVES4[i0 as usize], p0), @@ -401,22 +401,22 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { // STORE LEAVES ARE CORRECT ============================================================== // checks the leaves in the store corresponds to the expected values assert_eq!( - store.get_node(set.root(), NodeIndex::new(set.depth(), 0)), + store.get_node(set.root(), NodeIndex::make(set.depth(), 0)), Ok(LEAVES4[0]), "node 0 must be in the set" ); assert_eq!( - store.get_node(set.root(), NodeIndex::new(set.depth(), 1)), + store.get_node(set.root(), NodeIndex::make(set.depth(), 1)), Ok(LEAVES4[1]), "node 1 must be in the set" ); assert_eq!( - store.get_node(set.root(), NodeIndex::new(set.depth(), 2)), + store.get_node(set.root(), NodeIndex::make(set.depth(), 2)), Ok(LEAVES4[2]), "node 2 must be in the set" ); assert_eq!( - store.get_node(set.root(), NodeIndex::new(set.depth(), 3)), + store.get_node(set.root(), NodeIndex::make(set.depth(), 3)), Ok(LEAVES4[3]), "node 3 must be in the set" ); @@ -424,76 +424,76 @@ fn test_add_merkle_paths() -> Result<(), MerkleError> { // STORE LEAVES MATCH SET ================================================================ // sanity check the values returned by the store and the set assert_eq!( - set.get_node(NodeIndex::new(set.depth(), 0)), - store.get_node(set.root(), NodeIndex::new(set.depth(), 0)), + set.get_node(NodeIndex::make(set.depth(), 0)), + store.get_node(set.root(), NodeIndex::make(set.depth(), 0)), "node 0 must be the same for both SparseMerkleTree and MerkleStore" ); assert_eq!( - set.get_node(NodeIndex::new(set.depth(), 1)), - store.get_node(set.root(), NodeIndex::new(set.depth(), 1)), + set.get_node(NodeIndex::make(set.depth(), 1)), + store.get_node(set.root(), NodeIndex::make(set.depth(), 1)), "node 1 must be the same for both SparseMerkleTree and MerkleStore" ); assert_eq!( - set.get_node(NodeIndex::new(set.depth(), 2)), - store.get_node(set.root(), NodeIndex::new(set.depth(), 2)), + set.get_node(NodeIndex::make(set.depth(), 2)), + store.get_node(set.root(), NodeIndex::make(set.depth(), 2)), "node 2 must be the same for both SparseMerkleTree and MerkleStore" ); assert_eq!( - set.get_node(NodeIndex::new(set.depth(), 3)), - store.get_node(set.root(), NodeIndex::new(set.depth(), 3)), + set.get_node(NodeIndex::make(set.depth(), 3)), + store.get_node(set.root(), NodeIndex::make(set.depth(), 3)), "node 3 must be the same for both SparseMerkleTree and MerkleStore" ); // STORE MERKLE PATH MATCHS ============================================================== // assert the merkle path returned by the store is the same as the one in the set let result = store - .get_path(set.root(), NodeIndex::new(set.depth(), 0)) + .get_path(set.root(), NodeIndex::make(set.depth(), 0)) .unwrap(); assert_eq!( LEAVES4[0], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( - set.get_path(NodeIndex::new(set.depth(), 0)), + set.get_path(NodeIndex::make(set.depth(), 0)), Ok(result.path), "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" ); let result = store - .get_path(set.root(), NodeIndex::new(set.depth(), 1)) + .get_path(set.root(), NodeIndex::make(set.depth(), 1)) .unwrap(); assert_eq!( LEAVES4[1], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( - set.get_path(NodeIndex::new(set.depth(), 1)), + set.get_path(NodeIndex::make(set.depth(), 1)), Ok(result.path), "merkle path for index 1 must be the same for the MerkleTree and MerkleStore" ); let result = store - .get_path(set.root(), NodeIndex::new(set.depth(), 2)) + .get_path(set.root(), NodeIndex::make(set.depth(), 2)) .unwrap(); assert_eq!( LEAVES4[2], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( - set.get_path(NodeIndex::new(set.depth(), 2)), + set.get_path(NodeIndex::make(set.depth(), 2)), Ok(result.path), "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" ); let result = store - .get_path(set.root(), NodeIndex::new(set.depth(), 3)) + .get_path(set.root(), NodeIndex::make(set.depth(), 3)) .unwrap(); assert_eq!( LEAVES4[3], result.value, "Value for merkle path at index 0 must match leaf value" ); assert_eq!( - set.get_path(NodeIndex::new(set.depth(), 3)), + set.get_path(NodeIndex::make(set.depth(), 3)), Ok(result.path), "merkle path for index 0 must be the same for the MerkleTree and MerkleStore" ); @@ -550,7 +550,7 @@ fn store_path_opens_from_leaf() { .with_merkle_tree([a, b, c, d, e, f, g, h]) .unwrap(); let path = store - .get_path(root.into(), NodeIndex::new(3, 1)) + .get_path(root.into(), NodeIndex::make(3, 1)) .unwrap() .path; @@ -563,7 +563,7 @@ fn test_set_node() -> Result<(), MerkleError> { let mtree = MerkleTree::new(LEAVES4.to_vec())?; let mut store = MerkleStore::default().with_merkle_tree(LEAVES4)?; let value = int_to_node(42); - let index = NodeIndex::new(mtree.depth(), 0); + let index = NodeIndex::make(mtree.depth(), 0); let new_root = store.set_node(mtree.root(), index, value)?.root; assert_eq!( store.get_node(new_root, index), @@ -582,7 +582,7 @@ fn test_constructors() -> Result<(), MerkleError> { let depth = mtree.depth(); let leaves = 2u64.pow(depth.into()); for index in 0..leaves { - let index = NodeIndex::new(depth, index); + let index = NodeIndex::make(depth, index); let value_path = store.get_path(mtree.root(), index)?; assert_eq!(mtree.get_path(index)?, value_path.path); } @@ -597,34 +597,50 @@ fn test_constructors() -> Result<(), MerkleError> { let depth = smt.depth(); for key in KEYS4 { - let index = NodeIndex::new(depth, key); + let index = NodeIndex::make(depth, key); let value_path = store.get_path(smt.root(), index)?; assert_eq!(smt.get_path(index)?, value_path.path); } let d = 2; let paths = [ - (0, LEAVES4[0], mtree.get_path(NodeIndex::new(d, 0)).unwrap()), - (1, LEAVES4[1], mtree.get_path(NodeIndex::new(d, 1)).unwrap()), - (2, LEAVES4[2], mtree.get_path(NodeIndex::new(d, 2)).unwrap()), - (3, LEAVES4[3], mtree.get_path(NodeIndex::new(d, 3)).unwrap()), + ( + 0, + LEAVES4[0], + mtree.get_path(NodeIndex::make(d, 0)).unwrap(), + ), + ( + 1, + LEAVES4[1], + mtree.get_path(NodeIndex::make(d, 1)).unwrap(), + ), + ( + 2, + LEAVES4[2], + mtree.get_path(NodeIndex::make(d, 2)).unwrap(), + ), + ( + 3, + LEAVES4[3], + mtree.get_path(NodeIndex::make(d, 3)).unwrap(), + ), ]; let store1 = MerkleStore::default().with_merkle_paths(paths.clone())?; let store2 = MerkleStore::default() - .with_merkle_path(0, LEAVES4[0], mtree.get_path(NodeIndex::new(d, 0))?)? - .with_merkle_path(1, LEAVES4[1], mtree.get_path(NodeIndex::new(d, 1))?)? - .with_merkle_path(2, LEAVES4[2], mtree.get_path(NodeIndex::new(d, 2))?)? - .with_merkle_path(3, LEAVES4[3], mtree.get_path(NodeIndex::new(d, 3))?)?; + .with_merkle_path(0, LEAVES4[0], mtree.get_path(NodeIndex::make(d, 0))?)? + .with_merkle_path(1, LEAVES4[1], mtree.get_path(NodeIndex::make(d, 1))?)? + .with_merkle_path(2, LEAVES4[2], mtree.get_path(NodeIndex::make(d, 2))?)? + .with_merkle_path(3, LEAVES4[3], mtree.get_path(NodeIndex::make(d, 3))?)?; let set = MerklePathSet::new(d).with_paths(paths).unwrap(); for key in [0, 1, 2, 3] { - let index = NodeIndex::new(d, key); + let index = NodeIndex::make(d, key); let value_path1 = store1.get_path(set.root(), index)?; let value_path2 = store2.get_path(set.root(), index)?; assert_eq!(value_path1, value_path2); - let index = NodeIndex::new(d, key); + let index = NodeIndex::make(d, key); assert_eq!(set.get_path(index)?, value_path1.path); } From ab903a2229436ea2498c08359c14d753657dd7d4 Mon Sep 17 00:00:00 2001 From: "Augusto F. Hack" Date: Thu, 6 Apr 2023 15:41:24 +0200 Subject: [PATCH 08/10] mmr: optimized peak hash for Miden VM --- src/merkle/mmr/accumulator.rs | 39 +++++++++++++++------ src/merkle/mmr/full.rs | 4 +-- src/merkle/mmr/tests.rs | 64 +++++++++++++++++++++++++++++++++-- 3 files changed, 92 insertions(+), 15 deletions(-) diff --git a/src/merkle/mmr/accumulator.rs b/src/merkle/mmr/accumulator.rs index f27355f..2a5d522 100644 --- a/src/merkle/mmr/accumulator.rs +++ b/src/merkle/mmr/accumulator.rs @@ -1,4 +1,8 @@ -use super::{super::Vec, MmrProof, Rpo256, Word}; +use super::{ + super::Vec, + super::{WORD_SIZE, ZERO}, + MmrProof, Rpo256, Word, +}; #[derive(Debug, Clone, PartialEq)] pub struct MmrPeaks { @@ -8,18 +12,17 @@ pub struct MmrPeaks { /// the MMR has a power-of-two number of leaves there is a single peak. /// /// Every tree in the MMR forest has a distinct power-of-two size, this means only the right - /// most tree can have an odd number of elements (1). Additionally this means that the bits in + /// most tree can have an odd number of elements (e.g. `1`). Additionally this means that the bits in /// `num_leaves` conveniently encode the size of each individual tree. /// /// Examples: /// - /// Example 1: With 5 leaves, the binary 0b101. The number of set bits is equal the number - /// of peaks, in this case there are 2 peaks. The 0-indexed least-significant position of - /// the bit determines the number of elements of a tree, so the rightmost tree has 2**0 - /// elements and the left most has 2**2. - /// - /// Example 2: With 12 leaves, the binary is 0b1100, this case also has 2 peaks, the - /// leftmost tree has 2**3=8 elements, and the right most has 2**2=4 elements. + /// - With 5 leaves, the binary `0b101`. The number of set bits is equal the number + /// of peaks, in this case there are 2 peaks. The 0-indexed least-significant position of + /// the bit determines the number of elements of a tree, so the rightmost tree has `2**0` + /// elements and the left most has `2**2`. + /// - With 12 leaves, the binary is `0b1100`, this case also has 2 peaks, the + /// leftmost tree has `2**3=8` elements, and the right most has `2**2=4` elements. pub num_leaves: usize, /// All the peaks of every tree in the MMR forest. The peaks are always ordered by number of @@ -30,9 +33,23 @@ pub struct MmrPeaks { } impl MmrPeaks { - /// Hashes the peaks sequentially, compacting it to a single digest + /// Hashes the peaks. + /// + /// The hashing is optimized to work with the Miden VM, the procedure will: + /// + /// - Pad the peaks with ZERO to an even number of words, this removes the need to handle RPO padding. + /// - Pad the peaks to a minimum length of 16 words, which reduces the constant cost of + /// hashing. pub fn hash_peaks(&self) -> Word { - Rpo256::hash_elements(&self.peaks.as_slice().concat()).into() + let mut copy = self.peaks.clone(); + + if copy.len() < 16 { + copy.resize(16, [ZERO; WORD_SIZE]) + } else if copy.len() % 2 == 1 { + copy.push([ZERO; WORD_SIZE]) + } + + Rpo256::hash_elements(©.as_slice().concat()).into() } pub fn verify(&self, value: Word, opening: MmrProof) -> bool { diff --git a/src/merkle/mmr/full.rs b/src/merkle/mmr/full.rs index 76f1e3d..d100f57 100644 --- a/src/merkle/mmr/full.rs +++ b/src/merkle/mmr/full.rs @@ -174,7 +174,7 @@ impl Mmr { self.forest += 1; } - /// Returns an accumulator representing the current state of the MMMR. + /// Returns an accumulator representing the current state of the MMR. pub fn accumulator(&self) -> MmrPeaks { let peaks: Vec = TrueBitPositionIterator::new(self.forest) .rev() @@ -192,7 +192,7 @@ impl Mmr { } } - /// An iterator over inner nodes in the [Mmm]. The order of iteration is unspecified. + /// An iterator over inner nodes in the MMR. The order of iteration is unspecified. pub fn inner_nodes(&self) -> MmrNodes { MmrNodes { mmr: self, diff --git a/src/merkle/mmr/tests.rs b/src/merkle/mmr/tests.rs index 1fb879e..218a6dd 100644 --- a/src/merkle/mmr/tests.rs +++ b/src/merkle/mmr/tests.rs @@ -1,8 +1,8 @@ use super::bit::TrueBitPositionIterator; use super::full::{high_bitmask, leaf_to_corresponding_tree, nodes_in_forest}; use super::{ - super::{InnerNodeInfo, Vec}, - Mmr, Rpo256, Word, + super::{InnerNodeInfo, Vec, WORD_SIZE, ZERO}, + Mmr, MmrPeaks, Rpo256, Word, }; use crate::merkle::{int_to_node, MerklePath}; @@ -448,6 +448,66 @@ fn test_mmr_inner_nodes() { assert_eq!(postorder, nodes); } +#[test] +fn test_mmr_hash_peaks() { + let mmr: Mmr = LEAVES.into(); + let peaks = mmr.accumulator(); + + let first_peak = *Rpo256::merge(&[ + Rpo256::hash_elements(&[LEAVES[0], LEAVES[1]].concat()), + Rpo256::hash_elements(&[LEAVES[2], LEAVES[3]].concat()), + ]); + let second_peak = *Rpo256::hash_elements(&[LEAVES[4], LEAVES[5]].concat()); + let third_peak = LEAVES[6]; + + // minimum length is 16 + let mut expected_peaks = [first_peak, second_peak, third_peak].to_vec(); + expected_peaks.resize(16, [ZERO; WORD_SIZE]); + assert_eq!( + peaks.hash_peaks(), + *Rpo256::hash_elements(&expected_peaks.as_slice().concat()) + ); +} + +#[test] +fn test_mmr_peaks_hash_less_than_16() { + let mut peaks = Vec::new(); + + for i in 0..16 { + peaks.push(int_to_node(i)); + let accumulator = MmrPeaks { + num_leaves: (1 << peaks.len()) - 1, + peaks: peaks.clone(), + }; + + // minimum length is 16 + let mut expected_peaks = peaks.clone(); + expected_peaks.resize(16, [ZERO; WORD_SIZE]); + assert_eq!( + accumulator.hash_peaks(), + *Rpo256::hash_elements(&expected_peaks.as_slice().concat()) + ); + } +} + +#[test] +fn test_mmr_peaks_hash_odd() { + let peaks: Vec<_> = (0..=17).map(|i| int_to_node(i)).collect(); + + let accumulator = MmrPeaks { + num_leaves: (1 << peaks.len()) - 1, + peaks: peaks.clone(), + }; + + // odd length bigger than 16 is padded to the next even nubmer + let mut expected_peaks = peaks.clone(); + expected_peaks.resize(18, [ZERO; WORD_SIZE]); + assert_eq!( + accumulator.hash_peaks(), + *Rpo256::hash_elements(&expected_peaks.as_slice().concat()) + ); +} + mod property_tests { use super::leaf_to_corresponding_tree; use proptest::prelude::*; From 00ffc1568aefcc8bb0cf081e707a7186f87e110c Mon Sep 17 00:00:00 2001 From: Victor Lopez Date: Wed, 5 Apr 2023 16:02:39 +0200 Subject: [PATCH 09/10] feat: add `MerkleStore::get_leaf_depth` This commit introduces `get_leaf_depth`, a tiered SMT helpers that will retrieve the depth of a leaf for a given root, capped by `64`. closes #119 --- src/merkle/store/mod.rs | 75 +++++++++++++++++++ src/merkle/store/tests.rs | 151 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 226 insertions(+) diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index 3c52462..a938c4f 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -244,6 +244,81 @@ impl MerkleStore { }) } + /// Reconstructs a path from the root until a leaf or empty node and returns its depth. + /// + /// The `tree_depth` parameter defines up to which depth the tree will be traversed, starting + /// from `root`. The maximum value the argument accepts is [u64::BITS]. + /// + /// The traversed path from leaf to root will start at the least significant bit of `index`, + /// and will be executed for `tree_depth` bits. + /// + /// # Errors + /// Will return an error if: + /// - The provided root is not found. + /// - The path from the root continues to a depth greater than `tree_depth`. + /// - The provided `tree_depth` is greater than `64. + /// - The provided `index` is not valid for a depth equivalent to `tree_depth`. For more + /// information, check [NodeIndex::new]. + pub fn get_leaf_depth( + &self, + root: Word, + tree_depth: u8, + index: u64, + ) -> Result { + // validate depth and index + if tree_depth > 64 { + return Err(MerkleError::DepthTooBig(tree_depth as u64)); + } + NodeIndex::new(tree_depth, index)?; + + // it's not illegal to have a maximum depth of `0`; we should just return the root in that + // case. this check will simplify the implementation as we could overflow bits for depth + // `0`. + if tree_depth == 0 { + return Ok(0); + } + + // check if the root exists, providing the proper error report if it doesn't + let empty = EmptySubtreeRoots::empty_hashes(tree_depth); + let mut hash: RpoDigest = root.into(); + if !self.nodes.contains_key(&hash) { + return Err(MerkleError::RootNotInStore(hash.into())); + } + + // we traverse from root to leaf, so the path is reversed + let mut path = (index << (64 - tree_depth)).reverse_bits(); + + // iterate every depth and reconstruct the path from root to leaf + for depth in 0..tree_depth { + // we short-circuit if an empty node has been found + if hash == empty[depth as usize] { + return Ok(depth); + } + + // fetch the children pair, mapped by its parent hash + let children = match self.nodes.get(&hash) { + Some(node) => node, + None => return Ok(depth), + }; + + // traverse down + hash = if path & 1 == 0 { + children.left + } else { + children.right + }; + path >>= 1; + } + + // at max depth assert it doesn't have sub-trees + if self.nodes.contains_key(&hash) { + return Err(MerkleError::DepthTooBig(tree_depth as u64 + 1)); + } + + // depleted bits; return max depth + Ok(tree_depth) + } + // STATE MUTATORS // -------------------------------------------------------------------------------------------- diff --git a/src/merkle/store/tests.rs b/src/merkle/store/tests.rs index e6648a9..de65b1f 100644 --- a/src/merkle/store/tests.rs +++ b/src/merkle/store/tests.rs @@ -647,6 +647,157 @@ fn test_constructors() -> Result<(), MerkleError> { Ok(()) } +#[test] +fn node_path_should_be_truncated_by_midtier_insert() { + let key = 0b11010010_11001100_11001100_11001100_11001100_11001100_11001100_11001100_u64; + + let mut store = MerkleStore::new(); + let root: Word = EmptySubtreeRoots::empty_hashes(64)[0].into(); + + // insert first node - works as expected + let depth = 64; + let node = [Felt::new(key); WORD_SIZE]; + let index = NodeIndex::new(depth, key).unwrap(); + let root = store.set_node(root, index, node).unwrap().root; + let result = store.get_node(root, index).unwrap(); + let path = store.get_path(root, index).unwrap().path; + assert_eq!(node, result); + assert_eq!(path.depth(), depth); + assert!(path.verify(index.value(), result, &root)); + + // flip the first bit of the key and insert the second node on a different depth + let key = key ^ (1 << 63); + let key = key >> 8; + let depth = 56; + let node = [Felt::new(key); WORD_SIZE]; + let index = NodeIndex::new(depth, key).unwrap(); + let root = store.set_node(root, index, node).unwrap().root; + let result = store.get_node(root, index).unwrap(); + let path = store.get_path(root, index).unwrap().path; + assert_eq!(node, result); + assert_eq!(path.depth(), depth); + assert!(path.verify(index.value(), result, &root)); + + // attempt to fetch a path of the second node to depth 64 + // should fail because the previously inserted node will remove its sub-tree from the set + let key = key << 8; + let index = NodeIndex::new(64, key).unwrap(); + assert!(store.get_node(root, index).is_err()); +} + +#[test] +fn get_leaf_depth_works_depth_64() { + let mut store = MerkleStore::new(); + let mut root: Word = EmptySubtreeRoots::empty_hashes(64)[0].into(); + let key = u64::MAX; + + // this will create a rainbow tree and test all opening to depth 64 + for d in 0..64 { + let k = key & (u64::MAX >> d); + let node = [Felt::new(k); WORD_SIZE]; + let index = NodeIndex::new(64, k).unwrap(); + + // assert the leaf doesn't exist before the insert. the returned depth should always + // increment with the paths count of the set, as they are insersecting one another up to + // the first bits of the used key. + assert_eq!(d, store.get_leaf_depth(root, 64, k).unwrap()); + + // insert and assert the correct depth + root = store.set_node(root, index, node).unwrap().root; + assert_eq!(64, store.get_leaf_depth(root, 64, k).unwrap()); + } +} + +#[test] +fn get_leaf_depth_works_with_incremental_depth() { + let mut store = MerkleStore::new(); + let mut root: Word = EmptySubtreeRoots::empty_hashes(64)[0].into(); + + // insert some path to the left of the root and assert it + let key = 0b01001011_10110110_00001101_01110100_00111011_10101101_00000100_01000001_u64; + assert_eq!(0, store.get_leaf_depth(root, 64, key).unwrap()); + let depth = 64; + let index = NodeIndex::new(depth, key).unwrap(); + let node = [Felt::new(key); WORD_SIZE]; + root = store.set_node(root, index, node).unwrap().root; + assert_eq!(depth, store.get_leaf_depth(root, 64, key).unwrap()); + + // flip the key to the right of the root and insert some content on depth 16 + let key = 0b11001011_10110110_00000000_00000000_00000000_00000000_00000000_00000000_u64; + assert_eq!(1, store.get_leaf_depth(root, 64, key).unwrap()); + let depth = 16; + let index = NodeIndex::new(depth, key >> (64 - depth)).unwrap(); + let node = [Felt::new(key); WORD_SIZE]; + root = store.set_node(root, index, node).unwrap().root; + assert_eq!(depth, store.get_leaf_depth(root, 64, key).unwrap()); + + // attempt the sibling of the previous leaf + let key = 0b11001011_10110111_00000000_00000000_00000000_00000000_00000000_00000000_u64; + assert_eq!(16, store.get_leaf_depth(root, 64, key).unwrap()); + let index = NodeIndex::new(depth, key >> (64 - depth)).unwrap(); + let node = [Felt::new(key); WORD_SIZE]; + root = store.set_node(root, index, node).unwrap().root; + assert_eq!(depth, store.get_leaf_depth(root, 64, key).unwrap()); + + // move down to the next depth and assert correct behavior + let key = 0b11001011_10110100_00000000_00000000_00000000_00000000_00000000_00000000_u64; + assert_eq!(15, store.get_leaf_depth(root, 64, key).unwrap()); + let depth = 17; + let index = NodeIndex::new(depth, key >> (64 - depth)).unwrap(); + let node = [Felt::new(key); WORD_SIZE]; + root = store.set_node(root, index, node).unwrap().root; + assert_eq!(depth, store.get_leaf_depth(root, 64, key).unwrap()); +} + +#[test] +fn get_leaf_depth_works_with_depth_8() { + let mut store = MerkleStore::new(); + let mut root: Word = EmptySubtreeRoots::empty_hashes(8)[0].into(); + + // insert some random, 8 depth keys. `a` diverges from the first bit + let a = 0b01101001_u64; + let b = 0b10011001_u64; + let c = 0b10010110_u64; + let d = 0b11110110_u64; + + for k in [a, b, c, d] { + let index = NodeIndex::new(8, k).unwrap(); + let node = [Felt::new(k); WORD_SIZE]; + root = store.set_node(root, index, node).unwrap().root; + } + + // assert all leaves returns the inserted depth + for k in [a, b, c, d] { + assert_eq!(8, store.get_leaf_depth(root, 8, k).unwrap()); + } + + // flip last bit of a and expect it to return the the same depth, but for an empty node + assert_eq!(8, store.get_leaf_depth(root, 8, 0b01101000_u64).unwrap()); + + // flip fourth bit of a and expect an empty node on depth 4 + assert_eq!(4, store.get_leaf_depth(root, 8, 0b01111001_u64).unwrap()); + + // flip third bit of a and expect an empty node on depth 3 + assert_eq!(3, store.get_leaf_depth(root, 8, 0b01001001_u64).unwrap()); + + // flip second bit of a and expect an empty node on depth 2 + assert_eq!(2, store.get_leaf_depth(root, 8, 0b00101001_u64).unwrap()); + + // flip fourth bit of c and expect an empty node on depth 4 + assert_eq!(4, store.get_leaf_depth(root, 8, 0b10000110_u64).unwrap()); + + // flip second bit of d and expect an empty node on depth 3 as depth 2 conflicts with b and c + assert_eq!(3, store.get_leaf_depth(root, 8, 0b10110110_u64).unwrap()); + + // duplicate the tree on `a` and assert the depth is short-circuited by such sub-tree + let index = NodeIndex::new(8, a).unwrap(); + root = store.set_node(root, index, root).unwrap().root; + assert_eq!( + Err(MerkleError::DepthTooBig(9)), + store.get_leaf_depth(root, 8, a) + ); +} + #[cfg(std)] #[test] fn test_serialization() -> Result<(), Box> { From 61db888b2cab4898937e61696a5f70ec3f29b14f Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Fri, 7 Apr 2023 23:44:27 -0700 Subject: [PATCH 10/10] chore: update crate version to v0.3 --- CHANGELOG.md | 7 +++++++ Cargo.toml | 4 ++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index aa67da3..461680b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,10 @@ +## 0.3.0 (2023-04-08) + +- Added `depth` parameter to SMT constructors in `MerkleStore` (#115). +- Optimized MMR peak hashing for Miden VM (#120). +- Added `get_leaf_depth` method to `MerkleStore` (#119). +- Added inner node iterators to `MerkleTree`, `SimpleSmt`, and `Mmr` (#117, #118, #121). + ## 0.2.0 (2023-03-24) - Implemented `Mmr` and related structs (#67). diff --git a/Cargo.toml b/Cargo.toml index 069e286..13fc97f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,12 +1,12 @@ [package] name = "miden-crypto" -version = "0.2.0" +version = "0.3.0" description = "Miden Cryptographic primitives" authors = ["miden contributors"] readme = "README.md" license = "MIT" repository = "https://github.com/0xPolygonMiden/crypto" -documentation = "https://docs.rs/miden-crypto/0.2.0" +documentation = "https://docs.rs/miden-crypto/0.3.0" categories = ["cryptography", "no-std"] keywords = ["miden", "crypto", "hash", "merkle"] edition = "2021"