1use super::{
  2    get_common_prefix_tier_depth, get_key_prefix, hash_bottom_leaf, hash_upper_leaf,
  3    EmptySubtreeRoots, LeafNodeIndex, MerklePath, RpoDigest, TieredSmtProofError, 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/// Depths at which leaves can exist in a tiered SMT.
 16pub const TIER_DEPTHS: [u8; 4] = super::TieredSmt::TIER_DEPTHS;
 17
 18// TIERED SPARSE MERKLE TREE PROOF
 19// ================================================================================================
 20
 21/// A proof which can be used to assert membership (or non-membership) of key-value pairs in a
 22/// Tiered Sparse Merkle tree.
 23///
 24/// The proof consists of a Merkle path and one or more key-value entries which describe the node
 25/// located at the base of the path. If the node at the base of the path resolves to [ZERO; 4],
 26/// the entries will contain a single item with value set to [ZERO; 4].
 27#[derive(PartialEq, Eq, Debug)]
 28pub struct TieredSmtProof {
 29    path: MerklePath,
 30    entries: Vec<(RpoDigest, Word)>,
 31}
 32
 33impl TieredSmtProof {
 34    // CONSTRUCTOR
 35    // --------------------------------------------------------------------------------------------
 36    /// Returns a new instance of [TieredSmtProof] instantiated from the specified path and entries.
 37    ///
 38    /// # Panics
 39    /// Panics if:
 40    /// - The length of the path is greater than 64.
 41    /// - Entries is an empty vector.
 42    /// - Entries contains more than 1 item, but the length of the path is not 64.
 43    /// - Entries contains more than 1 item, and one of the items has value set to [ZERO; 4].
 44    /// - Entries contains multiple items with keys which don't share the same 64-bit prefix.
 45    pub fn new<I>(path: MerklePath, entries: I) -> Result<Self, TieredSmtProofError>
 46    where
 47        I: IntoIterator<Item = (RpoDigest, Word)>,
 48    {
 49        let entries: Vec<(RpoDigest, Word)> = entries.into_iter().collect();
 50
 51        if !TIER_DEPTHS.into_iter().any(|e| e == path.depth()) {
 52            return Err(TieredSmtProofError::NotATierPath(path.depth()));
 53        }
 54
 55        if entries.is_empty() {
 56            return Err(TieredSmtProofError::EntriesEmpty);
 57        }
 58
 59        if entries.len() > 1 {
 60            if path.depth() != MAX_DEPTH {
 61                return Err(TieredSmtProofError::MultipleEntriesOutsideLastTier);
 62            }
 63
 64            let prefix = get_key_prefix(&entries[0].0);
 65            for entry in entries.iter().skip(1) {
 66                if entry.1 == EMPTY_VALUE {
 67                    return Err(TieredSmtProofError::EmptyValueNotAllowed);
 68                }
 69                let current = get_key_prefix(&entry.0);
 70                if prefix != current {
 71                    return Err(TieredSmtProofError::MismatchedPrefixes(prefix, current));
 72                }
 73            }
 74        }
 75
 76        Ok(Self { path, entries })
 77    }
 78
 79    // PROOF VERIFIER
 80    // --------------------------------------------------------------------------------------------
 81
 82    /// Returns true if a Tiered Sparse Merkle tree with the specified root contains the provided
 83    /// key-value pair.
 84    ///
 85    /// Note: this method cannot be used to assert non-membership. That is, if false is returned,
 86    /// it does not mean that the provided key-value pair is not in the tree.
 87    pub fn verify_membership(&self, key: &RpoDigest, value: &Word, root: &RpoDigest) -> bool {
 88        // Handles the following scenarios:
 89        // - the value is set
 90        // - empty leaf, there is an explicit entry for the key with the empty value
 91        // - shared 64-bit prefix, the target key is not included in the entries list, the value is implicitly the empty word
 92        let v = match self.entries.iter().find(|(k, _)| k == key) {
 93            Some((_, v)) => v,
 94            None => &EMPTY_VALUE,
 95        };
 96
 97        // The value must match for the proof to be valid
 98        if v != value {
 99            return false;
100        }
101
102        // If the proof is for an empty value, we can verify it against any key which has a common
103        // prefix with the key storied in entries, but the prefix must be greater than the path
104        // length
105        if self.is_value_empty()
106            && get_common_prefix_tier_depth(key, &self.entries[0].0) < self.path.depth()
107        {
108            return false;
109        }
110
111        // make sure the Merkle path resolves to the correct root
112        root == &self.compute_root()
113    }
114
115    // PUBLIC ACCESSORS
116    // --------------------------------------------------------------------------------------------
117
118    /// Returns the value associated with the specific key according to this proof, or None if
119    /// this proof does not contain a value for the specified key.
120    ///
121    /// A key-value pair generated by using this method should pass the `verify_membership()` check.
122    pub fn get(&self, key: &RpoDigest) -> Option<Word> {
123        if self.is_value_empty() {
124            let common_prefix_tier = get_common_prefix_tier_depth(key, &self.entries[0].0);
125            if common_prefix_tier < self.path.depth() {
126                None
127            } else {
128                Some(EMPTY_VALUE)
129            }
130        } else {
131            self.entries.iter().find(|(k, _)| k == key).map(|(_, value)| *value)
132        }
133    }
134
135    /// Computes the root of a Tiered Sparse Merkle tree to which this proof resolve.
136    pub fn compute_root(&self) -> RpoDigest {
137        let node = self.build_node();
138        let index = LeafNodeIndex::from_key(&self.entries[0].0, self.path.depth());
139        self.path
140            .compute_root(index.value(), node)
141            .expect("failed to compute Merkle path root")
142    }
143
144    /// Consume the proof and returns its parts.
145    pub fn into_parts(self) -> (MerklePath, Vec<(RpoDigest, Word)>) {
146        (self.path, self.entries)
147    }
148
149    // HELPER METHODS
150    // --------------------------------------------------------------------------------------------
151
152    /// Returns true if the proof is for an empty value.
153    fn is_value_empty(&self) -> bool {
154        self.entries[0].1 == EMPTY_VALUE
155    }
156
157    /// Converts the entries contained in this proof into a node value for node at the base of the
158    /// path contained in this proof.
159    fn build_node(&self) -> RpoDigest {
160        let depth = self.path.depth();
161        if self.is_value_empty() {
162            EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[depth as usize]
163        } else if depth == MAX_DEPTH {
164            hash_bottom_leaf(&self.entries)
165        } else {
166            let (key, value) = self.entries[0];
167            hash_upper_leaf(key, value, depth)
168        }
169    }
170}