Merge pull request #212 from 0xPolygonMiden/hacka-tsmt-bug-fix-empty-…

…value-depth-64

bugfix: TSMT failed to verify empty word for depth 64.

src/merkle/tiered_smt/proof.rs

@@ -85,18 +85,26 @@ impl TieredSmtProof {
     /// Note: this method cannot be used to assert non-membership. That is, if false is returned,
     /// it does not mean that the provided key-value pair is not in the tree.
     pub fn verify_membership(&self, key: &RpoDigest, value: &Word, root: &RpoDigest) -> bool {
-        if self.is_value_empty() {
-            if value != &EMPTY_VALUE {
-                return false;
-            }
-            // if the proof is for an empty value, we can verify it against any key which has a
-            // common prefix with the key storied in entries, but the prefix must be greater than
-            // the path length
-            let common_prefix_tier = get_common_prefix_tier_depth(key, &self.entries[0].0);
-            if common_prefix_tier < self.path.depth() {
-                return false;
-            }
-        } else if !self.entries.contains(&(*key, *value)) {
+        // Handles the following scenarios:
+        // - the value is set
+        // - empty leaf, there is an explicit entry for the key with the empty value
+        // - shared 64-bit prefix, the target key is not included in the entries list, the value is implicitly the empty word
+        let v = match self.entries.iter().find(|(k, _)| k == key) {
+            Some((_, v)) => v,
+            None => &EMPTY_VALUE,
+        };
+
+        // The value must match for the proof to be valid
+        if v != value {
+            return false;
+        }
+
+        // If the proof is for an empty value, we can verify it against any key which has a common
+        // prefix with the key storied in entries, but the prefix must be greater than the path
+        // length
+        if self.is_value_empty()
+            && get_common_prefix_tier_depth(key, &self.entries[0].0) < self.path.depth()
+        {
             return false;
         }
 

src/merkle/tiered_smt/tests.rs

@@ -715,6 +715,38 @@ fn tsmt_bottom_tier_two() {
 // GET PROOF TESTS
 // ================================================================================================
 
+/// Tests the membership and non-membership proof for a single at depth 64
+#[test]
+fn tsmt_get_proof_single_element_64() {
+    let mut smt = TieredSmt::default();
+
+    let raw_a = 0b_00000000_00000001_00000000_00000001_00000000_00000001_00000000_00000001_u64;
+    let key_a = [ONE, ONE, ONE, raw_a.into()].into();
+    let value_a = [ONE, ONE, ONE, ONE];
+    smt.insert(key_a, value_a);
+
+    // push element `a` to depth 64, by inserting another value that shares the 48-bit prefix
+    let raw_b = 0b_00000000_00000001_00000000_00000001_00000000_00000001_00000000_00000000_u64;
+    let key_b = [ONE, ONE, ONE, raw_b.into()].into();
+    smt.insert(key_b, [ONE, ONE, ONE, ONE]);
+
+    // verify the proof for element `a`
+    let proof = smt.prove(key_a);
+    assert!(proof.verify_membership(&key_a, &value_a, &smt.root()));
+
+    // check that a value that is not inserted in the tree produces a valid membership proof for the
+    // empty word
+    let key = [ZERO, ZERO, ZERO, ZERO].into();
+    let proof = smt.prove(key);
+    assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root()));
+
+    // check that a key that shared the 64-bit prefix with `a`, but is not inserted, also has a
+    // valid membership proof for the empty word
+    let key = [ONE, ONE, ZERO, raw_a.into()].into();
+    let proof = smt.prove(key);
+    assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root()));
+}
+
 #[test]
 fn tsmt_get_proof() {
     let mut smt = TieredSmt::default();