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