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

feat: added handling of bottom tier to TieredSmt

Browse source
This commit is contained in:
Bobbin Threadbare 2023-05-14 02:29:42 -07:00
parent 51ce07cc34
commit b768eade4d
7 changed files with 654 additions and 217 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
src/hash/rpo/digest.rs
View file

@ -2,7 +2,7 @@ use super::{Digest, Felt, StarkField, DIGEST_SIZE, ZERO};
use crate::utils::{
string::String, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
};
use core::{cmp::Ordering, ops::Deref};
use core::{cmp::Ordering, fmt::Display, ops::Deref};
// DIGEST TRAIT IMPLEMENTATIONS
// ================================================================================================
@ -85,6 +85,28 @@ impl From<RpoDigest> for [Felt; DIGEST_SIZE] {
}
}
impl From<&RpoDigest> for [u64; DIGEST_SIZE] {
fn from(value: &RpoDigest) -> Self {
[
value.0[0].as_int(),
value.0[1].as_int(),
value.0[2].as_int(),
value.0[3].as_int(),
]
}
}
impl From<RpoDigest> for [u64; DIGEST_SIZE] {
fn from(value: RpoDigest) -> Self {
[
value.0[0].as_int(),
value.0[1].as_int(),
value.0[2].as_int(),
value.0[3].as_int(),
]
}
}
impl From<&RpoDigest> for [u8; 32] {
fn from(value: &RpoDigest) -> Self {
value.as_bytes()
@ -134,6 +156,15 @@ impl PartialOrd for RpoDigest {
}
}
impl Display for RpoDigest {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
for byte in self.as_bytes() {
write!(f, "{byte:02x}")?;
}
Ok(())
}
}
// TESTS
// ================================================================================================

28
src/merkle/index.rs
View file

@ -1,4 +1,5 @@
use super::{Felt, MerkleError, RpoDigest, StarkField};
use core::fmt::Display;
// NODE INDEX
// ================================================================================================
@ -40,6 +41,12 @@ impl NodeIndex {
}
}
/// Creates a new node index without checking its validity.
pub const fn new_unchecked(depth: u8, value: u64) -> Self {
debug_assert!((64 - value.leading_zeros()) <= depth as u32);
Self { depth, value }
}
/// Creates a new node index for testing purposes.
///
/// # Panics
@ -117,11 +124,26 @@ impl NodeIndex {
// STATE MUTATORS
// --------------------------------------------------------------------------------------------
/// Traverse one level towards the root, decrementing the depth by `1`.
pub fn move_up(&mut self) -> &mut Self {
/// Traverses one level towards the root, decrementing the depth by `1`.
pub fn move_up(&mut self) {
self.depth = self.depth.saturating_sub(1);
self.value >>= 1;
self
}
/// Traverses towards the root until the specified depth is reached.
///
/// Assumes that the specified depth is smaller than the current depth.
pub fn move_up_to(&mut self, depth: u8) {
debug_assert!(depth < self.depth);
let delta = self.depth.saturating_sub(depth);
self.depth = self.depth.saturating_sub(delta);
self.value >>= delta as u32;
}
}
impl Display for NodeIndex {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "depth={}, value={}", self.depth, self.value)
}
}

20
src/merkle/mod.rs
View file

