diff --git a/src/merkle/index.rs b/src/merkle/index.rs index 5a559a1..564d981 100644 --- a/src/merkle/index.rs +++ b/src/merkle/index.rs @@ -74,12 +74,26 @@ impl NodeIndex { Self { depth: 0, value: 0 } } - /// Computes the value of the sibling of the current node. - pub fn sibling(mut self) -> Self { + /// Computes sibling index of the current node. + pub const fn sibling(mut self) -> Self { self.value ^= 1; self } + /// Returns left child index of the current node. + pub const fn left_child(mut self) -> Self { + self.depth += 1; + self.value <<= 1; + self + } + + /// Returns right child index of the current node. + pub const fn right_child(mut self) -> Self { + self.depth += 1; + self.value = (self.value << 1) + 1; + self + } + // PROVIDERS // -------------------------------------------------------------------------------------------- diff --git a/src/merkle/store/mod.rs b/src/merkle/store/mod.rs index 40eaabd..d4c1ffb 100644 --- a/src/merkle/store/mod.rs +++ b/src/merkle/store/mod.rs @@ -1,6 +1,6 @@ use super::{ mmr::Mmr, BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerklePathSet, - MerkleTree, NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, ValuePath, Vec, Word, + MerkleTree, NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, TieredSmt, ValuePath, Vec, Word, }; use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; use core::borrow::Borrow; @@ -152,7 +152,6 @@ impl MerkleStore { /// The path starts at the sibling of the target leaf. /// /// # Errors - /// /// This method can return the following errors: /// - `RootNotInStore` if the `root` is not present in the store. /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store. @@ -440,6 +439,14 @@ impl From<&Mmr> for MerkleStore { } } +impl From<&TieredSmt> for MerkleStore { + fn from(value: &TieredSmt) -> Self { + let mut store = MerkleStore::new(); + store.extend(value.inner_nodes()); + store + } +} + impl FromIterator for MerkleStore { fn from_iter>(iter: T) -> Self { let mut store = MerkleStore::new(); diff --git a/src/merkle/tiered_smt/mod.rs b/src/merkle/tiered_smt/mod.rs index d3ca68a..eaf1cd8 100644 --- a/src/merkle/tiered_smt/mod.rs +++ b/src/merkle/tiered_smt/mod.rs @@ -1,6 +1,6 @@ use super::{ - BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, MerkleError, MerklePath, NodeIndex, Rpo256, - RpoDigest, StarkField, Vec, Word, EMPTY_WORD, ZERO, + BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, + Rpo256, RpoDigest, StarkField, Vec, Word, EMPTY_WORD, ZERO, }; use core::cmp; @@ -189,6 +189,52 @@ impl TieredSmt { old_value } + // ITERATORS + // -------------------------------------------------------------------------------------------- + + /// Returns an iterator over all inner nodes of this [TieredSmt] (i.e., nodes not at depths 16 + /// 32, 48, or 64). + /// + /// The iterator order is unspecified. + pub fn inner_nodes(&self) -> impl Iterator + '_ { + self.nodes.iter().filter_map(|(index, node)| { + if is_inner_node(index) { + Some(InnerNodeInfo { + value: node.into(), + left: self.get_node_unchecked(&index.left_child()).into(), + right: self.get_node_unchecked(&index.right_child()).into(), + }) + } else { + None + } + }) + } + + /// Returns an iterator over upper leaves (i.e., depth = 16, 32, or 48) for this [TieredSmt]. + /// + /// Each yielded item is a (node, key, value) tuple where key is a full un-truncated key (i.e., + /// with key[3] element unmodified). + /// + /// The iterator order is unspecified. + pub fn upper_leaves(&self) -> impl Iterator + '_ { + self.upper_leaves.iter().map(|(index, key)| { + let node = self.get_node_unchecked(index); + let value = self.get_value(*key); + (node, *key, value) + }) + } + + /// Returns an iterator over bottom leaves (i.e., depth = 64) of this [TieredSmt]. + /// + /// Each yielded item consists of the hash of the leaf and its contents, where contents is + /// a vector containing key-value pairs of entries storied in this leaf. Note that keys are + /// un-truncated keys (i.e., with key[3] element unmodified). + /// + /// The iterator order is unspecified. + pub fn bottom_leaves(&self) -> impl Iterator)> + '_ { + self.bottom_leaves.values().map(|leaf| (leaf.hash(), leaf.contents())) + } + // HELPER METHODS // -------------------------------------------------------------------------------------------- @@ -288,6 +334,7 @@ impl TieredSmt { } // update the root + self.nodes.insert(NodeIndex::root(), node); self.root = node; } } @@ -367,6 +414,12 @@ const fn get_index_tier(index: &NodeIndex) -> usize { } } +/// Returns true if the specified index is an index for an inner node (i.e., the depth is not 16, +/// 32, 48, or 64). +const fn is_inner_node(index: &NodeIndex) -> bool { + !matches!(index.depth(), 16 | 32 | 48 | 64) +} + // BOTTOM LEAF // ================================================================================================ @@ -378,16 +431,18 @@ const fn get_index_tier(index: &NodeIndex) -> usize { /// the same hash value. #[derive(Debug, Clone, PartialEq, Eq)] struct BottomLeaf { + prefix: u64, values: BTreeMap<[u64; 4], Word>, } impl BottomLeaf { /// Returns a new [BottomLeaf] with a single key-value pair added. pub fn new(key: RpoDigest, value: Word) -> Self { + let prefix = Word::from(key)[3].as_int(); let mut values = BTreeMap::new(); let key = get_remaining_path(key, TieredSmt::MAX_DEPTH as u32); values.insert(key.into(), value); - Self { values } + Self { prefix, values } } /// Adds a new key-value pair to this leaf. @@ -406,4 +461,22 @@ impl BottomLeaf { // TODO: hash in domain Rpo256::hash_elements(&elements) } + + /// Returns contents of this leaf as a vector of (key, value) pairs. + /// + /// The keys are returned in their un-truncated form. + pub fn contents(&self) -> Vec<(RpoDigest, Word)> { + self.values + .iter() + .map(|(key, val)| { + let key = RpoDigest::from([ + Felt::new(key[0]), + Felt::new(key[1]), + Felt::new(key[2]), + Felt::new(self.prefix), + ]); + (key, *val) + }) + .collect() + } } diff --git a/src/merkle/tiered_smt/tests.rs b/src/merkle/tiered_smt/tests.rs index 5b36434..f46a4ae 100644 --- a/src/merkle/tiered_smt/tests.rs +++ b/src/merkle/tiered_smt/tests.rs @@ -1,9 +1,7 @@ use super::{ - super::{ - super::{ONE, ZERO}, - Felt, MerkleStore, WORD_SIZE, - }, - get_remaining_path, EmptySubtreeRoots, NodeIndex, Rpo256, RpoDigest, TieredSmt, Vec, Word, + super::{super::ONE, Felt, MerkleStore, WORD_SIZE, ZERO}, + get_remaining_path, EmptySubtreeRoots, InnerNodeInfo, NodeIndex, Rpo256, RpoDigest, TieredSmt, + Vec, Word, }; #[test] @@ -32,6 +30,17 @@ fn tsmt_insert_one() { // make sure the paths we get from the store and the tree match let expected_path = store.get_path(tree_root, index).unwrap(); assert_eq!(smt.get_path(index).unwrap(), expected_path.path); + + // make sure inner nodes match + let expected_nodes = get_non_empty_nodes(&store); + let actual_nodes = smt.inner_nodes().collect::>(); + assert_eq!(actual_nodes.len(), expected_nodes.len()); + actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); + + // make sure leaves are returned correctly + let mut leaves = smt.upper_leaves(); + assert_eq!(leaves.next(), Some((leaf_node, key, value))); + assert_eq!(leaves.next(), None); } #[test] @@ -76,6 +85,18 @@ fn tsmt_insert_two_16() { assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); let expected_path = store.get_path(tree_root, index_b).unwrap().path; assert_eq!(smt.get_path(index_b).unwrap(), expected_path); + + // make sure inner nodes match - the store contains more entries because it keeps track of + // all prior state - so, we don't check that the number of inner nodes is the same in both + let expected_nodes = get_non_empty_nodes(&store); + let actual_nodes = smt.inner_nodes().collect::>(); + actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); + + // make sure leaves are returned correctly + let mut leaves = smt.upper_leaves(); + assert_eq!(leaves.next(), Some((leaf_node_a, key_a, val_a))); + assert_eq!(leaves.next(), Some((leaf_node_b, key_b, val_b))); + assert_eq!(leaves.next(), None); } #[test] @@ -120,6 +141,12 @@ fn tsmt_insert_two_32() { assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); let expected_path = store.get_path(tree_root, index_b).unwrap().path; assert_eq!(smt.get_path(index_b).unwrap(), expected_path); + + // make sure inner nodes match - the store contains more entries because it keeps track of + // all prior state - so, we don't check that the number of inner nodes is the same in both + let expected_nodes = get_non_empty_nodes(&store); + let actual_nodes = smt.inner_nodes().collect::>(); + actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); } #[test] @@ -182,6 +209,12 @@ fn tsmt_insert_three() { assert_eq!(smt.get_node(index_c).unwrap(), leaf_node_c); let expected_path = store.get_path(tree_root, index_c).unwrap().path; assert_eq!(smt.get_path(index_c).unwrap(), expected_path); + + // make sure inner nodes match - the store contains more entries because it keeps track of + // all prior state - so, we don't check that the number of inner nodes is the same in both + let expected_nodes = get_non_empty_nodes(&store); + let actual_nodes = smt.inner_nodes().collect::>(); + actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); } #[test] @@ -211,6 +244,12 @@ fn tsmt_update() { assert_eq!(smt.get_node(index).unwrap(), leaf_node); let expected_path = store.get_path(tree_root, index).unwrap().path; assert_eq!(smt.get_path(index).unwrap(), expected_path); + + // make sure inner nodes match - the store contains more entries because it keeps track of + // all prior state - so, we don't check that the number of inner nodes is the same in both + let expected_nodes = get_non_empty_nodes(&store); + let actual_nodes = smt.inner_nodes().collect::>(); + actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); } // BOTTOM TIER TESTS @@ -254,6 +293,17 @@ fn tsmt_bottom_tier() { assert_eq!(smt.get_node(index).unwrap(), leaf_node); let expected_path = store.get_path(tree_root, index).unwrap().path; assert_eq!(smt.get_path(index).unwrap(), expected_path); + + // make sure inner nodes match - the store contains more entries because it keeps track of + // all prior state - so, we don't check that the number of inner nodes is the same in both + let expected_nodes = get_non_empty_nodes(&store); + let actual_nodes = smt.inner_nodes().collect::>(); + actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); + + // make sure leaves are returned correctly + let mut leaves = smt.bottom_leaves(); + assert_eq!(leaves.next(), Some((leaf_node, vec![(key_b, val_b), (key_a, val_a)]))); + assert_eq!(leaves.next(), None); } #[test] @@ -298,6 +348,18 @@ fn tsmt_bottom_tier_two() { assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b); let expected_path = store.get_path(tree_root, index_b).unwrap().path; assert_eq!(smt.get_path(index_b).unwrap(), expected_path); + + // make sure inner nodes match - the store contains more entries because it keeps track of + // all prior state - so, we don't check that the number of inner nodes is the same in both + let expected_nodes = get_non_empty_nodes(&store); + let actual_nodes = smt.inner_nodes().collect::>(); + actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); + + // make sure leaves are returned correctly + let mut leaves = smt.bottom_leaves(); + assert_eq!(leaves.next(), Some((leaf_node_b, vec![(key_b, val_b)]))); + assert_eq!(leaves.next(), Some((leaf_node_a, vec![(key_a, val_a)]))); + assert_eq!(leaves.next(), None); } // ERROR TESTS @@ -366,3 +428,14 @@ fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest { Rpo256::hash_elements(&elements) } + +fn get_non_empty_nodes(store: &MerkleStore) -> Vec { + store + .inner_nodes() + .filter(|node| !is_empty_subtree(&RpoDigest::from(node.value))) + .collect::>() +} + +fn is_empty_subtree(node: &RpoDigest) -> bool { + EmptySubtreeRoots::empty_hashes(255).contains(&node) +}