diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 9d20629..9580973 100644 --- a/src/merkle/mod.rs +++ b/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}; diff --git a/src/merkle/tiered_smt/mod.rs b/src/merkle/tiered_smt/mod.rs index 5379d20..c4f30ed 100644 --- a/src/merkle/tiered_smt/mod.rs +++ b/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 // -------------------------------------------------------------------------------------------- diff --git a/src/merkle/tiered_smt/nodes.rs b/src/merkle/tiered_smt/nodes.rs index 7135c6c..9db4ba3 100644 --- a/src/merkle/tiered_smt/nodes.rs +++ b/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 diff --git a/src/merkle/tiered_smt/proof.rs b/src/merkle/tiered_smt/proof.rs new file mode 100644 index 0000000..3965faa --- /dev/null +++ b/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 { + 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) + } + } +} diff --git a/src/merkle/tiered_smt/tests.rs b/src/merkle/tiered_smt/tests.rs index e459c90..c1b7649 100644 --- a/src/merkle/tiered_smt/tests.rs +++ b/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 // ================================================================================================