@ -47,14 +47,15 @@ pub enum MerkleError {
ConflictingRoots(Vec<Word>),
DepthTooSmall(u8),
DepthTooBig(u64),
DuplicateValuesForKey(u64),
NodeNotInStore(Word, NodeIndex),
NumLeavesNotPowerOfTwo(usize),
DuplicateValuesForIndex(u64),
DuplicateValuesForKey(RpoDigest),
InvalidIndex { depth: u8, value: u64 },
InvalidDepth { expected: u8, provided: u8 },
InvalidPath(MerklePath),
InvalidNumEntries(usize, usize),
NodeNotInSet(u64),
NodeNotInSet(NodeIndex),
NodeNotInStore(Word, NodeIndex),
NumLeavesNotPowerOfTwo(usize),
RootNotInStore(Word),
}
@ -65,10 +66,8 @@ impl fmt::Display for MerkleError {
ConflictingRoots(roots) => write!(f, "the merkle paths roots do not match {roots:?}"),
DepthTooSmall(depth) => write!(f, "the provided depth {depth} is too small"),
DepthTooBig(depth) => write!(f, "the provided depth {depth} is too big"),
DuplicateValuesForIndex(key) => write!(f, "multiple values provided for key {key}"),
DuplicateValuesForKey(key) => write!(f, "multiple values provided for key {key}"),
NumLeavesNotPowerOfTwo(leaves) => {
write!(f, "the leaves count {leaves} is not a power of 2")
}
InvalidIndex{ depth, value} => write!(
f,
"the index value {value} is not valid for the depth {depth}"
@ -79,8 +78,11 @@ impl fmt::Display for MerkleError {
),
InvalidPath(_path) => write!(f, "the provided path is not valid"),
InvalidNumEntries(max, provided) => write!(f, "the provided number of entries is {provided}, but the maximum for the given depth is {max}"),
NodeNotInSet(index) => write!(f, "the node indexed by {index} is not in the set"),
NodeNotInStore(hash, index) => write!(f, "the node {:?} indexed by {} and depth {} is not in the store", hash, index.value(), index.depth(),),
NodeNotInSet(index) => write!(f, "the node with index ({index}) is not in the set"),
NodeNotInStore(hash, index) => write!(f, "the node {hash:?} with index ({index}) is not in the store"),
NumLeavesNotPowerOfTwo(leaves) => {
write!(f, "the leaves count {leaves} is not a power of 2")
}
RootNotInStore(root) => write!(f, "the root {:?} is not in the store", root),
}
}

11
src/merkle/path_set.rs
View file

