Merge pull request #174 from 0xPolygonMiden/bobbin-tsmt-proof

Implement ability to generate TSMT proofs

src/merkle/mod.rs

@@ -27,7 +27,7 @@ mod simple_smt;
 pub use simple_smt::SimpleSmt;
 
 mod tiered_smt;
-pub use tiered_smt::TieredSmt;
+pub use tiered_smt::{TieredSmt, TieredSmtProof};
 
 mod mmr;
 pub use mmr::{Mmr, MmrPeaks, MmrProof};

src/merkle/tiered_smt/mod.rs

@@ -2,6 +2,7 @@ use super::{
     BTreeMap, BTreeSet, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex,
     Rpo256, RpoDigest, StarkField, Vec, Word,
 };
+use crate::utils::vec;
 use core::{cmp, ops::Deref};
 
 mod nodes;
@@ -10,6 +11,9 @@ use nodes::NodeStore;
 mod values;
 use values::ValueStore;
 
+mod proof;
+pub use proof::TieredSmtProof;
+
 #[cfg(test)]
 mod tests;
 
@@ -134,6 +138,30 @@ impl TieredSmt {
         }
     }
 
+    /// Returns a proof for a key-value pair defined by the specified key.
+    ///
+    /// The proof can be used to attest membership of this key-value pair in a Tiered Sparse Merkle
+    /// Tree defined by the same root as this tree.
+    pub fn prove(&self, key: RpoDigest) -> TieredSmtProof {
+        let (path, index, leaf_exists) = self.nodes.get_proof(&key);
+
+        let entries = if index.depth() == Self::MAX_DEPTH {
+            match self.values.get_all(index.value()) {
+                Some(entries) => entries,
+                None => vec![(key, Self::EMPTY_VALUE)],
+            }
+        } else if leaf_exists {
+            let entry =
+                self.values.get_first(index_to_prefix(&index)).expect("leaf entry not found");
+            debug_assert_eq!(entry.0, key);
+            vec![*entry]
+        } else {
+            vec![(key, Self::EMPTY_VALUE)]
+        };
+
+        TieredSmtProof::new(path, entries)
+    }
+
     // STATE MUTATORS
     // --------------------------------------------------------------------------------------------
 

src/merkle/tiered_smt/nodes.rs

@@ -86,6 +86,15 @@ impl NodeStore {
         Ok(path.into())
     }
 
+    /// Returns a Merkle path to the node specified by the key together with a flag indicating,
+    /// whether this node is a leaf at depths 16, 32, or 48.
+    pub fn get_proof(&self, key: &RpoDigest) -> (MerklePath, NodeIndex, bool) {
+        let (index, leaf_exists) = self.get_leaf_index(key);
+        let index: NodeIndex = index.into();
+        let path = self.get_path(index).expect("failed to retrieve Merkle path for a node index");
+        (path, index, leaf_exists)
+    }
+
     /// Returns an index at which a leaf node for the specified key should be inserted.
     ///
     /// The second value in the returned tuple is set to true if the node at the returned index

src/merkle/tiered_smt/proof.rs

