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