@ -73,7 +73,7 @@ impl MerklePathSet {
let path_key = index.value() - parity;
self.paths
.get(&path_key)
.ok_or(MerkleError::NodeNotInSet(path_key))
.ok_or(MerkleError::NodeNotInSet(index))
.map(|path| path[parity as usize])
}
@ -104,11 +104,8 @@ impl MerklePathSet {
let parity = index.value() & 1;
let path_key = index.value() - parity;
let mut path = self
.paths
.get(&path_key)
.cloned()
.ok_or(MerkleError::NodeNotInSet(index.value()))?;
let mut path =
self.paths.get(&path_key).cloned().ok_or(MerkleError::NodeNotInSet(index))?;
path.remove(parity as usize);
Ok(path)
}
@ -200,7 +197,7 @@ impl MerklePathSet {
let path_key = index.value() - parity;
let path = match self.paths.get_mut(&path_key) {
Some(path) => path,
None => return Err(MerkleError::NodeNotInSet(base_index_value)),
None => return Err(MerkleError::NodeNotInSet(index)),
};
// Fill old_hashes vector -----------------------------------------------------------------

4
src/merkle/simple_smt/mod.rs
View file

@ -92,12 +92,12 @@ impl SimpleSmt {
for (key, value) in entries {
let old_value = tree.update_leaf(key, value)?;
if old_value != EMPTY_WORD || empty_entries.contains(&key) {
return Err(MerkleError::DuplicateValuesForKey(key));
return Err(MerkleError::DuplicateValuesForIndex(key));
}
// if we've processed an empty entry, add the key to the set of empty entry keys, and
// if this key was already in the set, return an error
if value == EMPTY_WORD && !empty_entries.insert(key) {
return Err(MerkleError::DuplicateValuesForKey(key));
return Err(MerkleError::DuplicateValuesForIndex(key));
}
}
Ok(tree)

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

@ -1,7 +1,8 @@
use super::{
BTreeMap, EmptySubtreeRoots, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, StarkField,
Vec, Word, EMPTY_WORD,
BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, MerkleError, MerklePath, NodeIndex, Rpo256,
RpoDigest, StarkField, Vec, Word, EMPTY_WORD, ZERO,
};
use core::cmp;
#[cfg(test)]
mod tests;
@ -9,12 +10,34 @@ mod tests;
// TIERED SPARSE MERKLE TREE
// ================================================================================================
/// Tiered (compacted) Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and
/// values are represented by 4 field elements.
///
/// Leaves in the tree can exist only on specific depths called "tiers". These depths are: 16, 32,
/// 48, and 64. Initially, when a tree is empty, it is equivalent to an empty Sparse Merkle tree
/// 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).
///
/// 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).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TieredSmt {
root: RpoDigest,
nodes: BTreeMap<NodeIndex, RpoDigest>,
upper_leaves: BTreeMap<NodeIndex, RpoDigest>,
bottom_leaves: BTreeMap<u64, Vec<RpoDigest>>,
upper_leaves: BTreeMap<NodeIndex, RpoDigest>, // node_index |-> key map
bottom_leaves: BTreeMap<u64, BottomLeaf>, // leaves of depth 64
values: BTreeMap<RpoDigest, Word>,
}
@ -22,12 +45,255 @@ impl TieredSmt {
// CONSTANTS
// --------------------------------------------------------------------------------------------
/// The number of levels between tiers.
const TIER_SIZE: u8 = 16;
/// Depths at which leaves can exist in a tiered SMT.
const TIER_DEPTHS: [u8; 4] = [16, 32, 48, 64];
/// Maximum node depth. This is also the bottom tier of the tree.
const MAX_DEPTH: u8 = 64;
// CONSTRUCTORS
// --------------------------------------------------------------------------------------------
pub fn new() -> Self {
/// Returns a new [TieredSmt] instantiated with the specified key-value pairs.
///
/// # Errors
/// Returns an error if the provided entries contain multiple values for the same key.
pub fn with_leaves<R, I>(entries: R) -> Result<Self, MerkleError>
where
R: IntoIterator<IntoIter = I>,
I: Iterator<Item = (RpoDigest, Word)> + ExactSizeIterator,
{
// create an empty tree
let mut tree = Self::default();
// append leaves to the tree returning an error if a duplicate entry for the same key
// is found
let mut empty_entries = BTreeSet::new();
for (key, value) in entries {
let old_value = tree.insert(key, value);
if old_value != EMPTY_WORD || empty_entries.contains(&key) {
return Err(MerkleError::DuplicateValuesForKey(key));
}
// if we've processed an empty entry, add the key to the set of empty entry keys, and
// if this key was already in the set, return an error
if value == EMPTY_WORD && !empty_entries.insert(key) {
return Err(MerkleError::DuplicateValuesForKey(key));
}
}
Ok(tree)
}
// PUBLIC ACCESSORS
// --------------------------------------------------------------------------------------------
/// Returns the root of this Merkle tree.
pub const fn root(&self) -> RpoDigest {
self.root
}
/// 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.into());
index.move_up();
}
Ok(path.into())
}
/// Returns the value associated with the specified key.
///
/// If nothing was inserted into this tree for the specified key, [ZERO; 4] is returned.
pub fn get_value(&self, key: RpoDigest) -> Word {
match self.values.get(&key) {
Some(value) => *value,
None => EMPTY_WORD,
}
}
// STATE MUTATORS
// --------------------------------------------------------------------------------------------
/// Inserts the provided value into the tree under the specified key and returns the value
/// previously stored under this key.
///
/// 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(EMPTY_WORD);
if old_value == value {
return old_value;
}
// 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 node and return the old value
self.insert_node(index, key, value);
old_value
}
// HELPER METHODS
// --------------------------------------------------------------------------------------------
/// 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() > Self::MAX_DEPTH {
return Err(MerkleError::DepthTooBig(index.depth() as u64));
} else {
// make sure that there are no leaf nodes in the ancestors of the index; since leaf
// nodes can live at specific depth, we just need to check these depths.
let tier = get_index_tier(&index);
let mut tier_index = index;
for &depth in Self::TIER_DEPTHS[..tier].iter().rev() {
tier_index.move_up_to(depth);
if self.upper_leaves.contains_key(&tier_index) {
return Err(MerkleError::NodeNotInSet(index));
}
}
}
Ok(())
}
/// Returns a node at the specified index. If the node does not exist at this index, a root
/// for an empty subtree at the index's depth is returned.
///
/// Unlike [TieredSmt::get_node()] this does not perform any checks to verify that the returned
/// node is valid in the context of this tree.
fn get_node_unchecked(&self, index: &NodeIndex) -> RpoDigest {
match self.nodes.get(index) {
Some(node) => *node,
None => EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[index.depth() as usize],
}
}
/// 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) {
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 {
// 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()
} 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();
}
// update the root
self.root = node;
}
}
impl Default for TieredSmt {
fn default() -> Self {
Self {
root: EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[0],
nodes: BTreeMap::new(),
@ -36,178 +302,108 @@ impl TieredSmt {
values: BTreeMap::new(),
}
}
// PUBLIC ACCESSORS
// --------------------------------------------------------------------------------------------
pub const fn root(&self) -> RpoDigest {
self.root
}
pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
if index.is_root() {
return Err(MerkleError::DepthTooSmall(index.depth()));
} else if index.depth() > Self::MAX_DEPTH {
return Err(MerkleError::DepthTooBig(index.depth() as u64));
} else if !self.is_node_available(index) {
todo!()
}
Ok(self.get_branch_node(&index))
}
pub fn get_path(&self, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
if index.is_root() {
return Err(MerkleError::DepthTooSmall(index.depth()));
} else if index.depth() > Self::MAX_DEPTH {
return Err(MerkleError::DepthTooBig(index.depth() as u64));
} else if !self.is_node_available(index) {
todo!()
}
let mut path = Vec::with_capacity(index.depth() as usize);
for _ in 0..index.depth() {
let node = self.get_branch_node(&index.sibling());
path.push(node.into());
index.move_up();
}
Ok(path.into())
}
pub fn get_value(&self, key: RpoDigest) -> Result<Word, MerkleError> {
match self.values.get(&key) {
Some(value) => Ok(*value),
None => Ok(EMPTY_WORD),
}
}
// STATE MUTATORS
// --------------------------------------------------------------------------------------------
pub fn insert(&mut self, key: RpoDigest, value: Word) -> Result<Word, MerkleError> {
let (mut index, leaf_key) = self.get_insert_location(&key);
if let Some(other_key) = leaf_key {
if other_key != key {
let common_prefix_len = get_common_prefix_length(&key, &other_key);
let depth = common_prefix_len + 16;
let other_index = key_to_index(&other_key, depth);
self.move_leaf_node(other_key, index, other_index);
index = key_to_index(&key, depth);
}
}
let old_value = self.values.insert(key, value).unwrap_or(EMPTY_WORD);
if value != old_value {
self.upper_leaves.insert(index, key);
let new_node = build_leaf_node(key, value, index.depth().into());
self.root = self.update_path(index, new_node);
}
Ok(old_value)
}
// HELPER METHODS
// --------------------------------------------------------------------------------------------
fn is_node_available(&self, index: NodeIndex) -> bool {
match index.depth() {
32 => true,
48 => true,
_ => true,
}
}
fn get_branch_node(&self, index: &NodeIndex) -> RpoDigest {
match self.nodes.get(index) {
Some(node) => *node,
None => EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[index.depth() as usize],
}
}
fn get_insert_location(&self, key: &RpoDigest) -> (NodeIndex, Option<RpoDigest>) {
let mse = Word::from(key)[3].as_int();
for depth in (16..64).step_by(16) {
let index = NodeIndex::new(depth, mse >> (Self::MAX_DEPTH - depth)).unwrap();
if let Some(leaf_key) = self.upper_leaves.get(&index) {
return (index, Some(*leaf_key));
} else if self.nodes.contains_key(&index) {
continue;
} else {
return (index, None);
}
}
// TODO: handle bottom tier
unimplemented!()
}
fn move_leaf_node(&mut self, key: RpoDigest, old_index: NodeIndex, new_index: NodeIndex) {
self.upper_leaves.remove(&old_index).unwrap();
self.upper_leaves.insert(new_index, key);
let value = *self.values.get(&key).unwrap();
let new_node = build_leaf_node(key, value, new_index.depth().into());
self.update_path(new_index, new_node);
}
fn update_path(&mut self, mut index: NodeIndex, mut node: RpoDigest) -> RpoDigest {
for _ in 0..index.depth() {
self.nodes.insert(index, node);
let sibling = self.get_branch_node(&index.sibling());
node = Rpo256::merge(&index.build_node(node, sibling));
index.move_up();
}
node
}
}
impl Default for TieredSmt {
fn default() -> Self {
Self::new()
}
}
// 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);
let remaining = (key[3].as_int() << depth) >> depth;
key[3] = remaining.into();
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()
}
fn build_leaf_node(key: RpoDigest, value: Word, depth: u32) -> RpoDigest {
let remaining_path = get_remaining_path(key, depth);
Rpo256::merge_in_domain(&[remaining_path, value.into()], depth.into())
}
fn get_common_prefix_length(key1: &RpoDigest, key2: &RpoDigest) -> u8 {
let e1 = Word::from(key1)[3].as_int();
let e2 = Word::from(key2)[3].as_int();
if e1 == e2 {
64
} else if e1 >> 16 == e2 >> 16 {
48
} else if e1 >> 32 == e2 >> 32 {
32
} else if e1 >> 48 == e2 >> 48 {
16
} else {
0
}
}
/// 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 => mse >> (depth as u32),
16 | 32 | 48 | 64 => mse >> ((TieredSmt::MAX_DEPTH - depth) as u32),
_ => unreachable!("invalid depth: {depth}"),
};
// TODO: use unchecked version?
NodeIndex::new(depth, value).unwrap()
NodeIndex::new_unchecked(depth, value)
}
/// Returns tiered common prefix length between the most significant elements of the provided keys.
///
/// Specifically:
/// - returns 64 if the most significant elements are equal.
/// - returns 48 if the common prefix is between 48 and 63 bits.
/// - 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();
let ex = (e1 ^ e2).leading_zeros() as u8;
(ex / 16) * 16
}
/// Returns a tier for the specified index.
///
/// 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,
}
}
// BOTTOM LEAF
// ================================================================================================
/// Stores contents of the bottom leaf (i.e., leaf at depth = 64) in a [TieredSmt].
///
/// 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 {
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 mut values = BTreeMap::new();
let key = get_remaining_path(key, TieredSmt::MAX_DEPTH as u32);
values.insert(key.into(), value);
Self { 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);
}
// TODO: hash in domain
Rpo256::hash_elements(&elements)
}
}

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

