miden-crypto/src/merkle/store.rs
src/merkle/store.rsRaw
1//! An in-memory data store for Merkle-lized data
2//!
3//! This is a in memory data store for Merkle trees, this store allows all the nodes of a tree
4//! (leaves or internal) to live as long as necessary and without duplication, this allows the
5//! implementation of efficient persistent data structures
6use super::{
7 BTreeMap, BTreeSet, EmptySubtreeRoots, MerkleError, MerklePath, MerklePathSet, MerkleTree,
8 NodeIndex, Rpo256, RpoDigest, SimpleSmt, Vec, Word,
9};
10
11#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
12pub struct Node {
13 left: RpoDigest,
14 right: RpoDigest,
15}
16
17#[derive(Debug, Clone, Eq, PartialEq)]
18pub struct MerkleStore {
19 nodes: BTreeMap<RpoDigest, Node>,
20}
21
22impl Default for MerkleStore {
23 fn default() -> Self {
24 Self::new()
25 }
26}
27
28impl MerkleStore {
29 // CONSTRUCTORS
30 // --------------------------------------------------------------------------------------------
31
32 /// Creates an empty `MerkleStore` instance.
33 pub fn new() -> MerkleStore {
34 // pre-populate the store with the empty hashes
35 let subtrees = EmptySubtreeRoots::empty_hashes(64);
36 let nodes = subtrees
37 .iter()
38 .copied()
39 .zip(subtrees.iter().skip(1).copied())
40 .map(|(child, parent)| {
41 (
42 parent,
43 Node {
44 left: child,
45 right: child,
46 },
47 )
48 })
49 .collect();
50
51 MerkleStore { nodes }
52 }
53
54 /// Appends the provided merkle tree represented by its `leaves` to the set.
55 pub fn with_merkle_tree<I>(mut self, leaves: I) -> Result<Self, MerkleError>
56 where
57 I: IntoIterator<Item = Word>,
58 {
59 self.add_merkle_tree(leaves)?;
60 Ok(self)
61 }
62
63 /// Appends the provided sparse merkle tree represented by its `entries` to the set.
64 pub fn with_sparse_merkle_tree<R, I>(mut self, entries: R) -> Result<Self, MerkleError>
65 where
66 R: IntoIterator<IntoIter = I>,
67 I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
68 {
69 self.add_sparse_merkle_tree(entries)?;
70 Ok(self)
71 }
72
73 /// Appends the provided merkle path set.
74 pub fn with_merkle_path(
75 mut self,
76 index_value: u64,
77 node: Word,
78 path: MerklePath,
79 ) -> Result<Self, MerkleError> {
80 self.add_merkle_path(index_value, node, path)?;
81 Ok(self)
82 }
83
84 // PUBLIC ACCESSORS
85 // --------------------------------------------------------------------------------------------
86
87 /// Returns the node at `index` rooted on the tree `root`.
88 ///
89 /// # Errors
90 ///
91 /// This will return `NodeNotInStorage` if the element is not present in the store.
92 pub fn get_node(&self, root: Word, index: NodeIndex) -> Result<Word, MerkleError> {
93 let mut hash: RpoDigest = root.into();
94
95 // Check the root is in the storage when called with `NodeIndex::root()`
96 self.nodes
97 .get(&hash)
98 .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?;
99
100 for bit in index.bit_iterator().rev() {
101 let node = self
102 .nodes
103 .get(&hash)
104 .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?;
105 hash = if bit { node.right } else { node.left }
106 }
107
108 Ok(hash.into())
109 }
110
111 /// Returns the path for the node at `index` rooted on the tree `root`.
112 ///
113 /// The path starts at the sibling of the target leaf.
114 ///
115 /// # Errors
116 ///
117 /// This will return `NodeNotInStorage` if the element is not present in the store.
118 pub fn get_path(&self, root: Word, mut index: NodeIndex) -> Result<MerklePath, MerkleError> {
119 let mut path = Vec::with_capacity(index.depth().saturating_sub(1) as usize);
120 while index.depth() > 0 {
121 let sibling = index.sibling();
122 index.move_up();
123 let node = self.get_node(root, sibling)?;
124 path.push(node);
125 }
126 Ok(MerklePath::new(path))
127 }
128
129 // STATE MUTATORS
130 // --------------------------------------------------------------------------------------------
131
132 /// Adds all the nodes of a Merkle tree represented by `leaves`.
133 ///
134 /// This will instantiate a Merkle tree using `leaves` and include all the nodes into the
135 /// storage.
136 ///
137 /// # Errors
138 ///
139 /// This method may return the following errors:
140 /// - `DepthTooSmall` if leaves is empty or contains only 1 element
141 /// - `NumLeavesNotPowerOfTwo` if the number of leaves is not a power-of-two
142 pub fn add_merkle_tree<I>(&mut self, leaves: I) -> Result<Word, MerkleError>
143 where
144 I: IntoIterator<Item = Word>,
145 {
146 let leaves: Vec<_> = leaves.into_iter().collect();
147 if leaves.len() < 2 {
148 return Err(MerkleError::DepthTooSmall(leaves.len() as u8));
149 }
150
151 let layers = leaves.len().ilog2();
152 let tree = MerkleTree::new(leaves)?;
153
154 let mut depth = 0;
155 let mut parent_offset = 1;
156 let mut child_offset = 2;
157 while depth < layers {
158 let layer_size = 1usize << depth;
159 for _ in 0..layer_size {
160 // merkle tree is using level form representation, so left and right siblings are
161 // next to each other
162 let left = tree.nodes[child_offset];
163 let right = tree.nodes[child_offset + 1];
164 self.nodes.insert(
165 tree.nodes[parent_offset].into(),
166 Node {
167 left: left.into(),
168 right: right.into(),
169 },
170 );
171 parent_offset += 1;
172 child_offset += 2;
173 }
174 depth += 1;
175 }
176
177 Ok(tree.nodes[1])
178 }
179
180 /// Adds all the nodes of a Sparse Merkle tree represented by `entries`.
181 ///
182 /// This will instantiate a Sparse Merkle tree using `entries` and include all the nodes into
183 /// the storage.
184 ///
185 /// # Errors
186 ///
187 /// This will return `InvalidEntriesCount` if the length of `entries` is not `63`.
188 pub fn add_sparse_merkle_tree<R, I>(&mut self, entries: R) -> Result<Word, MerkleError>
189 where
190 R: IntoIterator<IntoIter = I>,
191 I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
192 {
193 let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)?.with_leaves(entries)?;
194 for branch in smt.store.branches.values() {
195 let parent = Rpo256::merge(&[branch.left, branch.right]);
196 self.nodes.insert(
197 parent,
198 Node {
199 left: branch.left,
200 right: branch.right,
201 },
202 );
203 }
204
205 Ok(smt.root())
206 }
207
208 /// Adds all the nodes of a Merkle path represented by `path`.
209 ///
210 /// This will compute the sibling elements determined by the Merkle `path` and `node`, and
211 /// include all the nodes into the storage.
212 pub fn add_merkle_path(
213 &mut self,
214 index_value: u64,
215 node: Word,
216 path: MerklePath,
217 ) -> Result<Word, MerkleError> {
218 let mut node = node;
219 let mut index = NodeIndex::new(self.nodes.len() as u8, index_value);
220
221 for sibling in path {
222 let (left, right) = match index.is_value_odd() {
223 true => (sibling, node),
224 false => (node, sibling),
225 };
226 let parent = Rpo256::merge(&[left.into(), right.into()]);
227 self.nodes.insert(
228 parent,
229 Node {
230 left: left.into(),
231 right: right.into(),
232 },
233 );
234
235 index.move_up();
236 node = parent.into();
237 }
238
239 Ok(node)
240 }
241
242 /// Adds all the nodes of multiple Merkle paths into the store.
243 ///
244 /// This will compute the sibling elements for each Merkle `path` and include all the nodes
245 /// into the storage.
246 ///
247 /// # Errors
248 ///
249 /// Every path must resolve to the same root, otherwise this will return an `ConflictingRoots`
250 /// error.
251 pub fn add_merkle_paths<I>(&mut self, paths: I) -> Result<Word, MerkleError>
252 where
253 I: IntoIterator<Item = (u64, Word, MerklePath)>,
254 {
255 let paths: Vec<(u64, Word, MerklePath)> = paths.into_iter().collect();
256
257 let roots: BTreeSet<RpoDigest> = paths
258 .iter()
259 .map(|(index, node, path)| path.compute_root(*index, *node).into())
260 .collect();
261
262 if roots.len() != 1 {
263 return Err(MerkleError::ConflictingRoots(
264 roots.iter().map(|v| Word::from(*v)).collect(),
265 ));
266 }
267
268 for (index_value, node, path) in paths {
269 self.add_merkle_path(index_value, node, path)?;
270 }
271
272 // Returns the parent of the last paths (assumes all paths have the same parent) or empty
273 // The length of unique_roots is checked above, so this wont panic
274 Ok(roots.iter().next().unwrap().into())
275 }
276
277 /// Appends the provided [MerklePathSet] into the store.
278 pub fn add_merkle_path_set(&mut self, path_set: &MerklePathSet) -> Result<Word, MerkleError> {
279 let root = path_set.root();
280 path_set.indexes().try_fold(root, |_, index| {
281 let node = path_set.get_node(index)?;
282 let path = path_set.get_path(index)?;
283 self.add_merkle_path(index.value(), node, path)
284 })
285 }
286
287 pub fn set_node(
288 &mut self,
289 root: Word,
290 index: NodeIndex,
291 value: Word,
292 ) -> Result<Word, MerkleError> {
293 let current_node = self.get_node(root, index)?;
294 let path = self.get_path(root, index)?;
295 if current_node != value {
296 self.add_merkle_path(index.value(), value, path)
297 } else {
298 Ok(root)
299 }
300 }
301
302 pub fn merge_roots(&mut self, root1: Word, root2: Word) -> Result<Word, MerkleError> {
303 let root1: RpoDigest = root1.into();
304 let root2: RpoDigest = root2.into();
305
306 if !self.nodes.contains_key(&root1) {
307 Err(MerkleError::NodeNotInStorage(
308 root1.into(),
309 NodeIndex::new(0, 0),
310 ))
311 } else if !self.nodes.contains_key(&root1) {
312 Err(MerkleError::NodeNotInStorage(
313 root2.into(),
314 NodeIndex::new(0, 0),
315 ))
316 } else {
317 let parent: Word = Rpo256::merge(&[root1, root2]).into();
318 self.nodes.insert(
319 parent.into(),
320 Node {
321 left: root1,
322 right: root2,
323 },
324 );
325
326 Ok(parent)
327 }
328 }
329}
330
331#[cfg(test)]
332mod test {
333 use super::*;
334 use crate::{
335 hash::rpo::Rpo256,
336 merkle::{int_to_node, MerklePathSet},
337 Felt, Word,
338 };
339
340 const KEYS4: [u64; 4] = [0, 1, 2, 3];
341 const LEAVES4: [Word; 4] = [
342 int_to_node(1),
343 int_to_node(2),
344 int_to_node(3),
345 int_to_node(4),
346 ];
347
348 #[test]
349 fn test_add_merkle_tree() -> Result<(), MerkleError> {
350 let mut store = MerkleStore::default();
351
352 let mtree = MerkleTree::new(LEAVES4.to_vec())?;
353 store.add_merkle_tree(LEAVES4.to_vec())?;
354
355 assert!(
356 store
357 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0))
358 .is_ok(),
359 "node 0 must be in the tree"
360 );
361 assert!(
362 store
363 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1))
364 .is_ok(),
365 "node 1 must be in the tree"
366 );
367 assert!(
368 store
369 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2))
370 .is_ok(),
371 "node 2 must be in the tree"
372 );
373 assert!(
374 store
375 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3))
376 .is_ok(),
377 "node 3 must be in the tree"
378 );
379
380 store
381 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0))
382 .expect("node 0 must be in tree");
383 store
384 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1))
385 .expect("node 1 must be in tree");
386 store
387 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2))
388 .expect("node 2 must be in tree");
389 store
390 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3))
391 .expect("node 3 must be in tree");
392
393 Ok(())
394 }
395
396 #[test]
397 fn test_get_invalid_node() {
398 let mut store = MerkleStore::default();
399 let mtree = MerkleTree::new(LEAVES4.to_vec()).expect("creating a merkle tree must work");
400 store
401 .add_merkle_tree(LEAVES4.to_vec())
402 .expect("adding a merkle tree to the store must work");
403 let _ = store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3));
404 }
405
406 #[test]
407 fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
408 let mut store = MerkleStore::default();
409 let keys2: [u64; 2] = [0, 1];
410 let leaves2: [Word; 2] = [int_to_node(1), int_to_node(2)];
411 store.add_sparse_merkle_tree(keys2.into_iter().zip(leaves2.into_iter()))?;
412 let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
413 .unwrap()
414 .with_leaves(keys2.into_iter().zip(leaves2.into_iter()))
415 .unwrap();
416
417 let idx = NodeIndex::new(1, 0);
418 assert_eq!(
419 store.get_node(smt.root(), idx).unwrap(),
420 smt.get_node(&idx).unwrap()
421 );
422
423 Ok(())
424 }
425
426 #[test]
427 fn test_add_sparse_merkle_tree() -> Result<(), MerkleError> {
428 let mut store = MerkleStore::default();
429 store.add_sparse_merkle_tree(KEYS4.into_iter().zip(LEAVES4.into_iter()))?;
430
431 let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
432 .unwrap()
433 .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter()))
434 .unwrap();
435
436 let idx = NodeIndex::new(1, 0);
437 assert_eq!(
438 store.get_node(smt.root(), idx).unwrap(),
439 smt.get_node(&idx).unwrap()
440 );
441 let idx = NodeIndex::new(1, 1);
442 assert_eq!(
443 store.get_node(smt.root(), idx).unwrap(),
444 smt.get_node(&idx).unwrap()
445 );
446
447 Ok(())
448 }
449
450 #[test]
451 fn test_add_merkle_paths() -> Result<(), MerkleError> {
452 let mut store = MerkleStore::default();
453 let mtree = MerkleTree::new(LEAVES4.to_vec())?;
454
455 let i0 = 0;
456 let p0 = mtree.get_path(NodeIndex::new(2, i0)).unwrap();
457
458 let i1 = 1;
459 let p1 = mtree.get_path(NodeIndex::new(2, i1)).unwrap();
460
461 let i2 = 2;
462 let p2 = mtree.get_path(NodeIndex::new(2, i2)).unwrap();
463
464 let i3 = 3;
465 let p3 = mtree.get_path(NodeIndex::new(2, i3)).unwrap();
466
467 let paths = [
468 (i0, LEAVES4[i0 as usize], p0),
469 (i1, LEAVES4[i1 as usize], p1),
470 (i2, LEAVES4[i2 as usize], p2),
471 (i3, LEAVES4[i3 as usize], p3),
472 ];
473
474 store
475 .add_merkle_paths(paths.clone())
476 .expect("the valid paths must work");
477
478 let set = MerklePathSet::new(3).with_paths(paths).unwrap();
479
480 assert_eq!(
481 set.get_node(NodeIndex::new(3, 0)).unwrap(),
482 store.get_node(set.root(), NodeIndex::new(2, 0b00)).unwrap(),
483 );
484 assert_eq!(
485 set.get_node(NodeIndex::new(3, 1)).unwrap(),
486 store.get_node(set.root(), NodeIndex::new(2, 0b01)).unwrap(),
487 );
488 assert_eq!(
489 set.get_node(NodeIndex::new(3, 2)).unwrap(),
490 store.get_node(set.root(), NodeIndex::new(2, 0b10)).unwrap(),
491 );
492 assert_eq!(
493 set.get_node(NodeIndex::new(3, 3)).unwrap(),
494 store.get_node(set.root(), NodeIndex::new(2, 0b11)).unwrap(),
495 );
496
497 Ok(())
498 }
499
500 #[test]
501 fn wont_open_to_different_depth_root() {
502 let empty = EmptySubtreeRoots::empty_hashes(64);
503 let a = [Felt::new(1); 4];
504 let b = [Felt::new(2); 4];
505
506 // compute the root for a different depth
507 let mut root = Rpo256::merge(&[a.into(), b.into()]);
508 for depth in (1..=63).rev() {
509 root = Rpo256::merge(&[root, empty[depth]]);
510 }
511 let root = Word::from(root);
512
513 let store = MerkleStore::default().with_merkle_tree([a, b]).unwrap();
514 let index = NodeIndex::root();
515 let err = store.get_node(root, index).err().unwrap();
516 assert_eq!(err, MerkleError::NodeNotInStorage(root, index));
517 }
518
519 #[test]
520 fn store_path_opens_from_leaf() {
521 let a = [Felt::new(1); 4];
522 let b = [Felt::new(2); 4];
523 let c = [Felt::new(3); 4];
524 let d = [Felt::new(4); 4];
525 let e = [Felt::new(5); 4];
526 let f = [Felt::new(6); 4];
527 let g = [Felt::new(7); 4];
528 let h = [Felt::new(8); 4];
529
530 let i = Rpo256::merge(&[a.into(), b.into()]);
531 let j = Rpo256::merge(&[c.into(), d.into()]);
532 let k = Rpo256::merge(&[e.into(), f.into()]);
533 let l = Rpo256::merge(&[g.into(), h.into()]);
534
535 let m = Rpo256::merge(&[i.into(), j.into()]);
536 let n = Rpo256::merge(&[k.into(), l.into()]);
537
538 let root = Rpo256::merge(&[m.into(), n.into()]);
539
540 let store = MerkleStore::default()
541 .with_merkle_tree([a, b, c, d, e, f, g, h])
542 .unwrap();
543 let path = store.get_path(root.into(), NodeIndex::new(3, 1)).unwrap();
544
545 let expected = MerklePath::new([a.into(), j.into(), n.into()].to_vec());
546 assert_eq!(path, expected);
547 }
548}