@@ -0,0 +1,134 @@
+use super::{
+    get_common_prefix_tier_depth, get_key_prefix, hash_bottom_leaf, hash_upper_leaf,
+    EmptySubtreeRoots, LeafNodeIndex, MerklePath, RpoDigest, Vec, Word,
+};
+
+// CONSTANTS
+// ================================================================================================
+
+/// Maximum node depth. This is also the bottom tier of the tree.
+const MAX_DEPTH: u8 = super::TieredSmt::MAX_DEPTH;
+
+/// Value of an empty leaf.
+pub const EMPTY_VALUE: Word = super::TieredSmt::EMPTY_VALUE;
+
+// TIERED SPARSE MERKLE TREE PROOF
+// ================================================================================================
+
+/// A proof which can be used to assert membership (or non-membership) of key-value pairs in a
+/// Tiered Sparse Merkle tree.
+///
+/// The proof consists of a Merkle path and one or more key-value entries which describe the node
+/// located at the base of the path. If the node at the base of the path resolves to [ZERO; 4],
+/// the entries will contain a single item with value set to [ZERO; 4].
+pub struct TieredSmtProof {
+    path: MerklePath,
+    entries: Vec<(RpoDigest, Word)>,
+}
+
+impl TieredSmtProof {
+    // CONSTRUCTOR
+    // --------------------------------------------------------------------------------------------
+    /// Returns a new instance of [TieredSmtProof] instantiated from the specified path and entries.
+    ///
+    /// # Panics
+    /// Panics if:
+    /// - The length of the path is greater than 64.
+    /// - Entries is an empty vector.
+    /// - Entries contains more than 1 item, but the length of the path is not 64.
+    /// - Entries contains more than 1 item, and one of the items has value set to [ZERO; 4].
+    /// - Entries contains multiple items with keys which don't share the same 64-bit prefix.
+    pub fn new(path: MerklePath, entries: Vec<(RpoDigest, Word)>) -> Self {
+        assert!(path.depth() <= MAX_DEPTH);
+        assert!(!entries.is_empty());
+        if entries.len() > 1 {
+            assert!(path.depth() == MAX_DEPTH);
+            let prefix = get_key_prefix(&entries[0].0);
+            for entry in entries.iter().skip(1) {
+                assert_ne!(entry.1, EMPTY_VALUE);
+                assert_eq!(prefix, get_key_prefix(&entry.0));
+            }
+        }
+
+        Self { path, entries }
+    }
+
+    // PROOF VERIFIER
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns true if a Tiered Sparse Merkle tree with the specified root contains the provided
+    /// key-value pair.
+    ///
+    /// 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)) {
+            return false;
+        }
+
+        // make sure the Merkle path resolves to the correct root
+        root == &self.compute_root()
+    }
+
+    // PUBLIC ACCESSORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns the value associated with the specific key according to this proof, or None if
+    /// this proof does not contain a value for the specified key.
+    ///
+    /// A key-value pair generated by using this method should pass the `verify_membership()` check.
+    pub fn get(&self, key: &RpoDigest) -> Option<Word> {
+        if self.is_value_empty() {
+            let common_prefix_tier = get_common_prefix_tier_depth(key, &self.entries[0].0);
+            if common_prefix_tier < self.path.depth() {
+                None
+            } else {
+                Some(EMPTY_VALUE)
+            }
+        } else {
+            self.entries.iter().find(|(k, _)| k == key).map(|(_, value)| *value)
+        }
+    }
+
+    /// Computes the root of a Tiered Sparse Merkle tree to which this proof resolve.
+    pub fn compute_root(&self) -> RpoDigest {
+        let node = self.build_node();
+        let index = LeafNodeIndex::from_key(&self.entries[0].0, self.path.depth());
+        self.path
+            .compute_root(index.value(), node)
+            .expect("failed to compute Merkle path root")
+    }
+
+    // HELPER METHODS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns true if the proof is for an empty value.
+    fn is_value_empty(&self) -> bool {
+        self.entries[0].1 == EMPTY_VALUE
+    }
+
+    /// Converts the entries contained in this proof into a node value for node at the base of the
+    /// path contained in this proof.
+    fn build_node(&self) -> RpoDigest {
+        let depth = self.path.depth();
+        if self.is_value_empty() {
+            EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[depth as usize]
+        } else if depth == MAX_DEPTH {
+            hash_bottom_leaf(&self.entries)
+        } else {
+            let (key, value) = self.entries[0];
+            hash_upper_leaf(key, value, depth)
+        }
+    }
+}

src/merkle/tiered_smt/tests.rs

