diff --git a/src/merkle/tiered_smt/mod.rs b/src/merkle/tiered_smt/mod.rs index 3269517..7ec8afe 100644 --- a/src/merkle/tiered_smt/mod.rs +++ b/src/merkle/tiered_smt/mod.rs @@ -1,6 +1,6 @@ use super::{ BTreeMap, BTreeSet, EmptySubtreeRoots, Felt, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, - Rpo256, RpoDigest, StarkField, Vec, Word, ZERO, + Rpo256, RpoDigest, StarkField, Vec, Word, }; use core::cmp; @@ -27,8 +27,8 @@ mod tests; /// /// 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]). +/// - Leaf node at depths 16, 32, or 64: hash(key, value, domain=depth). +/// - Leaf node at depth 64: hash([key_0, value_0, ..., key_n, value_n, domain=64]). /// /// 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). @@ -36,7 +36,7 @@ mod tests; pub struct TieredSmt { root: RpoDigest, nodes: BTreeMap, - upper_leaves: BTreeMap, // node_index |-> key map + upper_leaves: BTreeMap, // node_index |-> key bottom_leaves: BTreeMap, // leaves of depth 64 values: BTreeMap, } @@ -180,7 +180,7 @@ impl TieredSmt { 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); + self.insert_leaf_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); @@ -188,7 +188,7 @@ impl TieredSmt { } // insert the node and return the old value - self.insert_node(index, key, value); + self.insert_leaf_node(index, key, value); old_value } @@ -307,8 +307,9 @@ impl TieredSmt { /// 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) { + fn insert_leaf_node(&mut self, mut index: NodeIndex, key: RpoDigest, value: Word) { let depth = index.depth(); + debug_assert!(Self::TIER_DEPTHS.contains(&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 { @@ -323,9 +324,8 @@ impl TieredSmt { // 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()) + // the node value is computed as: hash(key || value, domain = depth) + Rpo256::merge_in_domain(&[key, value.into()], depth.into()) }; // insert the node and update the path from the node to the root @@ -357,21 +357,6 @@ impl Default for TieredSmt { // HELPER FUNCTIONS // ================================================================================================ -/// Returns the remaining path for the specified key at the specified depth. -/// -/// Remaining path is computed by setting n most significant bits of the key to zeros, where n is -/// the specified depth. -fn get_remaining_path(key: RpoDigest, depth: u32) -> RpoDigest { - let mut key = Word::from(key); - key[3] = if depth == 64 { - ZERO - } else { - // remove `depth` bits from the most significant key element - ((key[3].as_int() << depth) >> depth).into() - }; - key.into() -} - /// Returns 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 @@ -443,14 +428,12 @@ impl BottomLeaf { pub fn new(key: RpoDigest, value: Word) -> Self { let prefix = Word::from(key)[3].as_int(); let mut values = BTreeMap::new(); - let key = get_remaining_path(key, TieredSmt::MAX_DEPTH as u32); values.insert(key.into(), value); Self { prefix, values } } /// Adds a new key-value pair to this leaf. pub fn add_value(&mut self, key: RpoDigest, value: Word) { - let key = get_remaining_path(key, TieredSmt::MAX_DEPTH as u32); self.values.insert(key.into(), value); } @@ -476,7 +459,7 @@ impl BottomLeaf { Felt::new(key[0]), Felt::new(key[1]), Felt::new(key[2]), - Felt::new(self.prefix), + Felt::new(key[3]), ]); (key, *val) }) diff --git a/src/merkle/tiered_smt/tests.rs b/src/merkle/tiered_smt/tests.rs index fbf26ea..d8e6723 100644 --- a/src/merkle/tiered_smt/tests.rs +++ b/src/merkle/tiered_smt/tests.rs @@ -1,7 +1,6 @@ use super::{ super::{super::ONE, Felt, MerkleStore, WORD_SIZE, ZERO}, - get_remaining_path, EmptySubtreeRoots, InnerNodeInfo, NodeIndex, Rpo256, RpoDigest, TieredSmt, - Vec, Word, + EmptySubtreeRoots, InnerNodeInfo, NodeIndex, Rpo256, RpoDigest, TieredSmt, Vec, Word, }; #[test] @@ -411,8 +410,7 @@ fn get_init_root() -> RpoDigest { } fn build_leaf_node(key: RpoDigest, value: Word, depth: u8) -> RpoDigest { - let remaining_path = get_remaining_path(key, depth as u32); - Rpo256::merge_in_domain(&[remaining_path, value.into()], depth.into()) + Rpo256::merge_in_domain(&[key, value.into()], depth.into()) } fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest { @@ -420,9 +418,7 @@ fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest { let mut elements = Vec::with_capacity(keys.len()); for (key, val) in keys.iter().zip(values.iter()) { - let mut key = Word::from(key); - key[3] = ZERO; - elements.extend_from_slice(&key); + elements.extend_from_slice(key.as_elements()); elements.extend_from_slice(val.as_slice()); }