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        if self.is_value_empty() {
 89            if value != &EMPTY_VALUE {
 90                return false;
 91            }
 92            // if the proof is for an empty value, we can verify it against any key which has a
 93            // common prefix with the key storied in entries, but the prefix must be greater than
 94            // the path length
 95            let common_prefix_tier = get_common_prefix_tier_depth(key, &self.entries[0].0);
 96            if common_prefix_tier < self.path.depth() {
 97                return false;
 98            }
 99        } else if !self.entries.contains(&(*key, *value)) {
100            return false;
101        }
102
103        // make sure the Merkle path resolves to the correct root
104        root == &self.compute_root()
105    }
106
107    // PUBLIC ACCESSORS
108    // --------------------------------------------------------------------------------------------
109
110    /// Returns the value associated with the specific key according to this proof, or None if
111    /// this proof does not contain a value for the specified key.
112    ///
113    /// A key-value pair generated by using this method should pass the `verify_membership()` check.
114    pub fn get(&self, key: &RpoDigest) -> Option<Word> {
115        if self.is_value_empty() {
116            let common_prefix_tier = get_common_prefix_tier_depth(key, &self.entries[0].0);
117            if common_prefix_tier < self.path.depth() {
118                None
119            } else {
120                Some(EMPTY_VALUE)
121            }
122        } else {
123            self.entries.iter().find(|(k, _)| k == key).map(|(_, value)| *value)
124        }
125    }
126
127    /// Computes the root of a Tiered Sparse Merkle tree to which this proof resolve.
128    pub fn compute_root(&self) -> RpoDigest {
129        let node = self.build_node();
130        let index = LeafNodeIndex::from_key(&self.entries[0].0, self.path.depth());
131        self.path
132            .compute_root(index.value(), node)
133            .expect("failed to compute Merkle path root")
134    }
135
136    /// Consume the proof and returns its parts.
137    pub fn into_parts(self) -> (MerklePath, Vec<(RpoDigest, Word)>) {
138        (self.path, self.entries)
139    }
140
141    // HELPER METHODS
142    // --------------------------------------------------------------------------------------------
143
144    /// Returns true if the proof is for an empty value.
145    fn is_value_empty(&self) -> bool {
146        self.entries[0].1 == EMPTY_VALUE
147    }
148
149    /// Converts the entries contained in this proof into a node value for node at the base of the
150    /// path contained in this proof.
151    fn build_node(&self) -> RpoDigest {
152        let depth = self.path.depth();
153        if self.is_value_empty() {
154            EmptySubtreeRoots::empty_hashes(MAX_DEPTH)[depth as usize]
155        } else if depth == MAX_DEPTH {
156            hash_bottom_leaf(&self.entries)
157        } else {
158            let (key, value) = self.entries[0];
159            hash_upper_leaf(key, value, depth)
160        }
161    }
162}