qyriad/miden-crypto
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #173 from 0xPolygonMiden/bobbin-tsmt-refactor

Browse source 
Implement value clearing in TSMT
This commit is contained in:
Augusto Hack 2023-08-03 11:28:54 +02:00 • committed by GitHub
commit 92bb3ac462
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 1445 additions and 292 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
src/merkle/partial_mt/mod.rs
View file

@ -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");

537
src/merkle/tiered_smt/mod.rs
View file

@ -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<NodeIndex, RpoDigest>,
upper_leaves: BTreeMap<NodeIndex, RpoDigest>, // node_index |-> key map
bottom_leaves: BTreeMap<u64, BottomLeaf>, // leaves of depth 64
values: BTreeMap<RpoDigest, Word>,
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<RpoDigest, MerkleError> {
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<MerklePath, MerkleError> {
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<MerklePath, MerkleError> {
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<Item = &(RpoDigest, Word)> {
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<Item = InnerNodeInfo> + '_ {
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<Item = (RpoDigest, RpoDigest, Word)> + '_ {
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<Item = (RpoDigest, Vec<(RpoDigest, Word)>)> + '_ {
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));
/// 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,
};
// 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;
};
}
// 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 {
// 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));
}
}
// 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);
}
Ok(())
old_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.
/// Builds and returns a leaf node value for the node located as the specified index.
///
/// 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],
}
}
/// 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<RpoDigest>) {
// 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 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)
}
/// 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) {
/// 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<NodeIndex> for LeafNodeIndex {
fn from(value: NodeIndex) -> Self {
Self::new(value)
}
}
impl From<LeafNodeIndex> 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,
}
/// 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())
}
/// 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
// ================================================================================================
/// 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)
}

374
src/merkle/tiered_smt/nodes.rs Normal file
View file

@ -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<NodeIndex, RpoDigest>,
upper_leaves: BTreeSet<NodeIndex>,
bottom_leaves: BTreeSet<u64>,
}
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<RpoDigest, MerkleError> {
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<MerklePath, MerkleError> {
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<Item = InnerNodeInfo> + '_ {
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<Item = (&NodeIndex, &RpoDigest)> {
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<Item = (&u64, &RpoDigest)> {
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)
}

238
src/merkle/tiered_smt/tests.rs
View file

@ -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::<Vec<_>>();
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());
}

585
src/merkle/tiered_smt/values.rs Normal file
View file

@ -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<u64, StoreEntry>,
}
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<Vec<(RpoDigest, Word)>> {
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<Item = &(RpoDigest, Word)> {
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<Word> {
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<Word> {
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<R: RangeBounds<u64>>(&self, bounds: R) -> impl Iterator<Item = &(RpoDigest, Word)> {
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<Item = &(RpoDigest, Word)> {
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<Word> {
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<Word>, 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<Self::Item> {
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);
}
}