arnaucube
arnaucube/miden-cryptoPublic Self-hosted repository
miden-crypto/src/merkle/smt/full/mod.rs
src/merkle/smt/full/mod.rsRaw
  1use super::{
  2    EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, MerklePath,
  3    NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD,
  4};
  5use crate::utils::collections::*;
  6
  7mod error;
  8pub use error::{SmtLeafError, SmtProofError};
  9
 10mod leaf;
 11pub use leaf::SmtLeaf;
 12
 13mod proof;
 14pub use proof::SmtProof;
 15
 16#[cfg(test)]
 17mod tests;
 18
 19// CONSTANTS
 20// ================================================================================================
 21
 22pub const SMT_DEPTH: u8 = 64;
 23
 24// SMT
 25// ================================================================================================
 26
 27/// Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and values are represented
 28/// by 4 field elements.
 29///
 30/// All leaves sit at depth 64. The most significant element of the key is used to identify the leaf to
 31/// which the key maps.
 32///
 33/// A leaf is either empty, or holds one or more key-value pairs. An empty leaf hashes to the empty
 34/// word. Otherwise, a leaf hashes to the hash of its key-value pairs, ordered by key first, value
 35/// second.
 36#[derive(Debug, Clone, PartialEq, Eq)]
 37#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
 38pub struct Smt {
 39    root: RpoDigest,
 40    leaves: BTreeMap<u64, SmtLeaf>,
 41    inner_nodes: BTreeMap<NodeIndex, InnerNode>,
 42}
 43
 44impl Smt {
 45    // CONSTANTS
 46    // --------------------------------------------------------------------------------------------
 47    /// The default value used to compute the hash of empty leaves
 48    pub const EMPTY_VALUE: Word = <Self as SparseMerkleTree<SMT_DEPTH>>::EMPTY_VALUE;
 49
 50    // CONSTRUCTORS
 51    // --------------------------------------------------------------------------------------------
 52
 53    /// Returns a new [Smt].
 54    ///
 55    /// All leaves in the returned tree are set to [Self::EMPTY_VALUE].
 56    pub fn new() -> Self {
 57        let root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
 58
 59        Self {
 60            root,
 61            leaves: BTreeMap::new(),
 62            inner_nodes: BTreeMap::new(),
 63        }
 64    }
 65
 66    /// Returns a new [Smt] instantiated with leaves set as specified by the provided entries.
 67    ///
 68    /// All leaves omitted from the entries list are set to [Self::EMPTY_VALUE].
 69    ///
 70    /// # Errors
 71    /// Returns an error if the provided entries contain multiple values for the same key.
 72    pub fn with_entries(
 73        entries: impl IntoIterator<Item = (RpoDigest, Word)>,
 74    ) -> Result<Self, MerkleError> {
 75        // create an empty tree
 76        let mut tree = Self::new();
 77
 78        // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
 79        // entries with the empty value need additional tracking.
 80        let mut key_set_to_zero = BTreeSet::new();
 81
 82        for (key, value) in entries {
 83            let old_value = tree.insert(key, value);
 84
 85            if old_value != EMPTY_WORD || key_set_to_zero.contains(&key) {
 86                return Err(MerkleError::DuplicateValuesForIndex(
 87                    LeafIndex::<SMT_DEPTH>::from(key).value(),
 88                ));
 89            }
 90
 91            if value == EMPTY_WORD {
 92                key_set_to_zero.insert(key);
 93            };
 94        }
 95        Ok(tree)
 96    }
 97
 98    // PUBLIC ACCESSORS
 99    // --------------------------------------------------------------------------------------------
100
101    /// Returns the depth of the tree
102    pub const fn depth(&self) -> u8 {
103        SMT_DEPTH
104    }
105
106    /// Returns the root of the tree
107    pub fn root(&self) -> RpoDigest {
108        <Self as SparseMerkleTree<SMT_DEPTH>>::root(self)
109    }
110
111    /// Returns the leaf to which `key` maps
112    pub fn get_leaf(&self, key: &RpoDigest) -> SmtLeaf {
113        <Self as SparseMerkleTree<SMT_DEPTH>>::get_leaf(self, key)
114    }
115
116    /// Returns the value associated with `key`
117    pub fn get_value(&self, key: &RpoDigest) -> Word {
118        let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
119
120        match self.leaves.get(&leaf_pos) {
121            Some(leaf) => leaf.get_value(key).unwrap_or_default(),
122            None => EMPTY_WORD,
123        }
124    }
125
126    /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
127    /// path to the leaf, as well as the leaf itself.
128    pub fn open(&self, key: &RpoDigest) -> SmtProof {
129        <Self as SparseMerkleTree<SMT_DEPTH>>::open(self, key)
130    }
131
132    // ITERATORS
133    // --------------------------------------------------------------------------------------------
134
135    /// Returns an iterator over the leaves of this [Smt].
136    pub fn leaves(&self) -> impl Iterator<Item = (LeafIndex<SMT_DEPTH>, &SmtLeaf)> {
137        self.leaves
138            .iter()
139            .map(|(leaf_index, leaf)| (LeafIndex::new_max_depth(*leaf_index), leaf))
140    }
141
142    /// Returns an iterator over the key-value pairs of this [Smt].
143    pub fn entries(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
144        self.leaves().flat_map(|(_, leaf)| leaf.entries())
145    }
146
147    /// Returns an iterator over the inner nodes of this [Smt].
148    pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
149        self.inner_nodes.values().map(|e| InnerNodeInfo {
150            value: e.hash(),
151            left: e.left,
152            right: e.right,
153        })
154    }
155
156    // STATE MUTATORS
157    // --------------------------------------------------------------------------------------------
158
159    /// Inserts a value at the specified key, returning the previous value associated with that key.
160    /// Recall that by definition, any key that hasn't been updated is associated with
161    /// [`Self::EMPTY_VALUE`].
162    ///
163    /// This also recomputes all hashes between the leaf (associated with the key) and the root,
164    /// updating the root itself.
165    pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word {
166        <Self as SparseMerkleTree<SMT_DEPTH>>::insert(self, key, value)
167    }
168
169    // HELPERS
170    // --------------------------------------------------------------------------------------------
171
172    /// Inserts `value` at leaf index pointed to by `key`. `value` is guaranteed to not be the empty
173    /// value, such that this is indeed an insertion.
174    fn perform_insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
175        debug_assert_ne!(value, Self::EMPTY_VALUE);
176
177        let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
178
179        match self.leaves.get_mut(&leaf_index.value()) {
180            Some(leaf) => leaf.insert(key, value),
181            None => {
182                self.leaves.insert(leaf_index.value(), SmtLeaf::Single((key, value)));
183
184                None
185            }
186        }
187    }
188
189    /// Removes key-value pair at leaf index pointed to by `key` if it exists.
190    fn perform_remove(&mut self, key: RpoDigest) -> Option<Word> {
191        let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
192
193        if let Some(leaf) = self.leaves.get_mut(&leaf_index.value()) {
194            let (old_value, is_empty) = leaf.remove(key);
195            if is_empty {
196                self.leaves.remove(&leaf_index.value());
197            }
198            old_value
199        } else {
200            // there's nothing stored at the leaf; nothing to update
201            None
202        }
203    }
204}
205
206impl SparseMerkleTree<SMT_DEPTH> for Smt {
207    type Key = RpoDigest;
208    type Value = Word;
209    type Leaf = SmtLeaf;
210    type Opening = SmtProof;
211
212    const EMPTY_VALUE: Self::Value = EMPTY_WORD;
213
214    fn root(&self) -> RpoDigest {
215        self.root
216    }
217
218    fn set_root(&mut self, root: RpoDigest) {
219        self.root = root;
220    }
221
222    fn get_inner_node(&self, index: NodeIndex) -> InnerNode {
223        self.inner_nodes.get(&index).cloned().unwrap_or_else(|| {
224            let node = EmptySubtreeRoots::entry(SMT_DEPTH, index.depth() + 1);
225
226            InnerNode { left: *node, right: *node }
227        })
228    }
229
230    fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) {
231        self.inner_nodes.insert(index, inner_node);
232    }
233
234    fn remove_inner_node(&mut self, index: NodeIndex) {
235        let _ = self.inner_nodes.remove(&index);
236    }
237
238    fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value> {
239        // inserting an `EMPTY_VALUE` is equivalent to removing any value associated with `key`
240        if value != Self::EMPTY_VALUE {
241            self.perform_insert(key, value)
242        } else {
243            self.perform_remove(key)
244        }
245    }
246
247    fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf {
248        let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
249
250        match self.leaves.get(&leaf_pos) {
251            Some(leaf) => leaf.clone(),
252            None => SmtLeaf::new_empty(key.into()),
253        }
254    }
255
256    fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest {
257        leaf.hash()
258    }
259
260    fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex<SMT_DEPTH> {
261        let most_significant_felt = key[3];
262        LeafIndex::new_max_depth(most_significant_felt.as_int())
263    }
264
265    fn path_and_leaf_to_opening(path: MerklePath, leaf: SmtLeaf) -> SmtProof {
266        SmtProof::new_unchecked(path, leaf)
267    }
268}
269
270impl Default for Smt {
271    fn default() -> Self {
272        Self::new()
273    }
274}
275
276// CONVERSIONS
277// ================================================================================================
278
279impl From<Word> for LeafIndex<SMT_DEPTH> {
280    fn from(value: Word) -> Self {
281        // We use the most significant `Felt` of a `Word` as the leaf index.
282        Self::new_max_depth(value[3].as_int())
283    }
284}
285
286impl From<RpoDigest> for LeafIndex<SMT_DEPTH> {
287    fn from(value: RpoDigest) -> Self {
288        Word::from(value).into()
289    }
290}
291
292impl From<&RpoDigest> for LeafIndex<SMT_DEPTH> {
293    fn from(value: &RpoDigest) -> Self {
294        Word::from(value).into()
295    }
296}