miden-crypto/src/merkle/store/mod.rs
src/merkle/store/mod.rsRaw
1use super::{
2 empty_roots::EMPTY_WORD, mmr::Mmr, BTreeMap, EmptySubtreeRoots, InnerNodeInfo, KvMap,
3 MerkleError, MerklePath, MerklePathSet, MerkleStoreDelta, MerkleTree, NodeIndex, RecordingMap,
4 RootPath, Rpo256, RpoDigest, SimpleSmt, TieredSmt, TryApplyDiff, ValuePath, Vec,
5};
6use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
7use core::borrow::Borrow;
8
9#[cfg(test)]
10mod tests;
11
12// MERKLE STORE
13// ================================================================================================
14
15/// A default [MerkleStore] which uses a simple [BTreeMap] as the backing storage.
16pub type DefaultMerkleStore = MerkleStore<BTreeMap<RpoDigest, StoreNode>>;
17
18/// A [MerkleStore] with recording capabilities which uses [RecordingMap] as the backing storage.
19pub type RecordingMerkleStore = MerkleStore<RecordingMap<RpoDigest, StoreNode>>;
20
21#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
22pub struct StoreNode {
23 left: RpoDigest,
24 right: RpoDigest,
25}
26
27/// An in-memory data store for Merkelized data.
28///
29/// This is a in memory data store for Merkle trees, this store allows all the nodes of multiple
30/// trees to live as long as necessary and without duplication, this allows the implementation of
31/// space efficient persistent data structures.
32///
33/// Example usage:
34///
35/// ```rust
36/// # use miden_crypto::{ZERO, Felt, Word};
37/// # use miden_crypto::merkle::{NodeIndex, MerkleStore, MerkleTree};
38/// # use miden_crypto::hash::rpo::Rpo256;
39/// # const fn int_to_node(value: u64) -> Word {
40/// # [Felt::new(value), ZERO, ZERO, ZERO]
41/// # }
42/// # let A = int_to_node(1);
43/// # let B = int_to_node(2);
44/// # let C = int_to_node(3);
45/// # let D = int_to_node(4);
46/// # let E = int_to_node(5);
47/// # let F = int_to_node(6);
48/// # let G = int_to_node(7);
49/// # let H0 = int_to_node(8);
50/// # let H1 = int_to_node(9);
51/// # let T0 = MerkleTree::new([A, B, C, D, E, F, G, H0].to_vec()).expect("even number of leaves provided");
52/// # let T1 = MerkleTree::new([A, B, C, D, E, F, G, H1].to_vec()).expect("even number of leaves provided");
53/// # let ROOT0 = T0.root();
54/// # let ROOT1 = T1.root();
55/// let mut store: MerkleStore = MerkleStore::new();
56///
57/// // the store is initialized with the SMT empty nodes
58/// assert_eq!(store.num_internal_nodes(), 255);
59///
60/// let tree1 = MerkleTree::new(vec![A, B, C, D, E, F, G, H0]).unwrap();
61/// let tree2 = MerkleTree::new(vec![A, B, C, D, E, F, G, H1]).unwrap();
62///
63/// // populates the store with two merkle trees, common nodes are shared
64/// store.extend(tree1.inner_nodes());
65/// store.extend(tree2.inner_nodes());
66///
67/// // every leaf except the last are the same
68/// for i in 0..7 {
69/// let idx0 = NodeIndex::new(3, i).unwrap();
70/// let d0 = store.get_node(ROOT0, idx0).unwrap();
71/// let idx1 = NodeIndex::new(3, i).unwrap();
72/// let d1 = store.get_node(ROOT1, idx1).unwrap();
73/// assert_eq!(d0, d1, "Both trees have the same leaf at pos {i}");
74/// }
75///
76/// // The leafs A-B-C-D are the same for both trees, so are their 2 immediate parents
77/// for i in 0..4 {
78/// let idx0 = NodeIndex::new(3, i).unwrap();
79/// let d0 = store.get_path(ROOT0, idx0).unwrap();
80/// let idx1 = NodeIndex::new(3, i).unwrap();
81/// let d1 = store.get_path(ROOT1, idx1).unwrap();
82/// assert_eq!(d0.path[0..2], d1.path[0..2], "Both sub-trees are equal up to two levels");
83/// }
84///
85/// // Common internal nodes are shared, the two added trees have a total of 30, but the store has
86/// // only 10 new entries, corresponding to the 10 unique internal nodes of these trees.
87/// assert_eq!(store.num_internal_nodes() - 255, 10);
88/// ```
89#[derive(Debug, Clone, Eq, PartialEq)]
90pub struct MerkleStore<T: KvMap<RpoDigest, StoreNode> = BTreeMap<RpoDigest, StoreNode>> {
91 nodes: T,
92}
93
94impl<T: KvMap<RpoDigest, StoreNode>> Default for MerkleStore<T> {
95 fn default() -> Self {
96 Self::new()
97 }
98}
99
100impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
101 // CONSTRUCTORS
102 // --------------------------------------------------------------------------------------------
103
104 /// Creates an empty `MerkleStore` instance.
105 pub fn new() -> MerkleStore<T> {
106 // pre-populate the store with the empty hashes
107 let nodes = empty_hashes().into_iter().collect();
108 MerkleStore { nodes }
109 }
110
111 // PUBLIC ACCESSORS
112 // --------------------------------------------------------------------------------------------
113
114 /// Return a count of the non-leaf nodes in the store.
115 pub fn num_internal_nodes(&self) -> usize {
116 self.nodes.len()
117 }
118
119 /// Returns the node at `index` rooted on the tree `root`.
120 ///
121 /// # Errors
122 /// This method can return the following errors:
123 /// - `RootNotInStore` if the `root` is not present in the store.
124 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in
125 /// the store.
126 pub fn get_node(&self, root: RpoDigest, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
127 let mut hash = root;
128
129 // corner case: check the root is in the store when called with index `NodeIndex::root()`
130 self.nodes.get(&hash).ok_or(MerkleError::RootNotInStore(hash))?;
131
132 for i in (0..index.depth()).rev() {
133 let node = self.nodes.get(&hash).ok_or(MerkleError::NodeNotInStore(hash, index))?;
134
135 let bit = (index.value() >> i) & 1;
136 hash = if bit == 0 { node.left } else { node.right }
137 }
138
139 Ok(hash)
140 }
141
142 /// Returns the node at the specified `index` and its opening to the `root`.
143 ///
144 /// The path starts at the sibling of the target leaf.
145 ///
146 /// # Errors
147 /// This method can return the following errors:
148 /// - `RootNotInStore` if the `root` is not present in the store.
149 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in
150 /// the store.
151 pub fn get_path(&self, root: RpoDigest, index: NodeIndex) -> Result<ValuePath, MerkleError> {
152 let mut hash = root;
153 let mut path = Vec::with_capacity(index.depth().into());
154
155 // corner case: check the root is in the store when called with index `NodeIndex::root()`
156 self.nodes.get(&hash).ok_or(MerkleError::RootNotInStore(hash))?;
157
158 for i in (0..index.depth()).rev() {
159 let node = self.nodes.get(&hash).ok_or(MerkleError::NodeNotInStore(hash, index))?;
160
161 let bit = (index.value() >> i) & 1;
162 hash = if bit == 0 {
163 path.push(node.right);
164 node.left
165 } else {
166 path.push(node.left);
167 node.right
168 }
169 }
170
171 // the path is computed from root to leaf, so it must be reversed
172 path.reverse();
173
174 Ok(ValuePath {
175 value: hash,
176 path: MerklePath::new(path),
177 })
178 }
179
180 /// Reconstructs a path from the root until a leaf or empty node and returns its depth.
181 ///
182 /// The `tree_depth` parameter defines up to which depth the tree will be traversed, starting
183 /// from `root`. The maximum value the argument accepts is [u64::BITS].
184 ///
185 /// The traversed path from leaf to root will start at the least significant bit of `index`,
186 /// and will be executed for `tree_depth` bits.
187 ///
188 /// # Errors
189 /// Will return an error if:
190 /// - The provided root is not found.
191 /// - The path from the root continues to a depth greater than `tree_depth`.
192 /// - The provided `tree_depth` is greater than `64.
193 /// - The provided `index` is not valid for a depth equivalent to `tree_depth`. For more
194 /// information, check [NodeIndex::new].
195 pub fn get_leaf_depth(
196 &self,
197 root: RpoDigest,
198 tree_depth: u8,
199 index: u64,
200 ) -> Result<u8, MerkleError> {
201 // validate depth and index
202 if tree_depth > 64 {
203 return Err(MerkleError::DepthTooBig(tree_depth as u64));
204 }
205 NodeIndex::new(tree_depth, index)?;
206
207 // it's not illegal to have a maximum depth of `0`; we should just return the root in that
208 // case. this check will simplify the implementation as we could overflow bits for depth
209 // `0`.
210 if tree_depth == 0 {
211 return Ok(0);
212 }
213
214 // check if the root exists, providing the proper error report if it doesn't
215 let empty = EmptySubtreeRoots::empty_hashes(tree_depth);
216 let mut hash = root;
217 if !self.nodes.contains_key(&hash) {
218 return Err(MerkleError::RootNotInStore(hash));
219 }
220
221 // we traverse from root to leaf, so the path is reversed
222 let mut path = (index << (64 - tree_depth)).reverse_bits();
223
224 // iterate every depth and reconstruct the path from root to leaf
225 for depth in 0..tree_depth {
226 // we short-circuit if an empty node has been found
227 if hash == empty[depth as usize] {
228 return Ok(depth);
229 }
230
231 // fetch the children pair, mapped by its parent hash
232 let children = match self.nodes.get(&hash) {
233 Some(node) => node,
234 None => return Ok(depth),
235 };
236
237 // traverse down
238 hash = if path & 1 == 0 { children.left } else { children.right };
239 path >>= 1;
240 }
241
242 // at max depth assert it doesn't have sub-trees
243 if self.nodes.contains_key(&hash) {
244 return Err(MerkleError::DepthTooBig(tree_depth as u64 + 1));
245 }
246
247 // depleted bits; return max depth
248 Ok(tree_depth)
249 }
250
251 // DATA EXTRACTORS
252 // --------------------------------------------------------------------------------------------
253
254 /// Returns a subset of this Merkle store such that the returned Merkle store contains all
255 /// nodes which are descendants of the specified roots.
256 ///
257 /// The roots for which no descendants exist in this Merkle store are ignored.
258 pub fn subset<I, R>(&self, roots: I) -> MerkleStore<T>
259 where
260 I: Iterator<Item = R>,
261 R: Borrow<RpoDigest>,
262 {
263 let mut store = MerkleStore::new();
264 for root in roots {
265 let root = *root.borrow();
266 store.clone_tree_from(root, self);
267 }
268 store
269 }
270
271 /// Iterator over the inner nodes of the [MerkleStore].
272 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
273 self.nodes.iter().map(|(r, n)| InnerNodeInfo {
274 value: *r,
275 left: n.left,
276 right: n.right,
277 })
278 }
279
280 /// Iterator over the non-empty leaves of the Merkle tree associated with the specified `root`
281 /// and `max_depth`.
282 pub fn non_empty_leaves(
283 &self,
284 root: RpoDigest,
285 max_depth: u8,
286 ) -> impl Iterator<Item = (NodeIndex, RpoDigest)> + '_ {
287 let empty_roots = EmptySubtreeRoots::empty_hashes(max_depth);
288 let mut stack = Vec::new();
289 stack.push((NodeIndex::new_unchecked(0, 0), root));
290
291 core::iter::from_fn(move || {
292 while let Some((index, node_hash)) = stack.pop() {
293 if index.depth() == max_depth {
294 return Some((index, node_hash));
295 }
296
297 if let Some(node) = self.nodes.get(&node_hash) {
298 if !empty_roots.contains(&node.left) {
299 stack.push((index.left_child(), node.left));
300 }
301 if !empty_roots.contains(&node.right) {
302 stack.push((index.right_child(), node.right));
303 }
304 }
305 }
306
307 None
308 })
309 }
310
311 // STATE MUTATORS
312 // --------------------------------------------------------------------------------------------
313
314 /// Adds all the nodes of a Merkle path represented by `path`, opening to `node`. Returns the
315 /// new root.
316 ///
317 /// This will compute the sibling elements determined by the Merkle `path` and `node`, and
318 /// include all the nodes into the store.
319 pub fn add_merkle_path(
320 &mut self,
321 index: u64,
322 node: RpoDigest,
323 path: MerklePath,
324 ) -> Result<RpoDigest, MerkleError> {
325 let root = path.inner_nodes(index, node)?.fold(RpoDigest::default(), |_, node| {
326 let value: RpoDigest = node.value;
327 let left: RpoDigest = node.left;
328 let right: RpoDigest = node.right;
329
330 debug_assert_eq!(Rpo256::merge(&[left, right]), value);
331 self.nodes.insert(value, StoreNode { left, right });
332
333 node.value
334 });
335 Ok(root)
336 }
337
338 /// Adds all the nodes of multiple Merkle paths into the store.
339 ///
340 /// This will compute the sibling elements for each Merkle `path` and include all the nodes
341 /// into the store.
342 ///
343 /// For further reference, check [MerkleStore::add_merkle_path].
344 pub fn add_merkle_paths<I>(&mut self, paths: I) -> Result<(), MerkleError>
345 where
346 I: IntoIterator<Item = (u64, RpoDigest, MerklePath)>,
347 {
348 for (index_value, node, path) in paths.into_iter() {
349 self.add_merkle_path(index_value, node, path)?;
350 }
351 Ok(())
352 }
353
354 /// Appends the provided [MerklePathSet] into the store.
355 ///
356 /// For further reference, check [MerkleStore::add_merkle_path].
357 pub fn add_merkle_path_set(
358 &mut self,
359 path_set: &MerklePathSet,
360 ) -> Result<RpoDigest, MerkleError> {
361 let root = path_set.root();
362 for (index, path) in path_set.to_paths() {
363 self.add_merkle_path(index, path.value, path.path)?;
364 }
365 Ok(root)
366 }
367
368 /// Sets a node to `value`.
369 ///
370 /// # Errors
371 /// This method can return the following errors:
372 /// - `RootNotInStore` if the `root` is not present in the store.
373 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in
374 /// the store.
375 pub fn set_node(
376 &mut self,
377 mut root: RpoDigest,
378 index: NodeIndex,
379 value: RpoDigest,
380 ) -> Result<RootPath, MerkleError> {
381 let node = value;
382 let ValuePath { value, path } = self.get_path(root, index)?;
383
384 // performs the update only if the node value differs from the opening
385 if node != value {
386 root = self.add_merkle_path(index.value(), node, path.clone())?;
387 }
388
389 Ok(RootPath { root, path })
390 }
391
392 /// Merges two elements and adds the resulting node into the store.
393 ///
394 /// Merges arbitrary values. They may be leafs, nodes, or a mixture of both.
395 pub fn merge_roots(
396 &mut self,
397 left_root: RpoDigest,
398 right_root: RpoDigest,
399 ) -> Result<RpoDigest, MerkleError> {
400 let parent = Rpo256::merge(&[left_root, right_root]);
401 self.nodes.insert(
402 parent,
403 StoreNode {
404 left: left_root,
405 right: right_root,
406 },
407 );
408
409 Ok(parent)
410 }
411
412 // DESTRUCTURING
413 // --------------------------------------------------------------------------------------------
414
415 /// Returns the inner storage of this MerkleStore while consuming `self`.
416 pub fn into_inner(self) -> T {
417 self.nodes
418 }
419
420 // HELPER METHODS
421 // --------------------------------------------------------------------------------------------
422
423 /// Recursively clones a tree with the specified root from the specified source into self.
424 ///
425 /// If the source store does not contain a tree with the specified root, this is a noop.
426 fn clone_tree_from(&mut self, root: RpoDigest, source: &Self) {
427 // process the node only if it is in the source
428 if let Some(node) = source.nodes.get(&root) {
429 // if the node has already been inserted, no need to process it further as all of its
430 // descendants should be already cloned from the source store
431 if self.nodes.insert(root, *node).is_none() {
432 self.clone_tree_from(node.left, source);
433 self.clone_tree_from(node.right, source);
434 }
435 }
436 }
437}
438
439// CONVERSIONS
440// ================================================================================================
441
442impl<T: KvMap<RpoDigest, StoreNode>> From<&MerkleTree> for MerkleStore<T> {
443 fn from(value: &MerkleTree) -> Self {
444 let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
445 Self { nodes }
446 }
447}
448
449impl<T: KvMap<RpoDigest, StoreNode>> From<&SimpleSmt> for MerkleStore<T> {
450 fn from(value: &SimpleSmt) -> Self {
451 let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
452 Self { nodes }
453 }
454}
455
456impl<T: KvMap<RpoDigest, StoreNode>> From<&Mmr> for MerkleStore<T> {
457 fn from(value: &Mmr) -> Self {
458 let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
459 Self { nodes }
460 }
461}
462
463impl<T: KvMap<RpoDigest, StoreNode>> From<&TieredSmt> for MerkleStore<T> {
464 fn from(value: &TieredSmt) -> Self {
465 let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
466 Self { nodes }
467 }
468}
469
470impl<T: KvMap<RpoDigest, StoreNode>> From<T> for MerkleStore<T> {
471 fn from(values: T) -> Self {
472 let nodes = values.into_iter().chain(empty_hashes()).collect();
473 Self { nodes }
474 }
475}
476
477impl<T: KvMap<RpoDigest, StoreNode>> FromIterator<InnerNodeInfo> for MerkleStore<T> {
478 fn from_iter<I: IntoIterator<Item = InnerNodeInfo>>(iter: I) -> Self {
479 let nodes = combine_nodes_with_empty_hashes(iter).collect();
480 Self { nodes }
481 }
482}
483
484impl<T: KvMap<RpoDigest, StoreNode>> FromIterator<(RpoDigest, StoreNode)> for MerkleStore<T> {
485 fn from_iter<I: IntoIterator<Item = (RpoDigest, StoreNode)>>(iter: I) -> Self {
486 let nodes = iter.into_iter().chain(empty_hashes()).collect();
487 Self { nodes }
488 }
489}
490
491// ITERATORS
492// ================================================================================================
493impl<T: KvMap<RpoDigest, StoreNode>> Extend<InnerNodeInfo> for MerkleStore<T> {
494 fn extend<I: IntoIterator<Item = InnerNodeInfo>>(&mut self, iter: I) {
495 self.nodes.extend(iter.into_iter().map(|info| {
496 (
497 info.value,
498 StoreNode {
499 left: info.left,
500 right: info.right,
501 },
502 )
503 }));
504 }
505}
506
507// DiffT & ApplyDiffT TRAIT IMPLEMENTATION
508// ================================================================================================
509impl<T: KvMap<RpoDigest, StoreNode>> TryApplyDiff<RpoDigest, StoreNode> for MerkleStore<T> {
510 type Error = MerkleError;
511 type DiffType = MerkleStoreDelta;
512
513 fn try_apply(&mut self, diff: Self::DiffType) -> Result<(), MerkleError> {
514 for (root, delta) in diff.0 {
515 let mut root = root;
516 for cleared_slot in delta.cleared_slots() {
517 root = self
518 .set_node(
519 root,
520 NodeIndex::new(delta.depth(), *cleared_slot)?,
521 EMPTY_WORD.into(),
522 )?
523 .root;
524 }
525 for (updated_slot, updated_value) in delta.updated_slots() {
526 root = self
527 .set_node(
528 root,
529 NodeIndex::new(delta.depth(), *updated_slot)?,
530 (*updated_value).into(),
531 )?
532 .root;
533 }
534 }
535
536 Ok(())
537 }
538}
539
540// SERIALIZATION
541// ================================================================================================
542
543impl Serializable for StoreNode {
544 fn write_into<W: ByteWriter>(&self, target: &mut W) {
545 self.left.write_into(target);
546 self.right.write_into(target);
547 }
548}
549
550impl Deserializable for StoreNode {
551 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
552 let left = RpoDigest::read_from(source)?;
553 let right = RpoDigest::read_from(source)?;
554 Ok(StoreNode { left, right })
555 }
556}
557
558impl<T: KvMap<RpoDigest, StoreNode>> Serializable for MerkleStore<T> {
559 fn write_into<W: ByteWriter>(&self, target: &mut W) {
560 target.write_u64(self.nodes.len() as u64);
561
562 for (k, v) in self.nodes.iter() {
563 k.write_into(target);
564 v.write_into(target);
565 }
566 }
567}
568
569impl<T: KvMap<RpoDigest, StoreNode>> Deserializable for MerkleStore<T> {
570 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
571 let len = source.read_u64()?;
572 let mut nodes: Vec<(RpoDigest, StoreNode)> = Vec::with_capacity(len as usize);
573
574 for _ in 0..len {
575 let key = RpoDigest::read_from(source)?;
576 let value = StoreNode::read_from(source)?;
577 nodes.push((key, value));
578 }
579
580 Ok(nodes.into_iter().collect())
581 }
582}
583
584// HELPER FUNCTIONS
585// ================================================================================================
586
587/// Creates empty hashes for all the subtrees of a tree with a max depth of 255.
588fn empty_hashes() -> impl IntoIterator<Item = (RpoDigest, StoreNode)> {
589 let subtrees = EmptySubtreeRoots::empty_hashes(255);
590 subtrees.iter().rev().copied().zip(subtrees.iter().rev().skip(1).copied()).map(
591 |(child, parent)| {
592 (
593 parent,
594 StoreNode {
595 left: child,
596 right: child,
597 },
598 )
599 },
600 )
601}
602
603/// Consumes an iterator of [InnerNodeInfo] and returns an iterator of `(value, node)` tuples
604/// which includes the nodes associate with roots of empty subtrees up to a depth of 255.
605fn combine_nodes_with_empty_hashes(
606 nodes: impl IntoIterator<Item = InnerNodeInfo>,
607) -> impl Iterator<Item = (RpoDigest, StoreNode)> {
608 nodes
609 .into_iter()
610 .map(|info| {
611 (
612 info.value,
613 StoreNode {
614 left: info.left,
615 right: info.right,
616 },
617 )
618 })
619 .chain(empty_hashes())
620}