diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs index 3f33535..a615e18 100644 --- a/src/merkle/partial_mt/mod.rs +++ b/src/merkle/partial_mt/mod.rs @@ -118,7 +118,7 @@ impl PartialMerkleTree { // fill layers without nodes with empty vector for depth in 0..max_depth { - layers.entry(depth).or_insert(vec![]); + layers.entry(depth).or_default(); } let mut layer_iter = layers.into_values().rev(); @@ -372,7 +372,6 @@ impl PartialMerkleTree { return Ok(old_value); } - let mut node_index = node_index; let mut value = value.into(); for _ in 0..node_index.depth() { let sibling = self.nodes.get(&node_index.sibling()).expect("sibling should exist"); diff --git a/src/merkle/tiered_smt/mod.rs b/src/merkle/tiered_smt/mod.rs index 3269517..5379d20 100644 --- a/src/merkle/tiered_smt/mod.rs +++ b/src/merkle/tiered_smt/mod.rs @@ -1,8 +1,14 @@ use super::{ - BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, - Rpo256, RpoDigest, StarkField, Vec, Word, ZERO, + BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, + Rpo256, RpoDigest, StarkField, Vec, Word, }; -use core::cmp; +use core::{cmp, ops::Deref}; + +mod nodes; +use nodes::NodeStore; + +mod values; +use values::ValueStore; #[cfg(test)] mod tests; @@ -18,27 +24,22 @@ mod tests; /// of depth 64 (i.e., leaves at depth 64 are set to [ZERO; 4]). As non-empty values are inserted /// into the tree they are added to the first available tier. /// -/// For example, when the first key-value is inserted, it will be stored in a node at depth 16 -/// such that the first 16 bits of the key determine the position of the node at depth 16. If -/// another value with a key sharing the same 16-bit prefix is inserted, both values move into -/// the next tier (depth 32). This process is repeated until values end up at tier 64. If multiple -/// values have keys with a common 64-bit prefix, such key-value pairs are stored in a sorted list -/// at the last tier (depth = 64). +/// For example, when the first key-value pair is inserted, it will be stored in a node at depth +/// 16 such that the 16 most significant bits of the key determine the position of the node at +/// depth 16. If another value with a key sharing the same 16-bit prefix is inserted, both values +/// move into the next tier (depth 32). This process is repeated until values end up at the bottom +/// tier (depth 64). If multiple values have keys with a common 64-bit prefix, such key-value pairs +/// are stored in a sorted list at the bottom tier. /// /// To differentiate between internal and leaf nodes, node values are computed as follows: /// - Internal nodes: hash(left_child, right_child). -/// - Leaf node at depths 16, 32, or 64: hash(rem_key, value, domain=depth). -/// - Leaf node at depth 64: hash([rem_key_0, value_0, ..., rem_key_n, value_n, domain=64]). -/// -/// Where rem_key is computed by replacing d most significant bits of the key with zeros where d -/// is depth (i.e., for a leaf at depth 16, we replace 16 most significant bits of the key with 0). +/// - Leaf node at depths 16, 32, or 64: hash(key, value, domain=depth). +/// - Leaf node at depth 64: hash([key_0, value_0, ..., key_n, value_n, domain=64]). #[derive(Debug, Clone, PartialEq, Eq)] pub struct TieredSmt { root: RpoDigest, - nodes: BTreeMap, - upper_leaves: BTreeMap, // node_index |-> key map - bottom_leaves: BTreeMap, // leaves of depth 64 - values: BTreeMap, + nodes: NodeStore, + values: ValueStore, } impl TieredSmt { @@ -106,8 +107,7 @@ impl TieredSmt { /// when a leaf node with the same index prefix exists at a tier higher than the requested /// node. pub fn get_node(&self, index: NodeIndex) -> Result { - self.validate_node_access(index)?; - Ok(self.get_node_unchecked(&index)) + self.nodes.get_node(index) } /// Returns a Merkle path from the node at the specified index to the root. @@ -120,17 +120,8 @@ impl TieredSmt { /// - The node with the specified index does not exists in the Merkle tree. This is possible /// when a leaf node with the same index prefix exists at a tier higher than the node to /// which the path is requested. - pub fn get_path(&self, mut index: NodeIndex) -> Result { - self.validate_node_access(index)?; - - let mut path = Vec::with_capacity(index.depth() as usize); - for _ in 0..index.depth() { - let node = self.get_node_unchecked(&index.sibling()); - path.push(node); - index.move_up(); - } - - Ok(path.into()) + pub fn get_path(&self, index: NodeIndex) -> Result { + self.nodes.get_path(index) } /// Returns the value associated with the specified key. @@ -151,238 +142,279 @@ impl TieredSmt { /// /// If the value for the specified key was not previously set, [ZERO; 4] is returned. pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word { - // insert the value into the key-value map, and if nothing has changed, return - let old_value = self.values.insert(key, value).unwrap_or(Self::EMPTY_VALUE); - if old_value == value { - return old_value; + // if an empty value is being inserted, remove the leaf node to make it look as if the + // value was never inserted + if value == Self::EMPTY_VALUE { + return self.remove_leaf_node(key); } - // determine the index for the value node; this index could have 3 different meanings: - // - it points to a root of an empty subtree (excluding depth = 64); in this case, we can - // replace the node with the value node immediately. - // - it points to a node at the bottom tier (i.e., depth = 64); in this case, we need to - // process bottom-tier insertion which will be handled by insert_node(). - // - it points to a leaf node; this node could be a node with the same key or a different - // key with a common prefix; in the latter case, we'll need to move the leaf to a lower - // tier; for this scenario the `leaf_key` will contain the key of the leaf node - let (mut index, leaf_key) = self.get_insert_location(&key); - - // if the returned index points to a leaf, and this leaf is for a different key, we need - // to move the leaf to a lower tier - if let Some(other_key) = leaf_key { - if other_key != key { - // determine how far down the tree should we move the existing leaf - let common_prefix_len = get_common_prefix_tier(&key, &other_key); - let depth = cmp::min(common_prefix_len + Self::TIER_SIZE, Self::MAX_DEPTH); - - // move the leaf to the new location; this requires first removing the existing - // index, re-computing node value, and inserting the node at a new location - let other_index = key_to_index(&other_key, depth); - let other_value = *self.values.get(&other_key).expect("no value for other key"); - self.upper_leaves.remove(&index).expect("other node key not in map"); - self.insert_node(other_index, other_key, other_value); - - // the new leaf also needs to move down to the same tier - index = key_to_index(&key, depth); + // insert the value into the value store, and if the key was already in the store, update + // it with the new value + if let Some(old_value) = self.values.insert(key, value) { + if old_value != value { + // if the new value is different from the old value, determine the location of + // the leaf node for this key, build the node, and update the root + let (index, leaf_exists) = self.nodes.get_leaf_index(&key); + debug_assert!(leaf_exists); + let node = self.build_leaf_node(index, key, value); + self.root = self.nodes.update_leaf_node(index, node); } - } + return old_value; + }; - // insert the node and return the old value - self.insert_node(index, key, value); - old_value + // determine the location for the leaf node; this index could have 3 different meanings: + // - it points to a root of an empty subtree or an empty node at depth 64; in this case, + // we can replace the node with the value node immediately. + // - it points to an existing leaf at the bottom tier (i.e., depth = 64); in this case, + // we need to process update the bottom leaf. + // - it points to an existing leaf node for a different key with the same prefix (same + // key case was handled above); in this case, we need to move the leaf to a lower tier + let (index, leaf_exists) = self.nodes.get_leaf_index(&key); + + self.root = if leaf_exists && index.depth() == Self::MAX_DEPTH { + // returned index points to a leaf at the bottom tier + let node = self.build_leaf_node(index, key, value); + self.nodes.update_leaf_node(index, node) + } else if leaf_exists { + // returned index pointes to a leaf for a different key with the same prefix + + // get the key-value pair for the key with the same prefix; since the key-value + // pair has already been inserted into the value store, we need to filter it out + // when looking for the other key-value pair + let (other_key, other_value) = self + .values + .get_first_filtered(index_to_prefix(&index), &key) + .expect("other key-value pair not found"); + + // determine how far down the tree should we move the leaves + let common_prefix_len = get_common_prefix_tier_depth(&key, other_key); + let depth = cmp::min(common_prefix_len + Self::TIER_SIZE, Self::MAX_DEPTH); + + // compute node locations for new and existing key-value paris + let new_index = LeafNodeIndex::from_key(&key, depth); + let other_index = LeafNodeIndex::from_key(other_key, depth); + + // compute node values for the new and existing key-value pairs + let new_node = self.build_leaf_node(new_index, key, value); + let other_node = self.build_leaf_node(other_index, *other_key, *other_value); + + // replace the leaf located at index with a subtree containing nodes for new and + // existing key-value paris + self.nodes.replace_leaf_with_subtree( + index, + [(new_index, new_node), (other_index, other_node)], + ) + } else { + // returned index points to an empty subtree or an empty leaf at the bottom tier + let node = self.build_leaf_node(index, key, value); + self.nodes.insert_leaf_node(index, node) + }; + + Self::EMPTY_VALUE } // ITERATORS // -------------------------------------------------------------------------------------------- + /// Returns an iterator over all key-value pairs in this [TieredSmt]. + pub fn iter(&self) -> impl Iterator { + self.values.iter() + } + /// 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, - left: self.get_node_unchecked(&index.left_child()), - right: self.get_node_unchecked(&index.right_child()), - }) - } else { - None - } - }) + self.nodes.inner_nodes() } - /// 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). + /// Returns an iterator over upper leaves (i.e., depth = 16, 32, or 48) for this [TieredSmt] + /// where each yielded item is a (node, key, value) tuple. /// /// 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) + self.nodes.upper_leaves().map(|(index, node)| { + let key_prefix = index_to_prefix(index); + let (key, value) = self.values.get_first(key_prefix).expect("upper leaf not found"); + debug_assert_eq!(*index, LeafNodeIndex::from_key(key, index.depth()).into()); + (*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). + /// a vector containing key-value pairs of entries storied in this leaf. /// /// The iterator order is unspecified. pub fn bottom_leaves(&self) -> impl Iterator)> + '_ { - self.bottom_leaves.values().map(|leaf| (leaf.hash(), leaf.contents())) + self.nodes.bottom_leaves().map(|(&prefix, node)| { + let values = self.values.get_all(prefix).expect("bottom leaf not found"); + (*node, values) + }) } // HELPER METHODS // -------------------------------------------------------------------------------------------- - /// Checks if the specified index is valid in the context of this Merkle tree. + /// Removes the node holding the key-value pair for the specified key from this tree, and + /// returns the value associated with the specified key. /// - /// # Errors - /// Returns an error if: - /// - The specified index depth is 0 or greater than 64. - /// - The node for the specified index does not exists in the Merkle tree. This is possible - /// when an ancestors of the specified index is a leaf node. - fn validate_node_access(&self, index: NodeIndex) -> Result<(), MerkleError> { - if index.is_root() { - return Err(MerkleError::DepthTooSmall(index.depth())); - } else if index.depth() > Self::MAX_DEPTH { - return Err(MerkleError::DepthTooBig(index.depth() as u64)); - } else { - // make sure that there are no leaf nodes in the ancestors of the index; since leaf - // nodes can live at specific depth, we just need to check these depths. - let tier = get_index_tier(&index); - let mut tier_index = index; - for &depth in Self::TIER_DEPTHS[..tier].iter().rev() { - tier_index.move_up_to(depth); - if self.upper_leaves.contains_key(&tier_index) { - return Err(MerkleError::NodeNotInSet(index)); - } - } - } - - Ok(()) - } + /// If no value was associated with the specified key, [ZERO; 4] is returned. + fn remove_leaf_node(&mut self, key: RpoDigest) -> Word { + // remove the key-value pair from the value store; if no value was associated with the + // specified key, return. + let old_value = match self.values.remove(&key) { + Some(old_value) => old_value, + None => return Self::EMPTY_VALUE, + }; - /// Returns a node at the specified index. If the node does not exist at this index, a root - /// for an empty subtree at the index's depth is returned. - /// - /// Unlike [TieredSmt::get_node()] this does not perform any checks to verify that the returned - /// node is valid in the context of this tree. - fn get_node_unchecked(&self, index: &NodeIndex) -> RpoDigest { - match self.nodes.get(index) { - Some(node) => *node, - None => EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[index.depth() as usize], + // determine the location of the leaf holding the key-value pair to be removed + let (index, leaf_exists) = self.nodes.get_leaf_index(&key); + debug_assert!(leaf_exists); + + // if the leaf is at the bottom tier and after removing the key-value pair from it, the + // leaf is still not empty, just recompute its hash and update the leaf node. + if index.depth() == Self::MAX_DEPTH { + if let Some(values) = self.values.get_all(index.value()) { + let node = hash_bottom_leaf(&values); + self.root = self.nodes.update_leaf_node(index, node); + return old_value; + }; } - } - /// Returns an index at which a node for the specified key should be inserted. If a leaf node - /// already exists at that index, returns the key associated with that leaf node. - /// - /// In case the index falls into the bottom tier (depth = 64), leaf node key is not returned - /// as the bottom tier may contain multiple key-value pairs in the same leaf. - fn get_insert_location(&self, key: &RpoDigest) -> (NodeIndex, Option) { - // traverse the tree from the root down checking nodes at tiers 16, 32, and 48. Return if - // a node at any of the tiers is either a leaf or a root of an empty subtree. - let mse = Word::from(key)[3].as_int(); - for depth in (Self::TIER_DEPTHS[0]..Self::MAX_DEPTH).step_by(Self::TIER_SIZE as usize) { - let index = NodeIndex::new_unchecked(depth, mse >> (Self::MAX_DEPTH - depth)); - if let Some(leaf_key) = self.upper_leaves.get(&index) { - return (index, Some(*leaf_key)); - } else if !self.nodes.contains_key(&index) { - return (index, None); - } + // if the removed key-value pair has a lone sibling at the current tier with a root at + // higher tier, we need to move the sibling to a higher tier + if let Some((sib_key, sib_val, new_sib_index)) = self.values.get_lone_sibling(index) { + // determine the current index of the sibling node + let sib_index = LeafNodeIndex::from_key(sib_key, index.depth()); + debug_assert!(sib_index.depth() > new_sib_index.depth()); + + // compute node value for the new location of the sibling leaf and replace the subtree + // with this leaf node + let node = self.build_leaf_node(new_sib_index, *sib_key, *sib_val); + let new_sib_depth = new_sib_index.depth(); + self.root = self.nodes.replace_subtree_with_leaf(index, sib_index, new_sib_depth, node); + } else { + // if the removed key-value pair did not have a sibling at the current tier with a + // root at higher tiers, just clear the leaf node + self.root = self.nodes.clear_leaf_node(index); } - // if we got here, that means all of the nodes checked so far are internal nodes, and - // the new node would need to be inserted in the bottom tier. - let index = NodeIndex::new_unchecked(Self::MAX_DEPTH, mse); - (index, None) + old_value } - /// Inserts the provided key-value pair at the specified index and updates the root of this - /// Merkle tree by recomputing the path to the root. - fn insert_node(&mut self, mut index: NodeIndex, key: RpoDigest, value: Word) { + /// Builds and returns a leaf node value for the node located as the specified index. + /// + /// This method assumes that the key-value pair for the node has already been inserted into + /// the value store, however, for depths 16, 32, and 48, the node is computed directly from + /// the passed-in values (for depth 64, the value store is queried to get all the key-value + /// pairs located at the specified index). + fn build_leaf_node(&self, index: LeafNodeIndex, key: RpoDigest, value: Word) -> RpoDigest { let depth = index.depth(); // insert the key into index-key map and compute the new value of the node - let mut node = if index.depth() == Self::MAX_DEPTH { + if index.depth() == Self::MAX_DEPTH { // for the bottom tier, we add the key-value pair to the existing leaf, or create a // new leaf with this key-value pair - self.bottom_leaves - .entry(index.value()) - .and_modify(|leaves| leaves.add_value(key, value)) - .or_insert(BottomLeaf::new(key, value)) - .hash() + let values = self.values.get_all(index.value()).unwrap(); + hash_bottom_leaf(&values) } else { - // for the upper tiers, we just update the index-key map and compute the value of the - // node - self.upper_leaves.insert(index, key); - // the node value is computed as: hash(remaining_key || value, domain = depth) - let remaining_path = get_remaining_path(key, depth.into()); - Rpo256::merge_in_domain(&[remaining_path, value.into()], depth.into()) - }; - - // insert the node and update the path from the node to the root - for _ in 0..index.depth() { - self.nodes.insert(index, node); - let sibling = self.get_node_unchecked(&index.sibling()); - node = Rpo256::merge(&index.build_node(node, sibling)); - index.move_up(); + debug_assert_eq!(self.values.get_first(index_to_prefix(&index)), Some(&(key, value))); + hash_upper_leaf(key, value, depth) } - - // update the root - self.nodes.insert(NodeIndex::root(), node); - self.root = node; } } impl Default for TieredSmt { fn default() -> Self { + let root = EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[0]; Self { - root: EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[0], - nodes: BTreeMap::new(), - upper_leaves: BTreeMap::new(), - bottom_leaves: BTreeMap::new(), - values: BTreeMap::new(), + root, + nodes: NodeStore::new(root), + values: ValueStore::default(), } } } +// LEAF NODE INDEX +// ================================================================================================ +/// A wrapper around [NodeIndex] to provide type-safe references to nodes at depths 16, 32, 48, and +/// 64. +#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)] +pub struct LeafNodeIndex(NodeIndex); + +impl LeafNodeIndex { + /// Returns a new [LeafNodeIndex] instantiated from the provided [NodeIndex]. + /// + /// In debug mode, panics if index depth is not 16, 32, 48, or 64. + pub fn new(index: NodeIndex) -> Self { + // check if the depth is 16, 32, 48, or 64; this works because for a valid depth, + // depth - 16, can be 0, 16, 32, or 48 - i.e., the value is either 0 or any of the 4th + // or 5th bits are set. We can test for this by computing a bitwise AND with a value + // which has all but the 4th and 5th bits set (which is !48). + debug_assert_eq!(((index.depth() - 16) & !48), 0, "invalid tier depth {}", index.depth()); + Self(index) + } + + /// Returns a new [LeafNodeIndex] instantiated from the specified key inserted at the specified + /// depth. + /// + /// The value for the key is computed by taking n most significant bits from the most significant + /// element of the key, where n is the specified depth. + pub fn from_key(key: &RpoDigest, depth: u8) -> Self { + let mse = get_key_prefix(key); + Self::new(NodeIndex::new_unchecked(depth, mse >> (TieredSmt::MAX_DEPTH - depth))) + } + + /// Returns a new [LeafNodeIndex] instantiated for testing purposes. + #[cfg(test)] + pub fn make(depth: u8, value: u64) -> Self { + Self::new(NodeIndex::make(depth, value)) + } + + /// Traverses towards the root until the specified depth is reached. + /// + /// The new depth must be a valid tier depth - i.e., 16, 32, 48, or 64. + pub fn move_up_to(&mut self, depth: u8) { + debug_assert_eq!(((depth - 16) & !48), 0, "invalid tier depth: {depth}"); + self.0.move_up_to(depth); + } +} + +impl Deref for LeafNodeIndex { + type Target = NodeIndex; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl From for LeafNodeIndex { + fn from(value: NodeIndex) -> Self { + Self::new(value) + } +} + +impl From for NodeIndex { + fn from(value: LeafNodeIndex) -> Self { + value.0 + } +} + // HELPER FUNCTIONS // ================================================================================================ -/// Returns the remaining path for the specified key at the specified depth. -/// -/// Remaining path is computed by setting n most significant bits of the key to zeros, where n is -/// the specified depth. -fn get_remaining_path(key: RpoDigest, depth: u32) -> RpoDigest { - let mut key = Word::from(key); - key[3] = if depth == 64 { - ZERO - } else { - // remove `depth` bits from the most significant key element - ((key[3].as_int() << depth) >> depth).into() - }; - key.into() +/// Returns the value representing the 64 most significant bits of the specified key. +fn get_key_prefix(key: &RpoDigest) -> u64 { + Word::from(key)[3].as_int() } -/// Returns index for the specified key inserted at the specified depth. -/// -/// The value for the key is computed by taking n most significant bits from the most significant -/// element of the key, where n is the specified depth. -fn key_to_index(key: &RpoDigest, depth: u8) -> NodeIndex { - let mse = Word::from(key)[3].as_int(); - let value = match depth { - 16 | 32 | 48 | 64 => mse >> ((TieredSmt::MAX_DEPTH - depth) as u32), - _ => unreachable!("invalid depth: {depth}"), - }; - NodeIndex::new_unchecked(depth, value) +/// Returns the index value shifted to be in the most significant bit positions of the returned +/// u64 value. +fn index_to_prefix(index: &NodeIndex) -> u64 { + index.value() << (TieredSmt::MAX_DEPTH - index.depth()) } /// Returns tiered common prefix length between the most significant elements of the provided keys. @@ -393,93 +425,34 @@ fn key_to_index(key: &RpoDigest, depth: u8) -> NodeIndex { /// - returns 32 if the common prefix is between 32 and 47 bits. /// - returns 16 if the common prefix is between 16 and 31 bits. /// - returns 0 if the common prefix is fewer than 16 bits. -fn get_common_prefix_tier(key1: &RpoDigest, key2: &RpoDigest) -> u8 { - let e1 = Word::from(key1)[3].as_int(); - let e2 = Word::from(key2)[3].as_int(); +fn get_common_prefix_tier_depth(key1: &RpoDigest, key2: &RpoDigest) -> u8 { + let e1 = get_key_prefix(key1); + let e2 = get_key_prefix(key2); let ex = (e1 ^ e2).leading_zeros() as u8; (ex / 16) * 16 } -/// Returns a tier for the specified index. +/// Computes node value for leaves at tiers 16, 32, or 48. /// -/// The tiers are defined as follows: -/// - Tier 0: depth 0 through 16 (inclusive). -/// - Tier 1: depth 17 through 32 (inclusive). -/// - Tier 2: depth 33 through 48 (inclusive). -/// - Tier 3: depth 49 through 64 (inclusive). -const fn get_index_tier(index: &NodeIndex) -> usize { - debug_assert!(index.depth() <= TieredSmt::MAX_DEPTH, "invalid depth"); - match index.depth() { - 0..=16 => 0, - 17..=32 => 1, - 33..=48 => 2, - _ => 3, - } -} - -/// 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) +/// Node value is computed as: hash(key || value, domain = depth). +pub fn hash_upper_leaf(key: RpoDigest, value: Word, depth: u8) -> RpoDigest { + const NUM_UPPER_TIERS: usize = TieredSmt::TIER_DEPTHS.len() - 1; + debug_assert!(TieredSmt::TIER_DEPTHS[..NUM_UPPER_TIERS].contains(&depth)); + Rpo256::merge_in_domain(&[key, value.into()], depth.into()) } -// BOTTOM LEAF -// ================================================================================================ - -/// Stores contents of the bottom leaf (i.e., leaf at depth = 64) in a [TieredSmt]. +/// Computes node value for leaves at the bottom tier (depth 64). /// -/// Bottom leaf can contain one or more key-value pairs all sharing the same 64-bit key prefix. -/// The values are sorted by key to make sure the structure of the leaf is independent of the -/// insertion order. This guarantees that a leaf with the same set of key-value pairs always has -/// 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 { prefix, values } - } - - /// Adds a new key-value pair to this leaf. - pub fn add_value(&mut self, key: RpoDigest, value: Word) { - let key = get_remaining_path(key, TieredSmt::MAX_DEPTH as u32); - self.values.insert(key.into(), value); - } - - /// Computes a hash of this leaf. - pub fn hash(&self) -> RpoDigest { - let mut elements = Vec::with_capacity(self.values.len() * 2); - for (key, val) in self.values.iter() { - key.iter().for_each(|&v| elements.push(Felt::new(v))); - elements.extend_from_slice(val.as_slice()); - } - // 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() +/// Node value is computed as: hash([key_0, value_0, ..., key_n, value_n], domain=64). +/// +/// TODO: when hashing in domain is implemented for `hash_elements()`, combine this function with +/// `hash_upper_leaf()` function. +pub fn hash_bottom_leaf(values: &[(RpoDigest, Word)]) -> RpoDigest { + let mut elements = Vec::with_capacity(values.len() * 8); + for (key, val) in values.iter() { + elements.extend_from_slice(key.as_elements()); + elements.extend_from_slice(val.as_slice()); } + // TODO: hash in domain + Rpo256::hash_elements(&elements) } diff --git a/src/merkle/tiered_smt/nodes.rs b/src/merkle/tiered_smt/nodes.rs new file mode 100644 index 0000000..7135c6c --- /dev/null +++ b/src/merkle/tiered_smt/nodes.rs @@ -0,0 +1,374 @@ +use super::{ + BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, LeafNodeIndex, MerkleError, MerklePath, + NodeIndex, Rpo256, RpoDigest, Vec, +}; + +// CONSTANTS +// ================================================================================================ + +/// The number of levels between tiers. +const TIER_SIZE: u8 = super::TieredSmt::TIER_SIZE; + +/// Depths at which leaves can exist in a tiered SMT. +const TIER_DEPTHS: [u8; 4] = super::TieredSmt::TIER_DEPTHS; + +/// Maximum node depth. This is also the bottom tier of the tree. +const MAX_DEPTH: u8 = super::TieredSmt::MAX_DEPTH; + +// NODE STORE +// ================================================================================================ + +/// A store of nodes for a Tiered Sparse Merkle tree. +/// +/// The store contains information about all nodes as well as information about which of the nodes +/// represent leaf nodes in a Tiered Sparse Merkle tree. In the current implementation, [BTreeSet]s +/// are used to determine the position of the leaves in the tree. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct NodeStore { + nodes: BTreeMap, + upper_leaves: BTreeSet, + bottom_leaves: BTreeSet, +} + +impl NodeStore { + // CONSTRUCTOR + // -------------------------------------------------------------------------------------------- + /// Returns a new instance of [NodeStore] instantiated with the specified root node. + /// + /// Root node is assumed to be a root of an empty sparse Merkle tree. + pub fn new(root_node: RpoDigest) -> Self { + let mut nodes = BTreeMap::default(); + nodes.insert(NodeIndex::root(), root_node); + + Self { + nodes, + upper_leaves: BTreeSet::default(), + bottom_leaves: BTreeSet::default(), + } + } + + // PUBLIC ACCESSORS + // -------------------------------------------------------------------------------------------- + + /// Returns a node at the specified index. + /// + /// # Errors + /// Returns an error if: + /// - The specified index depth is 0 or greater than 64. + /// - The node with the specified index does not exists in the Merkle tree. This is possible + /// when a leaf node with the same index prefix exists at a tier higher than the requested + /// node. + pub fn get_node(&self, index: NodeIndex) -> Result { + self.validate_node_access(index)?; + Ok(self.get_node_unchecked(&index)) + } + + /// Returns a Merkle path from the node at the specified index to the root. + /// + /// The node itself is not included in the path. + /// + /// # Errors + /// Returns an error if: + /// - The specified index depth is 0 or greater than 64. + /// - The node with the specified index does not exists in the Merkle tree. This is possible + /// when a leaf node with the same index prefix exists at a tier higher than the node to + /// which the path is requested. + pub fn get_path(&self, mut index: NodeIndex) -> Result { + self.validate_node_access(index)?; + + let mut path = Vec::with_capacity(index.depth() as usize); + for _ in 0..index.depth() { + let node = self.get_node_unchecked(&index.sibling()); + path.push(node); + index.move_up(); + } + + Ok(path.into()) + } + + /// Returns an index at which a leaf node for the specified key should be inserted. + /// + /// The second value in the returned tuple is set to true if the node at the returned index + /// is already a leaf node. + pub fn get_leaf_index(&self, key: &RpoDigest) -> (LeafNodeIndex, bool) { + // traverse the tree from the root down checking nodes at tiers 16, 32, and 48. Return if + // a node at any of the tiers is either a leaf or a root of an empty subtree. + const NUM_UPPER_TIERS: usize = TIER_DEPTHS.len() - 1; + for &tier_depth in TIER_DEPTHS[..NUM_UPPER_TIERS].iter() { + let index = LeafNodeIndex::from_key(key, tier_depth); + if self.upper_leaves.contains(&index) { + return (index, true); + } else if !self.nodes.contains_key(&index) { + return (index, false); + } + } + + // if we got here, that means all of the nodes checked so far are internal nodes, and + // the new node would need to be inserted in the bottom tier. + let index = LeafNodeIndex::from_key(key, MAX_DEPTH); + (index, self.bottom_leaves.contains(&index.value())) + } + + // ITERATORS + // -------------------------------------------------------------------------------------------- + + /// Returns an iterator over all inner nodes of the Tiered Sparse Merkle tree (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 self.is_internal_node(index) { + Some(InnerNodeInfo { + value: *node, + left: self.get_node_unchecked(&index.left_child()), + right: self.get_node_unchecked(&index.right_child()), + }) + } else { + None + } + }) + } + + /// Returns an iterator over the upper leaves (i.e., leaves with depths 16, 32, 48) of the + /// Tiered Sparse Merkle tree. + pub fn upper_leaves(&self) -> impl Iterator { + self.upper_leaves.iter().map(|index| (index, &self.nodes[index])) + } + + /// Returns an iterator over the bottom leaves (i.e., leaves with depth 64) of the Tiered + /// Sparse Merkle tree. + pub fn bottom_leaves(&self) -> impl Iterator { + self.bottom_leaves.iter().map(|value| { + let index = NodeIndex::new_unchecked(MAX_DEPTH, *value); + (value, &self.nodes[&index]) + }) + } + + // STATE MUTATORS + // -------------------------------------------------------------------------------------------- + + /// Replaces the leaf node at the specified index with a tree consisting of two leaves located + /// at the specified indexes. Recomputes and returns the new root. + pub fn replace_leaf_with_subtree( + &mut self, + leaf_index: LeafNodeIndex, + subtree_leaves: [(LeafNodeIndex, RpoDigest); 2], + ) -> RpoDigest { + debug_assert!(self.is_non_empty_leaf(&leaf_index)); + debug_assert!(!is_empty_root(&subtree_leaves[0].1)); + debug_assert!(!is_empty_root(&subtree_leaves[1].1)); + debug_assert_eq!(subtree_leaves[0].0.depth(), subtree_leaves[1].0.depth()); + debug_assert!(leaf_index.depth() < subtree_leaves[0].0.depth()); + + self.upper_leaves.remove(&leaf_index); + + if subtree_leaves[0].0 == subtree_leaves[1].0 { + // if the subtree is for a single node at depth 64, we only need to insert one node + debug_assert_eq!(subtree_leaves[0].0.depth(), MAX_DEPTH); + debug_assert_eq!(subtree_leaves[0].1, subtree_leaves[1].1); + self.insert_leaf_node(subtree_leaves[0].0, subtree_leaves[0].1) + } else { + self.insert_leaf_node(subtree_leaves[0].0, subtree_leaves[0].1); + self.insert_leaf_node(subtree_leaves[1].0, subtree_leaves[1].1) + } + } + + /// Replaces a subtree containing the retained and the removed leaf nodes, with a leaf node + /// containing the retained leaf. + /// + /// This has the effect of deleting the the node at the `removed_leaf` index from the tree, + /// moving the node at the `retained_leaf` index up to the tier specified by `new_depth`. + pub fn replace_subtree_with_leaf( + &mut self, + removed_leaf: LeafNodeIndex, + retained_leaf: LeafNodeIndex, + new_depth: u8, + node: RpoDigest, + ) -> RpoDigest { + debug_assert!(!is_empty_root(&node)); + debug_assert!(self.is_non_empty_leaf(&removed_leaf)); + debug_assert!(self.is_non_empty_leaf(&retained_leaf)); + debug_assert_eq!(removed_leaf.depth(), retained_leaf.depth()); + debug_assert!(removed_leaf.depth() > new_depth); + + // clear leaf flags + if removed_leaf.depth() == MAX_DEPTH { + self.bottom_leaves.remove(&removed_leaf.value()); + self.bottom_leaves.remove(&retained_leaf.value()); + } else { + self.upper_leaves.remove(&removed_leaf); + self.upper_leaves.remove(&retained_leaf); + } + + // remove the branches leading up to the tier to which the retained leaf is to be moved + self.remove_branch(removed_leaf, new_depth); + self.remove_branch(retained_leaf, new_depth); + + // compute the index of the common root for retained and removed leaves + let mut new_index = retained_leaf; + new_index.move_up_to(new_depth); + + // insert the node at the root index + self.insert_leaf_node(new_index, node) + } + + /// Inserts the specified node at the specified index; recomputes and returns the new root + /// of the Tiered Sparse Merkle tree. + /// + /// This method assumes that the provided node is a non-empty value, and that there is no node + /// at the specified index. + pub fn insert_leaf_node(&mut self, index: LeafNodeIndex, mut node: RpoDigest) -> RpoDigest { + debug_assert!(!is_empty_root(&node)); + debug_assert_eq!(self.nodes.get(&index), None); + + // mark the node as the leaf + if index.depth() == MAX_DEPTH { + self.bottom_leaves.insert(index.value()); + } else { + self.upper_leaves.insert(index.into()); + }; + + // insert the node and update the path from the node to the root + let mut index: NodeIndex = index.into(); + for _ in 0..index.depth() { + self.nodes.insert(index, node); + let sibling = self.get_node_unchecked(&index.sibling()); + node = Rpo256::merge(&index.build_node(node, sibling)); + index.move_up(); + } + + // update the root + self.nodes.insert(NodeIndex::root(), node); + node + } + + /// Updates the node at the specified index with the specified node value; recomputes and + /// returns the new root of the Tiered Sparse Merkle tree. + /// + /// This method can accept `node` as either an empty or a non-empty value. + pub fn update_leaf_node(&mut self, index: LeafNodeIndex, mut node: RpoDigest) -> RpoDigest { + debug_assert!(self.is_non_empty_leaf(&index)); + + // if the value we are updating the node to is a root of an empty tree, clear the leaf + // flag for this node + if node == EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize] { + if index.depth() == MAX_DEPTH { + self.bottom_leaves.remove(&index.value()); + } else { + self.upper_leaves.remove(&index); + } + } else { + debug_assert!(!is_empty_root(&node)); + } + + // update the path from the node to the root + let mut index: NodeIndex = index.into(); + for _ in 0..index.depth() { + if node == EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize] { + self.nodes.remove(&index); + } else { + self.nodes.insert(index, node); + } + + let sibling = self.get_node_unchecked(&index.sibling()); + node = Rpo256::merge(&index.build_node(node, sibling)); + index.move_up(); + } + + // update the root + self.nodes.insert(NodeIndex::root(), node); + node + } + + /// Replaces the leaf node at the specified index with a root of an empty subtree; recomputes + /// and returns the new root of the Tiered Sparse Merkle tree. + pub fn clear_leaf_node(&mut self, index: LeafNodeIndex) -> RpoDigest { + debug_assert!(self.is_non_empty_leaf(&index)); + let node = EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize]; + self.update_leaf_node(index, node) + } + + // HELPER METHODS + // -------------------------------------------------------------------------------------------- + + /// Returns true if the node at the specified index is a leaf node. + fn is_non_empty_leaf(&self, index: &LeafNodeIndex) -> bool { + if index.depth() == MAX_DEPTH { + self.bottom_leaves.contains(&index.value()) + } else { + self.upper_leaves.contains(index) + } + } + + /// Returns true if the node at the specified index is an internal node - i.e., there is + /// no leaf at that node and the node does not belong to the bottom tier. + fn is_internal_node(&self, index: &NodeIndex) -> bool { + if index.depth() == MAX_DEPTH { + false + } else { + !self.upper_leaves.contains(index) + } + } + + /// Checks if the specified index is valid in the context of this Merkle tree. + /// + /// # Errors + /// Returns an error if: + /// - The specified index depth is 0 or greater than 64. + /// - The node for the specified index does not exists in the Merkle tree. This is possible + /// when an ancestors of the specified index is a leaf node. + fn validate_node_access(&self, index: NodeIndex) -> Result<(), MerkleError> { + if index.is_root() { + return Err(MerkleError::DepthTooSmall(index.depth())); + } else if index.depth() > MAX_DEPTH { + return Err(MerkleError::DepthTooBig(index.depth() as u64)); + } else { + // make sure that there are no leaf nodes in the ancestors of the index; since leaf + // nodes can live at specific depth, we just need to check these depths. + let tier = ((index.depth() - 1) / TIER_SIZE) as usize; + let mut tier_index = index; + for &depth in TIER_DEPTHS[..tier].iter().rev() { + tier_index.move_up_to(depth); + if self.upper_leaves.contains(&tier_index) { + return Err(MerkleError::NodeNotInSet(index)); + } + } + } + + Ok(()) + } + + /// Returns a node at the specified index. If the node does not exist at this index, a root + /// for an empty subtree at the index's depth is returned. + /// + /// Unlike [NodeStore::get_node()] this does not perform any checks to verify that the + /// returned node is valid in the context of this tree. + fn get_node_unchecked(&self, index: &NodeIndex) -> RpoDigest { + match self.nodes.get(index) { + Some(node) => *node, + None => EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[index.depth() as usize], + } + } + + /// Removes a sequence of nodes starting at the specified index and traversing the + /// tree up to the specified depth. The node at the `end_depth` is also removed. + /// + /// This method does not update any other nodes and does not recompute the tree root. + fn remove_branch(&mut self, index: LeafNodeIndex, end_depth: u8) { + let mut index: NodeIndex = index.into(); + assert!(index.depth() > end_depth); + for _ in 0..(index.depth() - end_depth + 1) { + self.nodes.remove(&index); + index.move_up() + } + } +} + +// HELPER FUNCTIONS +// ================================================================================================ + +/// Returns true if the specified node is a root of an empty tree or an empty value ([ZERO; 4]). +fn is_empty_root(node: &RpoDigest) -> bool { + EmptySubtreeRoots::empty_hashes(MAX_DEPTH).contains(node) +} diff --git a/src/merkle/tiered_smt/tests.rs b/src/merkle/tiered_smt/tests.rs index fbf26ea..e459c90 100644 --- a/src/merkle/tiered_smt/tests.rs +++ b/src/merkle/tiered_smt/tests.rs @@ -1,9 +1,11 @@ use super::{ super::{super::ONE, Felt, MerkleStore, WORD_SIZE, ZERO}, - get_remaining_path, EmptySubtreeRoots, InnerNodeInfo, NodeIndex, Rpo256, RpoDigest, TieredSmt, - Vec, Word, + EmptySubtreeRoots, InnerNodeInfo, NodeIndex, Rpo256, RpoDigest, TieredSmt, Vec, Word, }; +// INSERTION TESTS +// ================================================================================================ + #[test] fn tsmt_insert_one() { let mut smt = TieredSmt::default(); @@ -217,6 +219,9 @@ fn tsmt_insert_three() { actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); } +// UPDATE TESTS +// ================================================================================================ + #[test] fn tsmt_update() { let mut smt = TieredSmt::default(); @@ -252,6 +257,209 @@ fn tsmt_update() { actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); } +// DELETION TESTS +// ================================================================================================ + +#[test] +fn tsmt_delete_16() { + let mut smt = TieredSmt::default(); + + // --- insert a value into the tree --------------------------------------- + let smt0 = smt.clone(); + let raw_a = 0b_01010101_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); + let value_a = [ONE, ONE, ONE, ONE]; + smt.insert(key_a, value_a); + + // --- insert another value into the tree --------------------------------- + let smt1 = smt.clone(); + let raw_b = 0b_01011111_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let value_b = [ONE, ONE, ONE, ZERO]; + smt.insert(key_b, value_b); + + // --- delete the last inserted value ------------------------------------- + assert_eq!(smt.insert(key_b, [ZERO; 4]), value_b); + assert_eq!(smt, smt1); + + // --- delete the first inserted value ------------------------------------ + assert_eq!(smt.insert(key_a, [ZERO; 4]), value_a); + assert_eq!(smt, smt0); +} + +#[test] +fn tsmt_delete_32() { + let mut smt = TieredSmt::default(); + + // --- insert a value into the tree --------------------------------------- + let smt0 = smt.clone(); + let raw_a = 0b_01010101_01101100_01111111_11111111_10010110_10010011_11100000_00000000_u64; + let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); + let value_a = [ONE, ONE, ONE, ONE]; + smt.insert(key_a, value_a); + + // --- insert another with the same 16-bit prefix into the tree ----------- + let smt1 = smt.clone(); + let raw_b = 0b_01010101_01101100_00111111_11111111_10010110_10010011_11100000_00000000_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let value_b = [ONE, ONE, ONE, ZERO]; + smt.insert(key_b, value_b); + + // --- insert the 3rd value with the same 16-bit prefix into the tree ----- + let smt2 = smt.clone(); + let raw_c = 0b_01010101_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); + let value_c = [ONE, ONE, ZERO, ZERO]; + smt.insert(key_c, value_c); + + // --- delete the last inserted value ------------------------------------- + assert_eq!(smt.insert(key_c, [ZERO; 4]), value_c); + assert_eq!(smt, smt2); + + // --- delete the last inserted value ------------------------------------- + assert_eq!(smt.insert(key_b, [ZERO; 4]), value_b); + assert_eq!(smt, smt1); + + // --- delete the first inserted value ------------------------------------ + assert_eq!(smt.insert(key_a, [ZERO; 4]), value_a); + assert_eq!(smt, smt0); +} + +#[test] +fn tsmt_delete_48_same_32_bit_prefix() { + let mut smt = TieredSmt::default(); + + // test the case when all values share the same 32-bit prefix + + // --- insert a value into the tree --------------------------------------- + let smt0 = smt.clone(); + let raw_a = 0b_01010101_01010101_11111111_11111111_10010110_10010011_11100000_00000000_u64; + let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); + let value_a = [ONE, ONE, ONE, ONE]; + smt.insert(key_a, value_a); + + // --- insert another with the same 32-bit prefix into the tree ----------- + let smt1 = smt.clone(); + let raw_b = 0b_01010101_01010101_11111111_11111111_11010110_10010011_11100000_00000000_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let value_b = [ONE, ONE, ONE, ZERO]; + smt.insert(key_b, value_b); + + // --- insert the 3rd value with the same 32-bit prefix into the tree ----- + let smt2 = smt.clone(); + let raw_c = 0b_01010101_01010101_11111111_11111111_11110110_10010011_11100000_00000000_u64; + let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); + let value_c = [ONE, ONE, ZERO, ZERO]; + smt.insert(key_c, value_c); + + // --- delete the last inserted value ------------------------------------- + assert_eq!(smt.insert(key_c, [ZERO; 4]), value_c); + assert_eq!(smt, smt2); + + // --- delete the last inserted value ------------------------------------- + assert_eq!(smt.insert(key_b, [ZERO; 4]), value_b); + assert_eq!(smt, smt1); + + // --- delete the first inserted value ------------------------------------ + assert_eq!(smt.insert(key_a, [ZERO; 4]), value_a); + assert_eq!(smt, smt0); +} + +#[test] +fn tsmt_delete_48_mixed_prefix() { + let mut smt = TieredSmt::default(); + + // test the case when some values share a 32-bit prefix and others share a 16-bit prefix + + // --- insert a value into the tree --------------------------------------- + let smt0 = smt.clone(); + let raw_a = 0b_01010101_01010101_11111111_11111111_10010110_10010011_11100000_00000000_u64; + let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); + let value_a = [ONE, ONE, ONE, ONE]; + smt.insert(key_a, value_a); + + // --- insert another with the same 16-bit prefix into the tree ----------- + let smt1 = smt.clone(); + let raw_b = 0b_01010101_01010101_01111111_11111111_10010110_10010011_11100000_00000000_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let value_b = [ONE, ONE, ONE, ZERO]; + smt.insert(key_b, value_b); + + // --- insert a value with the same 32-bit prefix as the first value ----- + let smt2 = smt.clone(); + let raw_c = 0b_01010101_01010101_11111111_11111111_11010110_10010011_11100000_00000000_u64; + let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); + let value_c = [ONE, ONE, ZERO, ZERO]; + smt.insert(key_c, value_c); + + // --- insert another value with the same 32-bit prefix as the first value + let smt3 = smt.clone(); + let raw_d = 0b_01010101_01010101_11111111_11111111_11110110_10010011_11100000_00000000_u64; + let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_d)]); + let value_d = [ONE, ZERO, ZERO, ZERO]; + smt.insert(key_d, value_d); + + // --- delete the inserted values one-by-one ------------------------------ + assert_eq!(smt.insert(key_d, [ZERO; 4]), value_d); + assert_eq!(smt, smt3); + + assert_eq!(smt.insert(key_c, [ZERO; 4]), value_c); + assert_eq!(smt, smt2); + + assert_eq!(smt.insert(key_b, [ZERO; 4]), value_b); + assert_eq!(smt, smt1); + + assert_eq!(smt.insert(key_a, [ZERO; 4]), value_a); + assert_eq!(smt, smt0); +} + +#[test] +fn tsmt_delete_64() { + let mut smt = TieredSmt::default(); + + // test the case when all values share the same 48-bit prefix + + // --- insert a value into the tree --------------------------------------- + let smt0 = smt.clone(); + let raw_a = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000000_u64; + let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); + let value_a = [ONE, ONE, ONE, ONE]; + smt.insert(key_a, value_a); + + // --- insert a value with the same 48-bit prefix into the tree ----------- + let smt1 = smt.clone(); + let raw_b = 0b_01010101_01010101_11111111_11111111_10110101_10101010_10111100_00000000_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let value_b = [ONE, ONE, ONE, ZERO]; + smt.insert(key_b, value_b); + + // --- insert a value with the same 32-bit prefix into the tree ----------- + let smt2 = smt.clone(); + let raw_c = 0b_01010101_01010101_11111111_11111111_11111101_10101010_10111100_00000000_u64; + let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); + let value_c = [ONE, ONE, ZERO, ZERO]; + smt.insert(key_c, value_c); + + let smt3 = smt.clone(); + let raw_d = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000000_u64; + let key_d = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_d)]); + let value_d = [ONE, ZERO, ZERO, ZERO]; + smt.insert(key_d, value_d); + + // --- delete the last inserted value ------------------------------------- + assert_eq!(smt.insert(key_d, [ZERO; 4]), value_d); + assert_eq!(smt, smt3); + + assert_eq!(smt.insert(key_c, [ZERO; 4]), value_c); + assert_eq!(smt, smt2); + + assert_eq!(smt.insert(key_b, [ZERO; 4]), value_b); + assert_eq!(smt, smt1); + + assert_eq!(smt.insert(key_a, [ZERO; 4]), value_a); + assert_eq!(smt, smt0); +} + // BOTTOM TIER TESTS // ================================================================================================ @@ -301,9 +509,26 @@ fn tsmt_bottom_tier() { actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node))); // make sure leaves are returned correctly - let mut leaves = smt.bottom_leaves(); + let smt_clone = smt.clone(); + let mut leaves = smt_clone.bottom_leaves(); assert_eq!(leaves.next(), Some((leaf_node, vec![(key_b, val_b), (key_a, val_a)]))); assert_eq!(leaves.next(), None); + + // --- update a leaf at the bottom tier ------------------------------------------------------- + + let val_a2 = [Felt::new(3); WORD_SIZE]; + assert_eq!(smt.insert(key_a, val_a2), val_a); + + let leaf_node = build_bottom_leaf_node(&[key_b, key_a], &[val_b, val_a2]); + store.set_node(tree_root, index, leaf_node).unwrap(); + + 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))); + + let mut leaves = smt.bottom_leaves(); + assert_eq!(leaves.next(), Some((leaf_node, vec![(key_b, val_b), (key_a, val_a2)]))); + assert_eq!(leaves.next(), None); } #[test] @@ -411,8 +636,7 @@ fn get_init_root() -> RpoDigest { } fn build_leaf_node(key: RpoDigest, value: Word, depth: u8) -> RpoDigest { - let remaining_path = get_remaining_path(key, depth as u32); - Rpo256::merge_in_domain(&[remaining_path, value.into()], depth.into()) + Rpo256::merge_in_domain(&[key, value.into()], depth.into()) } fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest { @@ -420,9 +644,7 @@ fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest { let mut elements = Vec::with_capacity(keys.len()); for (key, val) in keys.iter().zip(values.iter()) { - let mut key = Word::from(key); - key[3] = ZERO; - elements.extend_from_slice(&key); + elements.extend_from_slice(key.as_elements()); elements.extend_from_slice(val.as_slice()); } diff --git a/src/merkle/tiered_smt/values.rs b/src/merkle/tiered_smt/values.rs new file mode 100644 index 0000000..ec2a465 --- /dev/null +++ b/src/merkle/tiered_smt/values.rs @@ -0,0 +1,585 @@ +use super::{get_key_prefix, BTreeMap, LeafNodeIndex, RpoDigest, StarkField, Vec, Word}; +use crate::utils::vec; +use core::{ + cmp::{Ord, Ordering}, + ops::RangeBounds, +}; +use winter_utils::collections::btree_map::Entry; + +// CONSTANTS +// ================================================================================================ + +/// Depths at which leaves can exist in a tiered SMT. +const TIER_DEPTHS: [u8; 4] = super::TieredSmt::TIER_DEPTHS; + +/// Maximum node depth. This is also the bottom tier of the tree. +const MAX_DEPTH: u8 = super::TieredSmt::MAX_DEPTH; + +// VALUE STORE +// ================================================================================================ +/// A store for key-value pairs for a Tiered Sparse Merkle tree. +/// +/// The store is organized in a [BTreeMap] where keys are 64 most significant bits of a key, and +/// the values are the corresponding key-value pairs (or a list of key-value pairs if more that +/// a single key-value pair shares the same 64-bit prefix). +/// +/// The store supports lookup by the full key (i.e. [RpoDigest]) as well as by the 64-bit key +/// prefix. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub struct ValueStore { + values: BTreeMap, +} + +impl ValueStore { + // PUBLIC ACCESSORS + // -------------------------------------------------------------------------------------------- + + /// Returns a reference to the value stored under the specified key, or None if there is no + /// value associated with the specified key. + pub fn get(&self, key: &RpoDigest) -> Option<&Word> { + let prefix = get_key_prefix(key); + self.values.get(&prefix).and_then(|entry| entry.get(key)) + } + + /// Returns the first key-value pair such that the key prefix is greater than or equal to the + /// specified prefix. + pub fn get_first(&self, prefix: u64) -> Option<&(RpoDigest, Word)> { + self.range(prefix..).next() + } + + /// Returns the first key-value pair such that the key prefix is greater than or equal to the + /// specified prefix and the key value is not equal to the exclude_key value. + pub fn get_first_filtered( + &self, + prefix: u64, + exclude_key: &RpoDigest, + ) -> Option<&(RpoDigest, Word)> { + self.range(prefix..).find(|(key, _)| key != exclude_key) + } + + /// Returns a vector with key-value pairs for all keys with the specified 64-bit prefix, or + /// None if no keys with the specified prefix are present in this store. + pub fn get_all(&self, prefix: u64) -> Option> { + self.values.get(&prefix).map(|entry| match entry { + StoreEntry::Single(kv_pair) => vec![*kv_pair], + StoreEntry::List(kv_pairs) => kv_pairs.clone(), + }) + } + + /// Returns information about a sibling of a leaf node with the specified index, but only if + /// this is the only sibling the leaf has in some subtree starting at the first tier. + /// + /// For example, if `index` is an index at depth 32, and there is a leaf node at depth 32 with + /// the same root at depth 16 as `index`, we say that this leaf is a lone sibling. + /// + /// The returned tuple contains: they key-value pair of the sibling as well as the index of + /// the node for the root of the common subtree in which both nodes are leaves. + /// + /// This method assumes that the key-value pair for the specified index has already been + /// removed from the store. + pub fn get_lone_sibling( + &self, + index: LeafNodeIndex, + ) -> Option<(&RpoDigest, &Word, LeafNodeIndex)> { + // iterate over tiers from top to bottom, looking at the tiers which are strictly above + // the depth of the index. This implies that only tiers at depth 32 and 48 will be + // considered. For each tier, check if the parent of the index at the higher tier + // contains a single node. The fist tier (depth 16) is excluded because we cannot move + // nodes at depth 16 to a higher tier. This implies that nodes at the first tier will + // never have "lone siblings". + for &tier_depth in TIER_DEPTHS.iter().filter(|&t| index.depth() > *t) { + // compute the index of the root at a higher tier + let mut parent_index = index; + parent_index.move_up_to(tier_depth); + + // find the lone sibling, if any; we need to handle the "last node" at a given tier + // separately specify the bounds for the search correctly. + let start_prefix = parent_index.value() << (MAX_DEPTH - tier_depth); + let sibling = if start_prefix.leading_ones() as u8 == tier_depth { + let mut iter = self.range(start_prefix..); + iter.next().filter(|_| iter.next().is_none()) + } else { + let end_prefix = (parent_index.value() + 1) << (MAX_DEPTH - tier_depth); + let mut iter = self.range(start_prefix..end_prefix); + iter.next().filter(|_| iter.next().is_none()) + }; + + if let Some((key, value)) = sibling { + return Some((key, value, parent_index)); + } + } + + None + } + + /// Returns an iterator over all key-value pairs in this store. + pub fn iter(&self) -> impl Iterator { + self.values.iter().flat_map(|(_, entry)| entry.iter()) + } + + // STATE MUTATORS + // -------------------------------------------------------------------------------------------- + + /// Inserts the specified key-value pair into this store and returns the value previously + /// associated with the specified key. + /// + /// If no value was previously associated with the specified key, None is returned. + pub fn insert(&mut self, key: RpoDigest, value: Word) -> Option { + let prefix = get_key_prefix(&key); + match self.values.entry(prefix) { + Entry::Occupied(mut entry) => entry.get_mut().insert(key, value), + Entry::Vacant(entry) => { + entry.insert(StoreEntry::new(key, value)); + None + } + } + } + + /// Removes the key-value pair for the specified key from this store and returns the value + /// associated with this key. + /// + /// If no value was associated with the specified key, None is returned. + pub fn remove(&mut self, key: &RpoDigest) -> Option { + let prefix = get_key_prefix(key); + match self.values.entry(prefix) { + Entry::Occupied(mut entry) => { + let (value, remove_entry) = entry.get_mut().remove(key); + if remove_entry { + entry.remove_entry(); + } + value + } + Entry::Vacant(_) => None, + } + } + + // HELPER METHODS + // -------------------------------------------------------------------------------------------- + + /// Returns an iterator over all key-value pairs contained in this store such that the most + /// significant 64 bits of the key lay within the specified bounds. + /// + /// The order of iteration is from the smallest to the largest key. + fn range>(&self, bounds: R) -> impl Iterator { + self.values.range(bounds).flat_map(|(_, entry)| entry.iter()) + } +} + +// VALUE NODE +// ================================================================================================ + +/// An entry in the [ValueStore]. +/// +/// An entry can contain either a single key-value pair or a vector of key-value pairs sorted by +/// key. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum StoreEntry { + Single((RpoDigest, Word)), + List(Vec<(RpoDigest, Word)>), +} + +impl StoreEntry { + // CONSTRUCTOR + // -------------------------------------------------------------------------------------------- + /// Returns a new [StoreEntry] instantiated with a single key-value pair. + pub fn new(key: RpoDigest, value: Word) -> Self { + Self::Single((key, value)) + } + + // PUBLIC ACCESSORS + // -------------------------------------------------------------------------------------------- + + /// Returns the value associated with the specified key, or None if this entry does not contain + /// a value associated with the specified key. + pub fn get(&self, key: &RpoDigest) -> Option<&Word> { + match self { + StoreEntry::Single(kv_pair) => { + if kv_pair.0 == *key { + Some(&kv_pair.1) + } else { + None + } + } + StoreEntry::List(kv_pairs) => { + match kv_pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, key)) { + Ok(pos) => Some(&kv_pairs[pos].1), + Err(_) => None, + } + } + } + } + + /// Returns an iterator over all key-value pairs in this entry. + pub fn iter(&self) -> impl Iterator { + EntryIterator { + entry: self, + pos: 0, + } + } + + // STATE MUTATORS + // -------------------------------------------------------------------------------------------- + + /// Inserts the specified key-value pair into this entry and returns the value previously + /// associated with the specified key, or None if no value was associated with the specified + /// key. + /// + /// If a new key is inserted, this will also transform a `SingleEntry` into a `ListEntry`. + pub fn insert(&mut self, key: RpoDigest, value: Word) -> Option { + match self { + StoreEntry::Single(kv_pair) => { + // if the key is already in this entry, update the value and return + if kv_pair.0 == key { + let old_value = kv_pair.1; + kv_pair.1 = value; + return Some(old_value); + } + + // transform the entry into a list entry, and make sure the key-value pairs + // are sorted by key + let mut pairs = vec![*kv_pair, (key, value)]; + pairs.sort_by(|a, b| cmp_digests(&a.0, &b.0)); + + *self = StoreEntry::List(pairs); + None + } + StoreEntry::List(pairs) => { + match pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, &key)) { + Ok(pos) => { + let old_value = pairs[pos].1; + pairs[pos].1 = value; + Some(old_value) + } + Err(pos) => { + pairs.insert(pos, (key, value)); + None + } + } + } + } + } + + /// Removes the key-value pair with the specified key from this entry, and returns the value + /// of the removed pair. If the entry did not contain a key-value pair for the specified key, + /// None is returned. + /// + /// If the last last key-value pair was removed from the entry, the second tuple value will + /// be set to true. + pub fn remove(&mut self, key: &RpoDigest) -> (Option, bool) { + match self { + StoreEntry::Single(kv_pair) => { + if kv_pair.0 == *key { + (Some(kv_pair.1), true) + } else { + (None, false) + } + } + StoreEntry::List(kv_pairs) => { + match kv_pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, key)) { + Ok(pos) => { + let kv_pair = kv_pairs.remove(pos); + if kv_pairs.len() == 1 { + *self = StoreEntry::Single(kv_pairs[0]); + } + (Some(kv_pair.1), false) + } + Err(_) => (None, false), + } + } + } + } +} + +/// A custom iterator over key-value pairs of a [StoreEntry]. +/// +/// For a `SingleEntry` this returns only one value, but for `ListEntry`, this iterates over the +/// entire list of key-value pairs. +pub struct EntryIterator<'a> { + entry: &'a StoreEntry, + pos: usize, +} + +impl<'a> Iterator for EntryIterator<'a> { + type Item = &'a (RpoDigest, Word); + + fn next(&mut self) -> Option { + match self.entry { + StoreEntry::Single(kv_pair) => { + if self.pos == 0 { + self.pos = 1; + Some(kv_pair) + } else { + None + } + } + StoreEntry::List(kv_pairs) => { + if self.pos >= kv_pairs.len() { + None + } else { + let kv_pair = &kv_pairs[self.pos]; + self.pos += 1; + Some(kv_pair) + } + } + } + } +} + +// HELPER FUNCTIONS +// ================================================================================================ + +/// Compares two digests element-by-element using their integer representations starting with the +/// most significant element. +fn cmp_digests(d1: &RpoDigest, d2: &RpoDigest) -> Ordering { + let d1 = Word::from(d1); + let d2 = Word::from(d2); + + for (v1, v2) in d1.iter().zip(d2.iter()).rev() { + let v1 = v1.as_int(); + let v2 = v2.as_int(); + if v1 != v2 { + return v1.cmp(&v2); + } + } + + Ordering::Equal +} + +// TESTS +// ================================================================================================ + +#[cfg(test)] +mod tests { + use super::{LeafNodeIndex, RpoDigest, StoreEntry, ValueStore}; + use crate::{Felt, ONE, WORD_SIZE, ZERO}; + + #[test] + fn test_insert() { + let mut store = ValueStore::default(); + + // insert the first key-value pair into the store + let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]); + let value_a = [ONE; WORD_SIZE]; + + assert!(store.insert(key_a, value_a).is_none()); + assert_eq!(store.values.len(), 1); + + let entry = store.values.get(&raw_a).unwrap(); + let expected_entry = StoreEntry::Single((key_a, value_a)); + assert_eq!(entry, &expected_entry); + + // insert a key-value pair with a different key into the store; since the keys are + // different, another entry is added to the values map + let raw_b = 0b_11111110_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let value_b = [ONE, ZERO, ONE, ZERO]; + + assert!(store.insert(key_b, value_b).is_none()); + assert_eq!(store.values.len(), 2); + + let entry1 = store.values.get(&raw_a).unwrap(); + let expected_entry1 = StoreEntry::Single((key_a, value_a)); + assert_eq!(entry1, &expected_entry1); + + let entry2 = store.values.get(&raw_b).unwrap(); + let expected_entry2 = StoreEntry::Single((key_b, value_b)); + assert_eq!(entry2, &expected_entry2); + + // insert a key-value pair with the same 64-bit key prefix as the first key; this should + // transform the first entry into a List entry + let key_c = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]); + let value_c = [ONE, ONE, ZERO, ZERO]; + + assert!(store.insert(key_c, value_c).is_none()); + assert_eq!(store.values.len(), 2); + + let entry1 = store.values.get(&raw_a).unwrap(); + let expected_entry1 = StoreEntry::List(vec![(key_c, value_c), (key_a, value_a)]); + assert_eq!(entry1, &expected_entry1); + + let entry2 = store.values.get(&raw_b).unwrap(); + let expected_entry2 = StoreEntry::Single((key_b, value_b)); + assert_eq!(entry2, &expected_entry2); + + // replace values for keys a and b + let value_a2 = [ONE, ONE, ONE, ZERO]; + let value_b2 = [ZERO, ZERO, ZERO, ONE]; + + assert_eq!(store.insert(key_a, value_a2), Some(value_a)); + assert_eq!(store.values.len(), 2); + + assert_eq!(store.insert(key_b, value_b2), Some(value_b)); + assert_eq!(store.values.len(), 2); + + let entry1 = store.values.get(&raw_a).unwrap(); + let expected_entry1 = StoreEntry::List(vec![(key_c, value_c), (key_a, value_a2)]); + assert_eq!(entry1, &expected_entry1); + + let entry2 = store.values.get(&raw_b).unwrap(); + let expected_entry2 = StoreEntry::Single((key_b, value_b2)); + assert_eq!(entry2, &expected_entry2); + + // insert one more key-value pair with the same 64-bit key-prefix as the first key + let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); + let value_d = [ZERO, ONE, ZERO, ZERO]; + + assert!(store.insert(key_d, value_d).is_none()); + assert_eq!(store.values.len(), 2); + + let entry1 = store.values.get(&raw_a).unwrap(); + let expected_entry1 = + StoreEntry::List(vec![(key_c, value_c), (key_a, value_a2), (key_d, value_d)]); + assert_eq!(entry1, &expected_entry1); + + let entry2 = store.values.get(&raw_b).unwrap(); + let expected_entry2 = StoreEntry::Single((key_b, value_b2)); + assert_eq!(entry2, &expected_entry2); + } + + #[test] + fn test_remove() { + // populate the value store + let mut store = ValueStore::default(); + + let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]); + let value_a = [ONE; WORD_SIZE]; + store.insert(key_a, value_a); + + let raw_b = 0b_11111110_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let value_b = [ONE, ZERO, ONE, ZERO]; + store.insert(key_b, value_b); + + let key_c = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]); + let value_c = [ONE, ONE, ZERO, ZERO]; + store.insert(key_c, value_c); + + let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); + let value_d = [ZERO, ONE, ZERO, ZERO]; + store.insert(key_d, value_d); + + assert_eq!(store.values.len(), 2); + + let entry1 = store.values.get(&raw_a).unwrap(); + let expected_entry1 = + StoreEntry::List(vec![(key_c, value_c), (key_a, value_a), (key_d, value_d)]); + assert_eq!(entry1, &expected_entry1); + + let entry2 = store.values.get(&raw_b).unwrap(); + let expected_entry2 = StoreEntry::Single((key_b, value_b)); + assert_eq!(entry2, &expected_entry2); + + // remove non-existent keys + let key_e = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_a)]); + assert!(store.remove(&key_e).is_none()); + + let raw_f = 0b_11111110_11111111_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_f = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_f)]); + assert!(store.remove(&key_f).is_none()); + + // remove keys from the list entry + assert_eq!(store.remove(&key_c).unwrap(), value_c); + let entry1 = store.values.get(&raw_a).unwrap(); + let expected_entry1 = StoreEntry::List(vec![(key_a, value_a), (key_d, value_d)]); + assert_eq!(entry1, &expected_entry1); + + assert_eq!(store.remove(&key_a).unwrap(), value_a); + let entry1 = store.values.get(&raw_a).unwrap(); + let expected_entry1 = StoreEntry::Single((key_d, value_d)); + assert_eq!(entry1, &expected_entry1); + + assert_eq!(store.remove(&key_d).unwrap(), value_d); + assert!(store.values.get(&raw_a).is_none()); + assert_eq!(store.values.len(), 1); + + // remove a key from a single entry + assert_eq!(store.remove(&key_b).unwrap(), value_b); + assert!(store.values.get(&raw_b).is_none()); + assert_eq!(store.values.len(), 0); + } + + #[test] + fn test_range() { + // populate the value store + let mut store = ValueStore::default(); + + let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]); + let value_a = [ONE; WORD_SIZE]; + store.insert(key_a, value_a); + + let raw_b = 0b_11111110_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let value_b = [ONE, ZERO, ONE, ZERO]; + store.insert(key_b, value_b); + + let key_c = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]); + let value_c = [ONE, ONE, ZERO, ZERO]; + store.insert(key_c, value_c); + + let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); + let value_d = [ZERO, ONE, ZERO, ZERO]; + store.insert(key_d, value_d); + + let raw_e = 0b_10101000_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_e = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_e)]); + let value_e = [ZERO, ZERO, ZERO, ONE]; + store.insert(key_e, value_e); + + // check the entire range + let mut iter = store.range(..u64::MAX); + assert_eq!(iter.next(), Some(&(key_e, value_e))); + assert_eq!(iter.next(), Some(&(key_c, value_c))); + assert_eq!(iter.next(), Some(&(key_a, value_a))); + assert_eq!(iter.next(), Some(&(key_d, value_d))); + assert_eq!(iter.next(), Some(&(key_b, value_b))); + assert_eq!(iter.next(), None); + + // check all but e + let mut iter = store.range(raw_a..u64::MAX); + assert_eq!(iter.next(), Some(&(key_c, value_c))); + assert_eq!(iter.next(), Some(&(key_a, value_a))); + assert_eq!(iter.next(), Some(&(key_d, value_d))); + assert_eq!(iter.next(), Some(&(key_b, value_b))); + assert_eq!(iter.next(), None); + + // check all but e and b + let mut iter = store.range(raw_a..raw_b); + assert_eq!(iter.next(), Some(&(key_c, value_c))); + assert_eq!(iter.next(), Some(&(key_a, value_a))); + assert_eq!(iter.next(), Some(&(key_d, value_d))); + assert_eq!(iter.next(), None); + } + + #[test] + fn test_get_lone_sibling() { + // populate the value store + let mut store = ValueStore::default(); + + let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]); + let value_a = [ONE; WORD_SIZE]; + store.insert(key_a, value_a); + + let raw_b = 0b_11111111_11111111_00011111_11111111_10010110_10010011_11100000_00000000_u64; + let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]); + let value_b = [ONE, ZERO, ONE, ZERO]; + store.insert(key_b, value_b); + + // check sibling node for `a` + let index = LeafNodeIndex::make(32, 0b_10101010_10101010_00011111_11111110); + let parent_index = LeafNodeIndex::make(16, 0b_10101010_10101010); + assert_eq!(store.get_lone_sibling(index), Some((&key_a, &value_a, parent_index))); + + // check sibling node for `b` + let index = LeafNodeIndex::make(32, 0b_11111111_11111111_00011111_11111111); + let parent_index = LeafNodeIndex::make(16, 0b_11111111_11111111); + assert_eq!(store.get_lone_sibling(index), Some((&key_b, &value_b, parent_index))); + + // check some other sibling for some other index + let index = LeafNodeIndex::make(32, 0b_11101010_10101010); + assert_eq!(store.get_lone_sibling(index), None); + } +}