From 5c6a20cb6094fba6b2e915c87cf44409b2659b3f Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Fri, 4 Aug 2023 22:36:45 -0700 Subject: [PATCH 1/2] fix: bug in TSMT for depth 64 removal --- src/merkle/tiered_smt/mod.rs | 21 ++++-- src/merkle/tiered_smt/nodes.rs | 57 ++++++++++++--- src/merkle/tiered_smt/tests.rs | 125 +++++++++++++++++++++++++++++++++ 3 files changed, 188 insertions(+), 15 deletions(-) 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 // ================================================================================================ From b3e7578ab228c82936a57ae98e7d62987d4357aa Mon Sep 17 00:00:00 2001 From: Bobbin Threadbare Date: Fri, 4 Aug 2023 22:46:23 -0700 Subject: [PATCH 2/2] fix: misspelled variant name in TieredSmtProofError --- src/merkle/tiered_smt/error.rs | 31 +++++++++++++++---------------- src/merkle/tiered_smt/proof.rs | 2 +- 2 files changed, 16 insertions(+), 17 deletions(-) diff --git a/src/merkle/tiered_smt/error.rs b/src/merkle/tiered_smt/error.rs index fdc6123..92f7dea 100644 --- a/src/merkle/tiered_smt/error.rs +++ b/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" + ) } } } diff --git a/src/merkle/tiered_smt/proof.rs b/src/merkle/tiered_smt/proof.rs index eae8e38..28ac288 100644 --- a/src/merkle/tiered_smt/proof.rs +++ b/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)); } } }