diff --git a/src/hash/rpo/digest.rs b/src/hash/rpo/digest.rs index 2edfd27..0e6c310 100644 --- a/src/hash/rpo/digest.rs +++ b/src/hash/rpo/digest.rs @@ -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 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 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 // ================================================================================================ diff --git a/src/merkle/index.rs b/src/merkle/index.rs index b1c7389..5a559a1 100644 --- a/src/merkle/index.rs +++ b/src/merkle/index.rs @@ -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) } } diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 6c666c0..631b960 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -47,14 +47,15 @@ pub enum MerkleError { ConflictingRoots(Vec), 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), } } diff --git a/src/merkle/path_set.rs b/src/merkle/path_set.rs index 653fab1..ed945fb 100644 --- a/src/merkle/path_set.rs +++ b/src/merkle/path_set.rs @@ -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 ----------------------------------------------------------------- diff --git a/src/merkle/simple_smt/mod.rs b/src/merkle/simple_smt/mod.rs index d56f294..9d25316 100644 --- a/src/merkle/simple_smt/mod.rs +++ b/src/merkle/simple_smt/mod.rs @@ -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) diff --git a/src/merkle/tiered_smt/mod.rs b/src/merkle/tiered_smt/mod.rs index a8c9255..d3ca68a 100644 --- a/src/merkle/tiered_smt/mod.rs +++ b/src/merkle/tiered_smt/mod.rs @@ -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, - upper_leaves: BTreeMap, - bottom_leaves: BTreeMap>, + upper_leaves: BTreeMap, // node_index |-> key map + bottom_leaves: BTreeMap, // leaves of depth 64 values: BTreeMap, } @@ -22,192 +45,365 @@ 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 { - Self { - root: EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[0], - nodes: BTreeMap::new(), - upper_leaves: BTreeMap::new(), - bottom_leaves: BTreeMap::new(), - values: BTreeMap::new(), + /// 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(entries: R) -> Result + where + R: IntoIterator, + I: Iterator + 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 { - 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)) + 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 { - 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!() - } + self.validate_node_access(index)?; let mut path = Vec::with_capacity(index.depth() as usize); for _ in 0..index.depth() { - let node = self.get_branch_node(&index.sibling()); + let node = self.get_node_unchecked(&index.sibling()); path.push(node.into()); index.move_up(); } + Ok(path.into()) } - pub fn get_value(&self, key: RpoDigest) -> Result { + /// 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) => Ok(*value), - None => Ok(EMPTY_WORD), + Some(value) => *value, + None => EMPTY_WORD, } } // STATE MUTATORS // -------------------------------------------------------------------------------------------- - pub fn insert(&mut self, key: RpoDigest, value: Word) -> Result { + /// 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 { - let common_prefix_len = get_common_prefix_length(&key, &other_key); - let depth = common_prefix_len + 16; + // 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); - self.move_leaf_node(other_key, index, other_index); + 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); } } - 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) + // insert the node and return the old value + self.insert_node(index, key, value); + old_value } // HELPER METHODS // -------------------------------------------------------------------------------------------- - fn is_node_available(&self, index: NodeIndex) -> bool { - match index.depth() { - 32 => true, - 48 => true, - _ => true, + /// 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(()) } - fn get_branch_node(&self, index: &NodeIndex) -> RpoDigest { + /// 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) { + // 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 (16..64).step_by(16) { - let index = NodeIndex::new(depth, mse >> (Self::MAX_DEPTH - depth)).unwrap(); + 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) { - continue; - } else { + } else if !self.nodes.contains_key(&index) { return (index, None); } } - // TODO: handle bottom tier - unimplemented!() + // 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) } - 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); - } + /// 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()) + }; - fn update_path(&mut self, mut index: NodeIndex, mut node: RpoDigest) -> RpoDigest { + // 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_branch_node(&index.sibling()); + let sibling = self.get_node_unchecked(&index.sibling()); node = Rpo256::merge(&index.build_node(node, sibling)); index.move_up(); } - node + + // update the root + self.root = node; } } impl Default for TieredSmt { fn default() -> Self { - Self::new() + Self { + root: EmptySubtreeRoots::empty_hashes(Self::MAX_DEPTH)[0], + nodes: BTreeMap::new(), + upper_leaves: BTreeMap::new(), + bottom_leaves: BTreeMap::new(), + values: BTreeMap::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()) +/// 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) } -fn get_common_prefix_length(key1: &RpoDigest, key2: &RpoDigest) -> u8 { +/// 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 +} - 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 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, } } -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), - _ => unreachable!("invalid depth: {depth}"), - }; +// BOTTOM LEAF +// ================================================================================================ - // TODO: use unchecked version? - NodeIndex::new(depth, value).unwrap() +/// 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) + } } diff --git a/src/merkle/tiered_smt/tests.rs b/src/merkle/tiered_smt/tests.rs index 121871d..5b36434 100644 --- a/src/merkle/tiered_smt/tests.rs +++ b/src/merkle/tiered_smt/tests.rs @@ -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), 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).unwrap(), value_b); + 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) +}