diff --git a/src/merkle/tiered_smt/mod.rs b/src/merkle/tiered_smt/mod.rs index a0bd723..fbdf2d3 100644 --- a/src/merkle/tiered_smt/mod.rs +++ b/src/merkle/tiered_smt/mod.rs @@ -41,7 +41,7 @@ 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(key, value, domain=depth). -/// - Leaf node at depth 64: hash([key_0, value_0, ..., key_n, value_n, domain=64]). +/// - Leaf node at depth 64: hash([key_0, value_0, ..., key_n, value_n], domain=64). #[derive(Debug, Clone, PartialEq, Eq)] pub struct TieredSmt { root: RpoDigest, @@ -306,10 +306,23 @@ impl TieredSmt { debug_assert!(leaf_exists); // if the leaf is at the bottom tier and after removing the key-value pair from it, the - // leaf is still not empty, just recompute its hash and update the leaf node. + // leaf is still not empty, we either just update it, or move it up to a higher tier (if + // the leaf doesn't have siblings at lower tiers) if index.depth() == Self::MAX_DEPTH { - if let Some(values) = self.values.get_all(index.value()) { - let node = hash_bottom_leaf(&values); + if let Some(entries) = self.values.get_all(index.value()) { + // if there is only one key-value pair left at the bottom leaf, and it can be + // moved up to a higher tier, truncate the branch and return + if entries.len() == 1 { + let new_depth = self.nodes.get_last_single_child_parent_depth(index.value()); + if new_depth != Self::MAX_DEPTH { + let node = hash_upper_leaf(entries[0].0, entries[0].1, new_depth); + self.root = self.nodes.truncate_branch(index.value(), new_depth, node); + return old_value; + } + } + + // otherwise just recompute the leaf hash and update the leaf node + let node = hash_bottom_leaf(&entries); self.root = self.nodes.update_leaf_node(index, node); return old_value; }; diff --git a/src/merkle/tiered_smt/nodes.rs b/src/merkle/tiered_smt/nodes.rs index 9db4ba3..0d94091 100644 --- a/src/merkle/tiered_smt/nodes.rs +++ b/src/merkle/tiered_smt/nodes.rs @@ -118,6 +118,24 @@ impl NodeStore { (index, self.bottom_leaves.contains(&index.value())) } + /// Traverses the tree up from the bottom tier starting at the specified leaf index and + /// returns the depth of the first node which hash more than one child. The returned depth + /// is rounded up to the next tier. + pub fn get_last_single_child_parent_depth(&self, leaf_index: u64) -> u8 { + let mut index = NodeIndex::new_unchecked(MAX_DEPTH, leaf_index); + + for _ in (TIER_DEPTHS[0]..MAX_DEPTH).rev() { + let sibling_index = index.sibling(); + if self.nodes.contains_key(&sibling_index) { + break; + } + index.move_up(); + } + + let tier = (index.depth() - 1) / TIER_SIZE; + TIER_DEPTHS[tier as usize] + } + // ITERATORS // -------------------------------------------------------------------------------------------- @@ -201,15 +219,6 @@ impl NodeStore { debug_assert_eq!(removed_leaf.depth(), retained_leaf.depth()); debug_assert!(removed_leaf.depth() > new_depth); - // clear leaf flags - if removed_leaf.depth() == MAX_DEPTH { - self.bottom_leaves.remove(&removed_leaf.value()); - self.bottom_leaves.remove(&retained_leaf.value()); - } else { - self.upper_leaves.remove(&removed_leaf); - self.upper_leaves.remove(&retained_leaf); - } - // remove the branches leading up to the tier to which the retained leaf is to be moved self.remove_branch(removed_leaf, new_depth); self.remove_branch(retained_leaf, new_depth); @@ -298,6 +307,25 @@ impl NodeStore { self.update_leaf_node(index, node) } + /// Truncates a branch starting with specified leaf at the bottom tier to new depth. + /// + /// This involves removing the part of the branch below the new depth, and then inserting a new + /// // node at the new depth. + pub fn truncate_branch( + &mut self, + leaf_index: u64, + new_depth: u8, + node: RpoDigest, + ) -> RpoDigest { + debug_assert!(self.bottom_leaves.contains(&leaf_index)); + + let mut leaf_index = LeafNodeIndex::new(NodeIndex::new_unchecked(MAX_DEPTH, leaf_index)); + self.remove_branch(leaf_index, new_depth); + + leaf_index.move_up_to(new_depth); + self.insert_leaf_node(leaf_index, node) + } + // HELPER METHODS // -------------------------------------------------------------------------------------------- @@ -360,11 +388,18 @@ impl NodeStore { } } - /// Removes a sequence of nodes starting at the specified index and traversing the - /// tree up to the specified depth. The node at the `end_depth` is also removed. + /// Removes a sequence of nodes starting at the specified index and traversing the tree up to + /// the specified depth. The node at the `end_depth` is also removed, and the appropriate leaf + /// flag is cleared. /// /// This method does not update any other nodes and does not recompute the tree root. fn remove_branch(&mut self, index: LeafNodeIndex, end_depth: u8) { + if index.depth() == MAX_DEPTH { + self.bottom_leaves.remove(&index.value()); + } else { + self.upper_leaves.remove(&index); + } + let mut index: NodeIndex = index.into(); assert!(index.depth() > end_depth); for _ in 0..(index.depth() - end_depth + 1) { diff --git a/src/merkle/tiered_smt/tests.rs b/src/merkle/tiered_smt/tests.rs index c1b7649..61e6081 100644 --- a/src/merkle/tiered_smt/tests.rs +++ b/src/merkle/tiered_smt/tests.rs @@ -460,6 +460,131 @@ fn tsmt_delete_64() { assert_eq!(smt, smt0); } +#[test] +fn tsmt_delete_64_leaf_promotion() { + let mut smt = TieredSmt::default(); + + // --- delete from bottom tier (no promotion to upper tiers) -------------- + + // insert a value into the tree + let raw_a = 0b_01010101_01010101_11111111_11111111_10101010_10101010_11111111_00000000_u64; + let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]); + let value_a = [ONE, ONE, ONE, ONE]; + smt.insert(key_a, value_a); + + // insert another value with a key having the same 64-bit prefix + let key_b = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]); + let value_b = [ONE, ONE, ONE, ZERO]; + smt.insert(key_b, value_b); + + // insert a value with a key which shared the same 48-bit prefix + let raw_c = 0b_01010101_01010101_11111111_11111111_10101010_10101010_00111111_00000000_u64; + let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); + let value_c = [ONE, ONE, ZERO, ZERO]; + smt.insert(key_c, value_c); + + // delete entry A and compare to the tree which was built from B and C + smt.insert(key_a, EMPTY_WORD); + + let mut expected_smt = TieredSmt::default(); + expected_smt.insert(key_b, value_b); + expected_smt.insert(key_c, value_c); + assert_eq!(smt, expected_smt); + + // entries B and C should stay at depth 64 + assert_eq!(smt.nodes.get_leaf_index(&key_b).0.depth(), 64); + assert_eq!(smt.nodes.get_leaf_index(&key_c).0.depth(), 64); + + // --- delete from bottom tier (promotion to depth 48) -------------------- + + let mut smt = TieredSmt::default(); + smt.insert(key_a, value_a); + smt.insert(key_b, value_b); + + // insert a value with a key which shared the same 32-bit prefix + let raw_c = 0b_01010101_01010101_11111111_11111111_11101010_10101010_11111111_00000000_u64; + let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); + smt.insert(key_c, value_c); + + // delete entry A and compare to the tree which was built from B and C + smt.insert(key_a, EMPTY_WORD); + + let mut expected_smt = TieredSmt::default(); + expected_smt.insert(key_b, value_b); + expected_smt.insert(key_c, value_c); + assert_eq!(smt, expected_smt); + + // entry B moves to depth 48, entry C stays at depth 48 + assert_eq!(smt.nodes.get_leaf_index(&key_b).0.depth(), 48); + assert_eq!(smt.nodes.get_leaf_index(&key_c).0.depth(), 48); + + // --- delete from bottom tier (promotion to depth 32) -------------------- + + let mut smt = TieredSmt::default(); + smt.insert(key_a, value_a); + smt.insert(key_b, value_b); + + // insert a value with a key which shared the same 16-bit prefix + let raw_c = 0b_01010101_01010101_01111111_11111111_10101010_10101010_11111111_00000000_u64; + let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); + smt.insert(key_c, value_c); + + // delete entry A and compare to the tree which was built from B and C + smt.insert(key_a, EMPTY_WORD); + + let mut expected_smt = TieredSmt::default(); + expected_smt.insert(key_b, value_b); + expected_smt.insert(key_c, value_c); + assert_eq!(smt, expected_smt); + + // entry B moves to depth 32, entry C stays at depth 32 + assert_eq!(smt.nodes.get_leaf_index(&key_b).0.depth(), 32); + assert_eq!(smt.nodes.get_leaf_index(&key_c).0.depth(), 32); + + // --- delete from bottom tier (promotion to depth 16) -------------------- + + let mut smt = TieredSmt::default(); + smt.insert(key_a, value_a); + smt.insert(key_b, value_b); + + // insert a value with a key which shared prefix < 16 bits + let raw_c = 0b_01010101_01010100_11111111_11111111_10101010_10101010_11111111_00000000_u64; + let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]); + smt.insert(key_c, value_c); + + // delete entry A and compare to the tree which was built from B and C + smt.insert(key_a, EMPTY_WORD); + + let mut expected_smt = TieredSmt::default(); + expected_smt.insert(key_b, value_b); + expected_smt.insert(key_c, value_c); + assert_eq!(smt, expected_smt); + + // entry B moves to depth 16, entry C stays at depth 16 + assert_eq!(smt.nodes.get_leaf_index(&key_b).0.depth(), 16); + assert_eq!(smt.nodes.get_leaf_index(&key_c).0.depth(), 16); +} + +#[test] +fn test_order_sensitivity() { + let raw = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000001_u64; + let value = [ONE; WORD_SIZE]; + + let key_1 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]); + let key_2 = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw)]); + + let mut smt_1 = TieredSmt::default(); + + smt_1.insert(key_1, value); + smt_1.insert(key_2, value); + smt_1.insert(key_2, [ZERO; WORD_SIZE]); + + let mut smt_2 = TieredSmt::default(); + smt_2.insert(key_1, value); + + assert_eq!(smt_1.root(), smt_2.root()); +} + // BOTTOM TIER TESTS // ================================================================================================