@ -1,11 +1,14 @@
use super::{
super::{super::ONE, Felt, MerkleStore, WORD_SIZE},
get_remaining_path, EmptySubtreeRoots, NodeIndex, Rpo256, RpoDigest, TieredSmt, Word,
super::{
super::{ONE, ZERO},
Felt, MerkleStore, WORD_SIZE,
},
get_remaining_path, EmptySubtreeRoots, NodeIndex, Rpo256, RpoDigest, TieredSmt, Vec, Word,
};
#[test]
fn tsmt_insert_one() {
let mut smt = TieredSmt::new();
let mut smt = TieredSmt::default();
let mut store = MerkleStore::default();
let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
@ -15,15 +18,15 @@ fn tsmt_insert_one() {
// since the tree is empty, the first node will be inserted at depth 16 and the index will be
// 16 most significant bits of the key
let index = NodeIndex::make(16, raw >> 48);
let leaf_node = compute_leaf_node(key, value, 16);
let leaf_node = build_leaf_node(key, value, 16);
let tree_root = store.set_node(smt.root().into(), index, leaf_node.into()).unwrap().root;
smt.insert(key, value).unwrap();
smt.insert(key, value);
assert_eq!(smt.root(), tree_root.into());
// make sure the value was inserted, and the node is at the expected index
assert_eq!(smt.get_value(key).unwrap(), value);
assert_eq!(smt.get_value(key), value);
assert_eq!(smt.get_node(index).unwrap(), leaf_node);
// make sure the paths we get from the store and the tree match
@ -32,15 +35,15 @@ fn tsmt_insert_one() {
}
#[test]
fn tsmt_insert_two() {
let mut smt = TieredSmt::new();
fn tsmt_insert_two_16() {
let mut smt = TieredSmt::default();
let mut store = MerkleStore::default();
// --- insert the first value ---------------------------------------------
let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
let val_a = [ONE; WORD_SIZE];
smt.insert(key_a, val_a).unwrap();
smt.insert(key_a, val_a);
// --- insert the second value --------------------------------------------
// the key for this value has the same 16-bit prefix as the key for the first value,
@ -48,28 +51,72 @@ fn tsmt_insert_two() {
let raw_b = 0b_10101010_10101010_10011111_11111111_10010110_10010011_11100000_00000000_u64;
let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
let val_b = [Felt::new(2); WORD_SIZE];
smt.insert(key_b, val_b).unwrap();
smt.insert(key_b, val_b);
// --- build Merkle store with equivalent data ----------------------------
let mut tree_root = get_init_root();
let index_a = NodeIndex::make(32, raw_a >> 32);
let leaf_node_a = compute_leaf_node(key_a, val_a, 32);
let leaf_node_a = build_leaf_node(key_a, val_a, 32);
tree_root = store.set_node(tree_root, index_a, leaf_node_a.into()).unwrap().root;
let index_b = NodeIndex::make(32, raw_b >> 32);
let leaf_node_b = compute_leaf_node(key_b, val_b, 32);
let leaf_node_b = build_leaf_node(key_b, val_b, 32);
tree_root = store.set_node(tree_root, index_b, leaf_node_b.into()).unwrap().root;
// --- verify that data is consistent between store and tree --------------
assert_eq!(smt.root(), tree_root.into());
assert_eq!(smt.get_value(key_a).unwrap(), val_a);
assert_eq!(smt.get_value(key_a), val_a);
assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
let expected_path = store.get_path(tree_root, index_a).unwrap().path;
assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
assert_eq!(smt.get_value(key_b).unwrap(), val_b);
assert_eq!(smt.get_value(key_b), val_b);
assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
let expected_path = store.get_path(tree_root, index_b).unwrap().path;
assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
}
#[test]
fn tsmt_insert_two_32() {
let mut smt = TieredSmt::default();
let mut store = MerkleStore::default();
// --- insert the first value ---------------------------------------------
let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
let val_a = [ONE; WORD_SIZE];
smt.insert(key_a, val_a);
// --- insert the second value --------------------------------------------
// the key for this value has the same 32-bit prefix as the key for the first value,
// thus, on insertions, both values should be pushed to depth 48 tier
let raw_b = 0b_10101010_10101010_00011111_11111111_00010110_10010011_11100000_00000000_u64;
let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
let val_b = [Felt::new(2); WORD_SIZE];
smt.insert(key_b, val_b);
// --- build Merkle store with equivalent data ----------------------------
let mut tree_root = get_init_root();
let index_a = NodeIndex::make(48, raw_a >> 16);
let leaf_node_a = build_leaf_node(key_a, val_a, 48);
tree_root = store.set_node(tree_root, index_a, leaf_node_a.into()).unwrap().root;
let index_b = NodeIndex::make(48, raw_b >> 16);
let leaf_node_b = build_leaf_node(key_b, val_b, 48);
tree_root = store.set_node(tree_root, index_b, leaf_node_b.into()).unwrap().root;
// --- verify that data is consistent between store and tree --------------
assert_eq!(smt.root(), tree_root.into());
assert_eq!(smt.get_value(key_a), val_a);
assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
let expected_path = store.get_path(tree_root, index_a).unwrap().path;
assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
assert_eq!(smt.get_value(key_b), val_b);
assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
let expected_path = store.get_path(tree_root, index_b).unwrap().path;
assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
@ -77,14 +124,14 @@ fn tsmt_insert_two() {
#[test]
fn tsmt_insert_three() {
let mut smt = TieredSmt::new();
let mut smt = TieredSmt::default();
let mut store = MerkleStore::default();
// --- insert the first value ---------------------------------------------
let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
let val_a = [ONE; WORD_SIZE];
smt.insert(key_a, val_a).unwrap();
smt.insert(key_a, val_a);
// --- insert the second value --------------------------------------------
// the key for this value has the same 16-bit prefix as the key for the first value,
@ -92,7 +139,7 @@ fn tsmt_insert_three() {
let raw_b = 0b_10101010_10101010_10011111_11111111_10010110_10010011_11100000_00000000_u64;
let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
let val_b = [Felt::new(2); WORD_SIZE];
smt.insert(key_b, val_b).unwrap();
smt.insert(key_b, val_b);
// --- insert the third value ---------------------------------------------
// the key for this value has the same 16-bit prefix as the keys for the first two,
@ -101,37 +148,37 @@ fn tsmt_insert_three() {
let raw_c = 0b_10101010_10101010_11011111_11111111_10010110_10010011_11100000_00000000_u64;
let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
let val_c = [Felt::new(3); WORD_SIZE];
smt.insert(key_c, val_c).unwrap();
smt.insert(key_c, val_c);
// --- build Merkle store with equivalent data ----------------------------
let mut tree_root = get_init_root();
let index_a = NodeIndex::make(32, raw_a >> 32);
let leaf_node_a = compute_leaf_node(key_a, val_a, 32);
let leaf_node_a = build_leaf_node(key_a, val_a, 32);
tree_root = store.set_node(tree_root, index_a, leaf_node_a.into()).unwrap().root;
let index_b = NodeIndex::make(32, raw_b >> 32);
let leaf_node_b = compute_leaf_node(key_b, val_b, 32);
let leaf_node_b = build_leaf_node(key_b, val_b, 32);
tree_root = store.set_node(tree_root, index_b, leaf_node_b.into()).unwrap().root;
let index_c = NodeIndex::make(32, raw_c >> 32);
let leaf_node_c = compute_leaf_node(key_c, val_c, 32);
let leaf_node_c = build_leaf_node(key_c, val_c, 32);
tree_root = store.set_node(tree_root, index_c, leaf_node_c.into()).unwrap().root;
// --- verify that data is consistent between store and tree --------------
assert_eq!(smt.root(), tree_root.into());
assert_eq!(smt.get_value(key_a).unwrap(), val_a);
assert_eq!(smt.get_value(key_a), val_a);
assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
let expected_path = store.get_path(tree_root, index_a).unwrap().path;
assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
assert_eq!(smt.get_value(key_b).unwrap(), val_b);
assert_eq!(smt.get_value(key_b), val_b);
assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
let expected_path = store.get_path(tree_root, index_b).unwrap().path;
assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
assert_eq!(smt.get_value(key_c).unwrap(), val_c);
assert_eq!(smt.get_value(key_c), val_c);
assert_eq!(smt.get_node(index_c).unwrap(), leaf_node_c);
let expected_path = store.get_path(tree_root, index_c).unwrap().path;
assert_eq!(smt.get_path(index_c).unwrap(), expected_path);
@ -139,33 +186,161 @@ fn tsmt_insert_three() {
#[test]
fn tsmt_update() {
let mut smt = TieredSmt::new();
let mut smt = TieredSmt::default();
let mut store = MerkleStore::default();
// --- insert a value into the tree ---------------------------------------
let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
let value_a = [ONE; WORD_SIZE];
smt.insert(key, value_a).unwrap();
smt.insert(key, value_a);
// --- update value ---------------------------------------
// --- update the value ---------------------------------------------------
let value_b = [Felt::new(2); WORD_SIZE];
smt.insert(key, value_b).unwrap();
smt.insert(key, value_b);
// --- verify consistency -------------------------------------------------
let mut tree_root = get_init_root();
let index = NodeIndex::make(16, raw >> 48);
let leaf_node = compute_leaf_node(key, value_b, 16);
let leaf_node = build_leaf_node(key, value_b, 16);
tree_root = store.set_node(tree_root, index, leaf_node.into()).unwrap().root;
assert_eq!(smt.root(), tree_root.into());
assert_eq!(smt.get_value(key).unwrap(), value_b);
assert_eq!(smt.get_value(key), value_b);
assert_eq!(smt.get_node(index).unwrap(), leaf_node);
let expected_path = store.get_path(tree_root, index).unwrap().path;
assert_eq!(smt.get_path(index).unwrap(), expected_path);
}
// BOTTOM TIER TESTS
// ================================================================================================
#[test]
fn tsmt_bottom_tier() {
let mut smt = TieredSmt::default();
let mut store = MerkleStore::default();
// common prefix for the keys
let prefix = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
// --- insert the first value ---------------------------------------------
let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(prefix)]);
let val_a = [ONE; WORD_SIZE];
smt.insert(key_a, val_a);
// --- insert the second value --------------------------------------------
// this key has the same 64-bit prefix and thus both values should end up in the same
// node at depth 64
let key_b = RpoDigest::from([ZERO, ONE, ONE, Felt::new(prefix)]);
let val_b = [Felt::new(2); WORD_SIZE];
smt.insert(key_b, val_b);
// --- build Merkle store with equivalent data ----------------------------
let index = NodeIndex::make(64, prefix);
// to build bottom leaf we sort by key starting with the least significant element, thus
// key_b is smaller than key_a.
let leaf_node = build_bottom_leaf_node(&[key_b, key_a], &[val_b, val_a]);
let mut tree_root = get_init_root();
tree_root = store.set_node(tree_root, index, leaf_node.into()).unwrap().root;
// --- verify that data is consistent between store and tree --------------
assert_eq!(smt.root(), tree_root.into());
assert_eq!(smt.get_value(key_a), val_a);
assert_eq!(smt.get_value(key_b), val_b);
assert_eq!(smt.get_node(index).unwrap(), leaf_node);
let expected_path = store.get_path(tree_root, index).unwrap().path;
assert_eq!(smt.get_path(index).unwrap(), expected_path);
}
#[test]
fn tsmt_bottom_tier_two() {
let mut smt = TieredSmt::default();
let mut store = MerkleStore::default();
// --- insert the first value ---------------------------------------------
let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
let val_a = [ONE; WORD_SIZE];
smt.insert(key_a, val_a);
// --- insert the second value --------------------------------------------
// the key for this value has the same 48-bit prefix as the key for the first value,
// thus, on insertions, both should end up in different nodes at depth 64
let raw_b = 0b_10101010_10101010_00011111_11111111_10010110_10010011_01100000_00000000_u64;
let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
let val_b = [Felt::new(2); WORD_SIZE];
smt.insert(key_b, val_b);
// --- build Merkle store with equivalent data ----------------------------
let mut tree_root = get_init_root();
let index_a = NodeIndex::make(64, raw_a);
let leaf_node_a = build_bottom_leaf_node(&[key_a], &[val_a]);
tree_root = store.set_node(tree_root, index_a, leaf_node_a.into()).unwrap().root;
let index_b = NodeIndex::make(64, raw_b);
let leaf_node_b = build_bottom_leaf_node(&[key_b], &[val_b]);
tree_root = store.set_node(tree_root, index_b, leaf_node_b.into()).unwrap().root;
// --- verify that data is consistent between store and tree --------------
assert_eq!(smt.root(), tree_root.into());
assert_eq!(smt.get_value(key_a), val_a);
assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
let expected_path = store.get_path(tree_root, index_a).unwrap().path;
assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
assert_eq!(smt.get_value(key_b), val_b);
assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
let expected_path = store.get_path(tree_root, index_b).unwrap().path;
assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
}
// ERROR TESTS
// ================================================================================================
#[test]
fn tsmt_node_not_available() {
let mut smt = TieredSmt::default();
let raw = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
let value = [ONE; WORD_SIZE];
// build an index which is just below the inserted leaf node
let index = NodeIndex::make(17, raw >> 47);
// since we haven't inserted the node yet, we should be able to get node and path to this index
assert!(smt.get_node(index).is_ok());
assert!(smt.get_path(index).is_ok());
smt.insert(key, value);
// but once the node is inserted, everything under it should be unavailable
assert!(smt.get_node(index).is_err());
assert!(smt.get_path(index).is_err());
let index = NodeIndex::make(32, raw >> 32);
assert!(smt.get_node(index).is_err());
assert!(smt.get_path(index).is_err());
let index = NodeIndex::make(34, raw >> 30);
assert!(smt.get_node(index).is_err());
assert!(smt.get_path(index).is_err());
let index = NodeIndex::make(50, raw >> 14);
assert!(smt.get_node(index).is_err());
assert!(smt.get_path(index).is_err());
let index = NodeIndex::make(64, raw);
assert!(smt.get_node(index).is_err());
assert!(smt.get_path(index).is_err());
}
// HELPER FUNCTIONS
// ================================================================================================
@ -173,7 +348,21 @@ fn get_init_root() -> Word {
EmptySubtreeRoots::empty_hashes(64)[0].into()
}
fn compute_leaf_node(key: RpoDigest, value: Word, depth: u8) -> 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())
}
fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest {
assert_eq!(keys.len(), values.len());
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(val);
}
Rpo256::hash_elements(&elements)
}