@@ -1,5 +1,5 @@
 use super::{
-    super::{super::ONE, Felt, MerkleStore, WORD_SIZE, ZERO},
+    super::{super::ONE, empty_roots::EMPTY_WORD, Felt, MerkleStore, WORD_SIZE, ZERO},
     EmptySubtreeRoots, InnerNodeInfo, NodeIndex, Rpo256, RpoDigest, TieredSmt, Vec, Word,
 };
 
@@ -587,6 +587,154 @@ fn tsmt_bottom_tier_two() {
     assert_eq!(leaves.next(), None);
 }
 
+// GET PROOF TESTS
+// ================================================================================================
+
+#[test]
+fn tsmt_get_proof() {
+    let mut smt = TieredSmt::default();
+
+    // --- insert a value into the tree ---------------------------------------
+    let raw_a = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_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 a value with the same 48-bit prefix into the tree -----------
+    let raw_b = 0b_01010101_01010101_11111111_11111111_10110101_10101010_10111100_00000000_u64;
+    let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
+    let value_b = [ONE, ONE, ONE, ZERO];
+    smt.insert(key_b, value_b);
+
+    let smt_alt = smt.clone();
+
+    // --- insert a value with the same 32-bit prefix into the tree -----------
+    let raw_c = 0b_01010101_01010101_11111111_11111111_11111101_10101010_10111100_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);
+
+    // --- insert a value with the same 64-bit prefix as A into the tree ------
+    let raw_d = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000000_u64;
+    let key_d = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_d)]);
+    let value_d = [ONE, ZERO, ZERO, ZERO];
+    smt.insert(key_d, value_d);
+
+    // at this point the tree looks as follows:
+    // - A and D are located in the same node at depth 64.
+    // - B is located at depth 64 and shares the same 48-bit prefix with A and D.
+    // - C is located at depth 48 and shares the same 32-bit prefix with A, B, and D.
+
+    // --- generate proof for key A and test that it verifies correctly -------
+    let proof = smt.prove(key_a);
+    assert!(proof.verify_membership(&key_a, &value_a, &smt.root()));
+
+    assert!(!proof.verify_membership(&key_a, &value_b, &smt.root()));
+    assert!(!proof.verify_membership(&key_a, &EMPTY_WORD, &smt.root()));
+    assert!(!proof.verify_membership(&key_b, &value_a, &smt.root()));
+    assert!(!proof.verify_membership(&key_a, &value_a, &smt_alt.root()));
+
+    assert_eq!(proof.get(&key_a), Some(value_a));
+    assert_eq!(proof.get(&key_b), None);
+
+    // since A and D are stored in the same node, we should be able to use the proof to verify
+    // membership of D
+    assert!(proof.verify_membership(&key_d, &value_d, &smt.root()));
+    assert_eq!(proof.get(&key_d), Some(value_d));
+
+    // --- generate proof for key B and test that it verifies correctly -------
+    let proof = smt.prove(key_b);
+    assert!(proof.verify_membership(&key_b, &value_b, &smt.root()));
+
+    assert!(!proof.verify_membership(&key_b, &value_a, &smt.root()));
+    assert!(!proof.verify_membership(&key_b, &EMPTY_WORD, &smt.root()));
+    assert!(!proof.verify_membership(&key_a, &value_b, &smt.root()));
+    assert!(!proof.verify_membership(&key_b, &value_b, &smt_alt.root()));
+
+    assert_eq!(proof.get(&key_b), Some(value_b));
+    assert_eq!(proof.get(&key_a), None);
+
+    // --- generate proof for key C and test that it verifies correctly -------
+    let proof = smt.prove(key_c);
+    assert!(proof.verify_membership(&key_c, &value_c, &smt.root()));
+
+    assert!(!proof.verify_membership(&key_c, &value_a, &smt.root()));
+    assert!(!proof.verify_membership(&key_c, &EMPTY_WORD, &smt.root()));
+    assert!(!proof.verify_membership(&key_a, &value_c, &smt.root()));
+    assert!(!proof.verify_membership(&key_c, &value_c, &smt_alt.root()));
+
+    assert_eq!(proof.get(&key_c), Some(value_c));
+    assert_eq!(proof.get(&key_b), None);
+
+    // --- generate proof for key D and test that it verifies correctly -------
+    let proof = smt.prove(key_d);
+    assert!(proof.verify_membership(&key_d, &value_d, &smt.root()));
+
+    assert!(!proof.verify_membership(&key_d, &value_b, &smt.root()));
+    assert!(!proof.verify_membership(&key_d, &EMPTY_WORD, &smt.root()));
+    assert!(!proof.verify_membership(&key_b, &value_d, &smt.root()));
+    assert!(!proof.verify_membership(&key_d, &value_d, &smt_alt.root()));
+
+    assert_eq!(proof.get(&key_d), Some(value_d));
+    assert_eq!(proof.get(&key_b), None);
+
+    // since A and D are stored in the same node, we should be able to use the proof to verify
+    // membership of A
+    assert!(proof.verify_membership(&key_a, &value_a, &smt.root()));
+    assert_eq!(proof.get(&key_a), Some(value_a));
+
+    // --- generate proof for an empty key at depth 64 ------------------------
+    // this key has the same 48-bit prefix as A but is different from B
+    let raw = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000011_u64;
+    let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
+
+    let proof = smt.prove(key);
+    assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root()));
+
+    assert!(!proof.verify_membership(&key, &value_a, &smt.root()));
+    assert!(!proof.verify_membership(&key, &EMPTY_WORD, &smt_alt.root()));
+
+    assert_eq!(proof.get(&key), Some(EMPTY_WORD));
+    assert_eq!(proof.get(&key_b), None);
+
+    // the same proof should verify against any key with the same 64-bit prefix
+    let key2 = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw)]);
+    assert!(proof.verify_membership(&key2, &EMPTY_WORD, &smt.root()));
+    assert_eq!(proof.get(&key2), Some(EMPTY_WORD));
+
+    // but verifying if against a key with the same 63-bit prefix (or smaller) should fail
+    let raw3 = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000010_u64;
+    let key3 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw3)]);
+    assert!(!proof.verify_membership(&key3, &EMPTY_WORD, &smt.root()));
+    assert_eq!(proof.get(&key3), None);
+
+    // --- generate proof for an empty key at depth 48 ------------------------
+    // this key has the same 32-prefix as A, B, C, and D, but is different from C
+    let raw = 0b_01010101_01010101_11111111_11111111_00110101_10101010_11111100_00000000_u64;
+    let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
+
+    let proof = smt.prove(key);
+    assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root()));
+
+    assert!(!proof.verify_membership(&key, &value_a, &smt.root()));
+    assert!(!proof.verify_membership(&key, &EMPTY_WORD, &smt_alt.root()));
+
+    assert_eq!(proof.get(&key), Some(EMPTY_WORD));
+    assert_eq!(proof.get(&key_b), None);
+
+    // the same proof should verify against any key with the same 48-bit prefix
+    let raw2 = 0b_01010101_01010101_11111111_11111111_00110101_10101010_01111100_00000000_u64;
+    let key2 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw2)]);
+    assert!(proof.verify_membership(&key2, &EMPTY_WORD, &smt.root()));
+    assert_eq!(proof.get(&key2), Some(EMPTY_WORD));
+
+    // but verifying against a key with the same 47-bit prefix (or smaller) should fail
+    let raw3 = 0b_01010101_01010101_11111111_11111111_00110101_10101011_11111100_00000000_u64;
+    let key3 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw3)]);
+    assert!(!proof.verify_membership(&key3, &EMPTY_WORD, &smt.root()));
+    assert_eq!(proof.get(&key3), None);
+}
+
 // ERROR TESTS
 // ================================================================================================