miden-crypto/src/merkle/smt/simple/mod.rs
src/merkle/smt/simple/mod.rsRaw
1use alloc::collections::{BTreeMap, BTreeSet};
2
3use super::{
4 super::ValuePath, EmptySubtreeRoots, InnerNode, InnerNodeInfo, LeafIndex, MerkleError,
5 MerklePath, MutationSet, NodeIndex, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD,
6 SMT_MAX_DEPTH, SMT_MIN_DEPTH,
7};
8
9#[cfg(test)]
10mod tests;
11
12// SPARSE MERKLE TREE
13// ================================================================================================
14
15/// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction.
16///
17/// The root of the tree is recomputed on each new leaf update.
18#[derive(Debug, Clone, PartialEq, Eq)]
19#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
20pub struct SimpleSmt<const DEPTH: u8> {
21 root: RpoDigest,
22 leaves: BTreeMap<u64, Word>,
23 inner_nodes: BTreeMap<NodeIndex, InnerNode>,
24}
25
26impl<const DEPTH: u8> SimpleSmt<DEPTH> {
27 // CONSTANTS
28 // --------------------------------------------------------------------------------------------
29
30 /// The default value used to compute the hash of empty leaves
31 pub const EMPTY_VALUE: Word = <Self as SparseMerkleTree<DEPTH>>::EMPTY_VALUE;
32
33 // CONSTRUCTORS
34 // --------------------------------------------------------------------------------------------
35
36 /// Returns a new [SimpleSmt].
37 ///
38 /// All leaves in the returned tree are set to [ZERO; 4].
39 ///
40 /// # Errors
41 /// Returns an error if DEPTH is 0 or is greater than 64.
42 pub fn new() -> Result<Self, MerkleError> {
43 // validate the range of the depth.
44 if DEPTH < SMT_MIN_DEPTH {
45 return Err(MerkleError::DepthTooSmall(DEPTH));
46 } else if SMT_MAX_DEPTH < DEPTH {
47 return Err(MerkleError::DepthTooBig(DEPTH as u64));
48 }
49
50 let root = *EmptySubtreeRoots::entry(DEPTH, 0);
51
52 Ok(Self {
53 root,
54 leaves: BTreeMap::new(),
55 inner_nodes: BTreeMap::new(),
56 })
57 }
58
59 /// Returns a new [SimpleSmt] instantiated with leaves set as specified by the provided entries.
60 ///
61 /// All leaves omitted from the entries list are set to [ZERO; 4].
62 ///
63 /// # Errors
64 /// Returns an error if:
65 /// - If the depth is 0 or is greater than 64.
66 /// - The number of entries exceeds the maximum tree capacity, that is 2^{depth}.
67 /// - The provided entries contain multiple values for the same key.
68 pub fn with_leaves(
69 entries: impl IntoIterator<Item = (u64, Word)>,
70 ) -> Result<Self, MerkleError> {
71 // create an empty tree
72 let mut tree = Self::new()?;
73
74 // compute the max number of entries. We use an upper bound of depth 63 because we consider
75 // passing in a vector of size 2^64 infeasible.
76 let max_num_entries = 2_usize.pow(DEPTH.min(63).into());
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 (idx, (key, value)) in entries.into_iter().enumerate() {
83 if idx >= max_num_entries {
84 return Err(MerkleError::TooManyEntries(max_num_entries));
85 }
86
87 let old_value = tree.insert(LeafIndex::<DEPTH>::new(key)?, value);
88
89 if old_value != Self::EMPTY_VALUE || key_set_to_zero.contains(&key) {
90 return Err(MerkleError::DuplicateValuesForIndex(key));
91 }
92
93 if value == Self::EMPTY_VALUE {
94 key_set_to_zero.insert(key);
95 };
96 }
97 Ok(tree)
98 }
99
100 /// Wrapper around [`SimpleSmt::with_leaves`] which inserts leaves at contiguous indices
101 /// starting at index 0.
102 pub fn with_contiguous_leaves(
103 entries: impl IntoIterator<Item = Word>,
104 ) -> Result<Self, MerkleError> {
105 Self::with_leaves(
106 entries
107 .into_iter()
108 .enumerate()
109 .map(|(idx, word)| (idx.try_into().expect("tree max depth is 2^8"), word)),
110 )
111 }
112
113 // PUBLIC ACCESSORS
114 // --------------------------------------------------------------------------------------------
115
116 /// Returns the depth of the tree
117 pub const fn depth(&self) -> u8 {
118 DEPTH
119 }
120
121 /// Returns the root of the tree
122 pub fn root(&self) -> RpoDigest {
123 <Self as SparseMerkleTree<DEPTH>>::root(self)
124 }
125
126 /// Returns the number of non-empty leaves in this tree.
127 pub fn num_leaves(&self) -> usize {
128 self.leaves.len()
129 }
130
131 /// Returns the leaf at the specified index.
132 pub fn get_leaf(&self, key: &LeafIndex<DEPTH>) -> Word {
133 <Self as SparseMerkleTree<DEPTH>>::get_leaf(self, key)
134 }
135
136 /// Returns a node at the specified index.
137 ///
138 /// # Errors
139 /// Returns an error if the specified index has depth set to 0 or the depth is greater than
140 /// the depth of this Merkle tree.
141 pub fn get_node(&self, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
142 if index.is_root() {
143 Err(MerkleError::DepthTooSmall(index.depth()))
144 } else if index.depth() > DEPTH {
145 Err(MerkleError::DepthTooBig(index.depth() as u64))
146 } else if index.depth() == DEPTH {
147 let leaf = self.get_leaf(&LeafIndex::<DEPTH>::try_from(index)?);
148
149 Ok(leaf.into())
150 } else {
151 Ok(self.get_inner_node(index).hash())
152 }
153 }
154
155 /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
156 /// path to the leaf, as well as the leaf itself.
157 pub fn open(&self, key: &LeafIndex<DEPTH>) -> ValuePath {
158 <Self as SparseMerkleTree<DEPTH>>::open(self, key)
159 }
160
161 /// Returns a boolean value indicating whether the SMT is empty.
162 pub fn is_empty(&self) -> bool {
163 debug_assert_eq!(self.leaves.is_empty(), self.root == Self::EMPTY_ROOT);
164 self.root == Self::EMPTY_ROOT
165 }
166
167 // ITERATORS
168 // --------------------------------------------------------------------------------------------
169
170 /// Returns an iterator over the leaves of this [SimpleSmt].
171 pub fn leaves(&self) -> impl Iterator<Item = (u64, &Word)> {
172 self.leaves.iter().map(|(i, w)| (*i, w))
173 }
174
175 /// Returns an iterator over the inner nodes of this [SimpleSmt].
176 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
177 self.inner_nodes.values().map(|e| InnerNodeInfo {
178 value: e.hash(),
179 left: e.left,
180 right: e.right,
181 })
182 }
183
184 // STATE MUTATORS
185 // --------------------------------------------------------------------------------------------
186
187 /// Inserts a value at the specified key, returning the previous value associated with that key.
188 /// Recall that by definition, any key that hasn't been updated is associated with
189 /// [`EMPTY_WORD`].
190 ///
191 /// This also recomputes all hashes between the leaf (associated with the key) and the root,
192 /// updating the root itself.
193 pub fn insert(&mut self, key: LeafIndex<DEPTH>, value: Word) -> Word {
194 <Self as SparseMerkleTree<DEPTH>>::insert(self, key, value)
195 }
196
197 /// Computes what changes are necessary to insert the specified key-value pairs into this
198 /// Merkle tree, allowing for validation before applying those changes.
199 ///
200 /// This method returns a [`MutationSet`], which contains all the information for inserting
201 /// `kv_pairs` into this Merkle tree already calculated, including the new root hash, which can
202 /// be queried with [`MutationSet::root()`]. Once a mutation set is returned,
203 /// [`SimpleSmt::apply_mutations()`] can be called in order to commit these changes to the
204 /// Merkle tree, or [`drop()`] to discard them.
205 ///
206 /// # Example
207 /// ```
208 /// # use miden_crypto::{hash::rpo::RpoDigest, Felt, Word};
209 /// # use miden_crypto::merkle::{LeafIndex, SimpleSmt, EmptySubtreeRoots, SMT_DEPTH};
210 /// let mut smt: SimpleSmt<3> = SimpleSmt::new().unwrap();
211 /// let pair = (LeafIndex::default(), Word::default());
212 /// let mutations = smt.compute_mutations(vec![pair]);
213 /// assert_eq!(mutations.root(), *EmptySubtreeRoots::entry(3, 0));
214 /// smt.apply_mutations(mutations);
215 /// assert_eq!(smt.root(), *EmptySubtreeRoots::entry(3, 0));
216 /// ```
217 pub fn compute_mutations(
218 &self,
219 kv_pairs: impl IntoIterator<Item = (LeafIndex<DEPTH>, Word)>,
220 ) -> MutationSet<DEPTH, LeafIndex<DEPTH>, Word> {
221 <Self as SparseMerkleTree<DEPTH>>::compute_mutations(self, kv_pairs)
222 }
223
224 /// Apply the prospective mutations computed with [`SimpleSmt::compute_mutations()`] to this
225 /// tree.
226 ///
227 /// # Errors
228 /// If `mutations` was computed on a tree with a different root than this one, returns
229 /// [`MerkleError::ConflictingRoots`] with a two-item [`alloc::vec::Vec`]. The first item is the
230 /// root hash the `mutations` were computed against, and the second item is the actual
231 /// current root of this tree.
232 pub fn apply_mutations(
233 &mut self,
234 mutations: MutationSet<DEPTH, LeafIndex<DEPTH>, Word>,
235 ) -> Result<(), MerkleError> {
236 <Self as SparseMerkleTree<DEPTH>>::apply_mutations(self, mutations)
237 }
238
239 /// Inserts a subtree at the specified index. The depth at which the subtree is inserted is
240 /// computed as `DEPTH - SUBTREE_DEPTH`.
241 ///
242 /// Returns the new root.
243 pub fn set_subtree<const SUBTREE_DEPTH: u8>(
244 &mut self,
245 subtree_insertion_index: u64,
246 subtree: SimpleSmt<SUBTREE_DEPTH>,
247 ) -> Result<RpoDigest, MerkleError> {
248 if SUBTREE_DEPTH > DEPTH {
249 return Err(MerkleError::SubtreeDepthExceedsDepth {
250 subtree_depth: SUBTREE_DEPTH,
251 tree_depth: DEPTH,
252 });
253 }
254
255 // Verify that `subtree_insertion_index` is valid.
256 let subtree_root_insertion_depth = DEPTH - SUBTREE_DEPTH;
257 let subtree_root_index =
258 NodeIndex::new(subtree_root_insertion_depth, subtree_insertion_index)?;
259
260 // add leaves
261 // --------------
262
263 // The subtree's leaf indices live in their own context - i.e. a subtree of depth `d`. If we
264 // insert the subtree at `subtree_insertion_index = 0`, then the subtree leaf indices are
265 // valid as they are. However, consider what happens when we insert at
266 // `subtree_insertion_index = 1`. The first leaf of our subtree now will have index `2^d`;
267 // you can see it as there's a full subtree sitting on its left. In general, for
268 // `subtree_insertion_index = i`, there are `i` subtrees sitting before the subtree we want
269 // to insert, so we need to adjust all its leaves by `i * 2^d`.
270 let leaf_index_shift: u64 = subtree_insertion_index * 2_u64.pow(SUBTREE_DEPTH.into());
271 for (subtree_leaf_idx, leaf_value) in subtree.leaves() {
272 let new_leaf_idx = leaf_index_shift + subtree_leaf_idx;
273 debug_assert!(new_leaf_idx < 2_u64.pow(DEPTH.into()));
274
275 self.leaves.insert(new_leaf_idx, *leaf_value);
276 }
277
278 // add subtree's branch nodes (which includes the root)
279 // --------------
280 for (branch_idx, branch_node) in subtree.inner_nodes {
281 let new_branch_idx = {
282 let new_depth = subtree_root_insertion_depth + branch_idx.depth();
283 let new_value = subtree_insertion_index * 2_u64.pow(branch_idx.depth().into())
284 + branch_idx.value();
285
286 NodeIndex::new(new_depth, new_value).expect("index guaranteed to be valid")
287 };
288
289 self.inner_nodes.insert(new_branch_idx, branch_node);
290 }
291
292 // recompute nodes starting from subtree root
293 // --------------
294 self.recompute_nodes_from_index_to_root(subtree_root_index, subtree.root);
295
296 Ok(self.root)
297 }
298}
299
300impl<const DEPTH: u8> SparseMerkleTree<DEPTH> for SimpleSmt<DEPTH> {
301 type Key = LeafIndex<DEPTH>;
302 type Value = Word;
303 type Leaf = Word;
304 type Opening = ValuePath;
305
306 const EMPTY_VALUE: Self::Value = EMPTY_WORD;
307 const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(DEPTH, 0);
308
309 fn root(&self) -> RpoDigest {
310 self.root
311 }
312
313 fn set_root(&mut self, root: RpoDigest) {
314 self.root = root;
315 }
316
317 fn get_inner_node(&self, index: NodeIndex) -> InnerNode {
318 self.inner_nodes
319 .get(&index)
320 .cloned()
321 .unwrap_or_else(|| EmptySubtreeRoots::get_inner_node(DEPTH, index.depth()))
322 }
323
324 fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) {
325 self.inner_nodes.insert(index, inner_node);
326 }
327
328 fn remove_inner_node(&mut self, index: NodeIndex) {
329 let _ = self.inner_nodes.remove(&index);
330 }
331
332 fn insert_value(&mut self, key: LeafIndex<DEPTH>, value: Word) -> Option<Word> {
333 if value == Self::EMPTY_VALUE {
334 self.leaves.remove(&key.value())
335 } else {
336 self.leaves.insert(key.value(), value)
337 }
338 }
339
340 fn get_value(&self, key: &LeafIndex<DEPTH>) -> Word {
341 self.get_leaf(key)
342 }
343
344 fn get_leaf(&self, key: &LeafIndex<DEPTH>) -> Word {
345 let leaf_pos = key.value();
346 match self.leaves.get(&leaf_pos) {
347 Some(word) => *word,
348 None => Self::EMPTY_VALUE,
349 }
350 }
351
352 fn hash_leaf(leaf: &Word) -> RpoDigest {
353 // `SimpleSmt` takes the leaf value itself as the hash
354 leaf.into()
355 }
356
357 fn construct_prospective_leaf(
358 &self,
359 _existing_leaf: Word,
360 _key: &LeafIndex<DEPTH>,
361 value: &Word,
362 ) -> Word {
363 *value
364 }
365
366 fn key_to_leaf_index(key: &LeafIndex<DEPTH>) -> LeafIndex<DEPTH> {
367 *key
368 }
369
370 fn path_and_leaf_to_opening(path: MerklePath, leaf: Word) -> ValuePath {
371 (path, leaf).into()
372 }
373}