1use super::{
  2    get_common_prefix_tier_depth, get_key_prefix, hash_bottom_leaf, hash_upper_leaf,
  3    EmptySubtreeRoots, LeafNodeIndex, MerklePath, RpoDigest, Vec, Word,
  4};
  5
  6// CONSTANTS
  7// ================================================================================================
  8
  9/// Maximum node depth. This is also the bottom tier of the tree.
 10const MAX_DEPTH: u8 = super::TieredSmt::MAX_DEPTH;
 11
 12/// Value of an empty leaf.
 13pub const EMPTY_VALUE: Word = super::TieredSmt::EMPTY_VALUE;
 14
 15// TIERED SPARSE MERKLE TREE PROOF
 16// ================================================================================================
 17
 18/// A proof which can be used to assert membership (or non-membership) of key-value pairs in a
 19/// Tiered Sparse Merkle tree.
 20///
 21/// The proof consists of a Merkle path and one or more key-value entries which describe the node
 22/// located at the base of the path. If the node at the base of the path resolves to [ZERO; 4],
 23/// the entries will contain a single item with value set to [ZERO; 4].
 24pub struct TieredSmtProof {
 25    path: MerklePath,
 26    entries: Vec<(RpoDigest, Word)>,
 27}
 28
 29impl TieredSmtProof {
 30    // CONSTRUCTOR
 31    // --------------------------------------------------------------------------------------------
 32    /// Returns a new instance of [TieredSmtProof] instantiated from the specified path and entries.
 33    ///
 34    /// # Panics
 35    /// Panics if:
 36    /// - The length of the path is greater than 64.
 37    /// - Entries is an empty vector.
 38    /// - Entries contains more than 1 item, but the length of the path is not 64.
 39    /// - Entries contains more than 1 item, and one of the items has value set to [ZERO; 4].
 40    /// - Entries contains multiple items with keys which don't share the same 64-bit prefix.
 41    pub fn new(path: MerklePath, entries: Vec<(RpoDigest, Word)>) -> Self {
 42        assert!(path.depth() <= MAX_DEPTH);
 43        assert!(!entries.is_empty());
 44        if entries.len() > 1 {
 45            assert!(path.depth() == MAX_DEPTH);
 46            let prefix = get_key_prefix(&entries[0].0);
 47            for entry in entries.iter().skip(1) {
 48                assert_ne!(entry.1, EMPTY_VALUE);
 49                assert_eq!(prefix, get_key_prefix(&entry.0));
 50            }
 51        }
 52
 53        Self { path, entries }
 54    }
 55
 56    // PROOF VERIFIER
 57    // --------------------------------------------------------------------------------------------
 58
 59    /// Returns true if a Tiered Sparse Merkle tree with the specified root contains the provided
 60    /// key-value pair.
 61    ///
 62    /// Note: this method cannot be used to assert non-membership. That is, if false is returned,
 63    /// it does not mean that the provided key-value pair is not in the tree.
 64    pub fn verify_membership(&self, key: &RpoDigest, value: &Word, root: &RpoDigest) -> bool {
 65        if self.is_value_empty() {
 66            if value != &EMPTY_VALUE {
 67                return false;
 68            }
 69            // if the proof is for an empty value, we can verify it against any key which has a
 70            // common prefix with the key storied in entries, but the prefix must be greater than
 71            // the path length
 72            let common_prefix_tier = get_common_prefix_tier_depth(key, &self.entries[0].0);
 73            if common_prefix_tier < self.path.depth() {
 74                return false;
 75            }
 76        } else if !self.entries.contains(&(*key, *value)) {
 77            return false;
 78        }
 79
 80        // make sure the Merkle path resolves to the correct root
 81        root == &self.compute_root()
 82    }
 83
 84    // PUBLIC ACCESSORS
 85    // --------------------------------------------------------------------------------------------
 86
 87    /// Returns the value associated with the specific key according to this proof, or None if
 88    /// this proof does not contain a value for the specified key.
 89    ///
 90    /// A key-value pair generated by using this method should pass the `verify_membership()` check.
 91    pub fn get(&self, key: &RpoDigest) -> Option<Word> {
 92        if self.is_value_empty() {
 93            let common_prefix_tier = get_common_prefix_tier_depth(key, &self.entries[0].0);
 94            if common_prefix_tier < self.path.depth() {
 95                None
 96            } else {
 97                Some(EMPTY_VALUE)
 98            }
 99        } else {
100            self.entries.iter().find(|(k, _)| k == key).map(|(_, value)| *value)
101        }
102    }
103
104    /// Computes the root of a Tiered Sparse Merkle tree to which this proof resolve.
105    pub fn compute_root(&self) -> RpoDigest {
106        let node = self.build_node();
107        let index = LeafNodeIndex::from_key(&self.entries[0].0, self.path.depth());
108        self.path
109            .compute_root(index.value(), node)
110            .expect("failed to compute Merkle path root")
111    }
112
113    // HELPER METHODS
114    // --------------------------------------------------------------------------------------------
115
116    /// Returns true if the proof is for an empty value.
117    fn is_value_empty(&self) -> bool {
118        self.entries[0].1 == EMPTY_VALUE
119    }
120
121    /// Converts the entries contained in this proof into a node value for node at the base of the
122    /// path contained in this proof.
123    fn build_node(&self) -> RpoDigest {
124        let depth = self.path.depth();
125        if self.is_value_empty() {
126            EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[depth as usize]
127        } else if depth == MAX_DEPTH {
128            hash_bottom_leaf(&self.entries)
129        } else {
130            let (key, value) = self.entries[0];
131            hash_upper_leaf(key, value, depth)
132        }
133    }
134}