Merge pull request #177 from 0xPolygonMiden/bobbin-tsmt-delete-64

Bug fix in TSMT for depth 64 removal

src/merkle/tiered_smt/error.rs

@@ -3,11 +3,11 @@ use core::fmt::Display;
 #[derive(Debug, PartialEq, Eq)]
 pub enum TieredSmtProofError {
     EntriesEmpty,
-    PathTooLong,
-    NotATierPath(u8),
-    MultipleEntriesOutsideLastTier,
     EmptyValueNotAllowed,
-    UnmatchingPrefixes(u64, u64),
+    MismatchedPrefixes(u64, u64),
+    MultipleEntriesOutsideLastTier,
+    NotATierPath(u8),
+    PathTooLong,
 }
 
 impl Display for TieredSmtProofError {
@@ -16,30 +16,29 @@ impl Display for TieredSmtProofError {
             TieredSmtProofError::EntriesEmpty => {
                 write!(f, "Missing entries for tiered sparse merkle tree proof")
             }
-            TieredSmtProofError::PathTooLong => {
+            TieredSmtProofError::EmptyValueNotAllowed => {
                 write!(
                     f,
-                    "Path longer than maximum depth of 64 for tiered sparse merkle tree proof"
+                    "The empty value [0, 0, 0, 0] is not allowed inside a tiered sparse merkle tree"
                 )
             }
-            TieredSmtProofError::NotATierPath(got) => {
-                write!(
-                    f,
-                    "Path length does not correspond to a tier. Got {} Expected one of 16,32,48,64",
-                    got
-                )
+            TieredSmtProofError::MismatchedPrefixes(first, second) => {
+                write!(f, "Not all leaves have the same prefix. First {first} second {second}")
             }
             TieredSmtProofError::MultipleEntriesOutsideLastTier => {
                 write!(f, "Multiple entries are only allowed for the last tier (depth 64)")
             }
-            TieredSmtProofError::EmptyValueNotAllowed => {
+            TieredSmtProofError::NotATierPath(got) => {
                 write!(
                     f,
-                    "The empty value [0,0,0,0] is not allowed inside a tiered sparse merkle tree"
+                    "Path length does not correspond to a tier. Got {got} Expected one of 16, 32, 48, 64"
                 )
             }
-            TieredSmtProofError::UnmatchingPrefixes(first, second) => {
-                write!(f, "Not all leaves have the same prefix. First {} second {}", first, second)
+            TieredSmtProofError::PathTooLong => {
+                write!(
+                    f,
+                    "Path longer than maximum depth of 64 for tiered sparse merkle tree proof"
+                )
             }
         }
     }

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;
             };

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) {

src/merkle/tiered_smt/proof.rs

@@ -68,7 +68,7 @@ impl TieredSmtProof {
                 }
                 let current = get_key_prefix(&entry.0);
                 if prefix != current {
-                    return Err(TieredSmtProofError::UnmatchingPrefixes(prefix, current));
+                    return Err(TieredSmtProofError::MismatchedPrefixes(prefix, current));
                 }
             }
         }

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
 // ================================================================================================