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