Browse Source
Merge pull request #173 from 0xPolygonMiden/bobbin-tsmt-refactor
Implement value clearing in TSMTal-gkr-basic-workflow
Augusto Hack
1 year ago
committed by
GitHub
GitHub
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 1443 additions and 290 deletions
Split View
Diff Options
-
+1 -2src/merkle/partial_mt/mod.rs
-
+253 -280src/merkle/tiered_smt/mod.rs
-
+374 -0src/merkle/tiered_smt/nodes.rs
-
+230 -8src/merkle/tiered_smt/tests.rs
-
+585 -0src/merkle/tiered_smt/values.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<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)
|
|||
}
|
|||
@ -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);
|
|||
}
|
|||
}
|
|||