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" 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/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/benches/store.rs b/benches/store.rs index 793118a..804faa6 100644 --- a/benches/store.rs +++ b/benches/store.rs @@ -18,17 +18,18 @@ 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. fn get_empty_leaf_simplesmt(c: &mut Criterion) { let mut group = c.benchmark_group("get_empty_leaf_simplesmt"); - let depth = 63u8; - let size = 2u64.pow(depth as u32); + let depth = SimpleSmt::MAX_DEPTH; + 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, ) }); @@ -103,12 +104,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(); @@ -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, ) }); @@ -136,8 +137,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 size = 2u64.pow(depth as u32); + let depth = SimpleSmt::MAX_DEPTH; // 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, ) }); @@ -216,29 +217,29 @@ 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; 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, ) }); @@ -295,12 +296,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(); @@ -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, ) }); @@ -366,7 +367,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 +383,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, ) }, @@ -409,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, ) @@ -418,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, @@ -450,12 +452,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(); @@ -463,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, ) @@ -472,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/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) +} 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 434c8e7..3587fc6 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 { @@ -77,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; @@ -97,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 @@ -124,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 @@ -158,9 +152,50 @@ 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, } -/// Utility to vizualize a [MerkleTree] in text. +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 visualize a [MerkleTree] in text. pub fn tree_to_text(tree: &MerkleTree) -> Result { let indent = " "; let mut s = String::new(); @@ -169,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); @@ -186,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('['); @@ -212,7 +244,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::*; @@ -258,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] @@ -279,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] @@ -323,6 +355,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::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![ + 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/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 92c350c..d100f57 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")] @@ -172,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() @@ -190,6 +192,16 @@ impl Mmr { } } + /// An iterator over inner nodes in the MMR. 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..218a6dd 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, WORD_SIZE, ZERO}, + Mmr, MmrPeaks, Rpo256, Word, +}; use crate::merkle::{int_to_node, MerklePath}; #[test] @@ -410,6 +413,101 @@ 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); +} + +#[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::*; diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index b0c7b2d..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; @@ -32,6 +32,9 @@ pub use mmr::{Mmr, MmrPeaks}; mod store; pub use store::MerkleStore; +mod node; +pub use node::InnerNodeInfo; + // ERRORS // ================================================================================================ @@ -42,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), @@ -60,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/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/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 e330a97..b9eda8a 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)] @@ -8,14 +9,27 @@ mod tests; // SPARSE MERKLE TREE // ================================================================================================ -/// A sparse Merkle tree with 63-bit keys and 4-element leaf values, without compaction. -/// Manipulation and retrieval of leaves and internal nodes is provided by its internal `Store`. +/// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction. /// 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, + branches: BTreeMap, + empty_hashes: Vec, +} + +#[derive(Debug, Default, Clone, PartialEq, Eq)] +struct BranchNode { + left: RpoDigest, + right: RpoDigest, +} + +impl BranchNode { + fn parent(&self) -> RpoDigest { + Rpo256::merge(&[self.left, self.right]) + } } impl SimpleSmt { @@ -26,7 +40,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 // -------------------------------------------------------------------------------------------- @@ -40,8 +54,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. @@ -57,7 +79,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())); } @@ -72,8 +94,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,32 +116,30 @@ 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 + /// 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() { 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)) + .ok_or(MerkleError::NodeNotInSet(index.value())) } 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 +161,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); } @@ -156,19 +175,31 @@ 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]. + 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 // -------------------------------------------------------------------------------------------- - /// 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> { - if !self.store.check_leaf_node_exists(key) { - return Err(MerkleError::InvalidIndex(NodeIndex::new(self.depth(), key))); + let index = NodeIndex::new(self.depth(), key)?; + if !self.check_leaf_node_exists(key) { + return Err(MerkleError::NodeNotInSet(index.value())); } self.insert_leaf(key, value)?; @@ -177,67 +208,29 @@ 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); + 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(); 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 +262,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/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs index 6abd343..82938d3 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::*; @@ -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,69 @@ 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] +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::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::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![ + 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] @@ -218,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); } @@ -242,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 f98ae6f..a938c4f 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -1,6 +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}; @@ -52,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"); /// } /// @@ -115,13 +120,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) } @@ -145,6 +156,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 // -------------------------------------------------------------------------------------------- @@ -168,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()) @@ -197,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 } } @@ -221,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 // -------------------------------------------------------------------------------------------- @@ -243,56 +341,44 @@ 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 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)?; - for branch in smt.store.branches.values() { - let parent = Rpo256::merge(&[branch.left, branch.right]); + let smt = SimpleSmt::new(depth)?.with_leaves(entries)?; + 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(), }, ); } @@ -311,7 +397,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() { @@ -340,33 +426,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. @@ -380,6 +447,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 @@ -409,15 +495,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 9ccd3ec..de65b1f 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] @@ -205,24 +205,24 @@ 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())) .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(()) @@ -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() @@ -241,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" ); @@ -269,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" ); @@ -369,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), @@ -398,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" ); @@ -421,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" ); @@ -547,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; @@ -560,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), @@ -579,54 +582,222 @@ 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); } + 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(); 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); } 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> {