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