miden-crypto/src/merkle/store/tests.rs
src/merkle/store/tests.rsRaw
1use seq_macro::seq;
2#[cfg(feature = "std")]
3use {
4 super::{Deserializable, Serializable},
5 alloc::boxed::Box,
6 std::error::Error,
7};
8
9use super::{
10 DefaultMerkleStore as MerkleStore, EmptySubtreeRoots, MerkleError, MerklePath, NodeIndex,
11 PartialMerkleTree, RecordingMerkleStore, Rpo256, RpoDigest,
12};
13use crate::{
14 merkle::{
15 digests_to_words, int_to_leaf, int_to_node, LeafIndex, MerkleTree, SimpleSmt, SMT_MAX_DEPTH,
16 },
17 Felt, Word, ONE, WORD_SIZE, ZERO,
18};
19
20// TEST DATA
21// ================================================================================================
22
23const KEYS4: [u64; 4] = [0, 1, 2, 3];
24const VALUES4: [RpoDigest; 4] = [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)];
25
26const KEYS8: [u64; 8] = [0, 1, 2, 3, 4, 5, 6, 7];
27const VALUES8: [RpoDigest; 8] = [
28 int_to_node(1),
29 int_to_node(2),
30 int_to_node(3),
31 int_to_node(4),
32 int_to_node(5),
33 int_to_node(6),
34 int_to_node(7),
35 int_to_node(8),
36];
37
38// TESTS
39// ================================================================================================
40
41#[test]
42fn test_root_not_in_store() -> Result<(), MerkleError> {
43 let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
44 let store = MerkleStore::from(&mtree);
45 assert_eq!(
46 store.get_node(VALUES4[0], NodeIndex::make(mtree.depth(), 0)),
47 Err(MerkleError::RootNotInStore(VALUES4[0])),
48 "Leaf 0 is not a root"
49 );
50 assert_eq!(
51 store.get_path(VALUES4[0], NodeIndex::make(mtree.depth(), 0)),
52 Err(MerkleError::RootNotInStore(VALUES4[0])),
53 "Leaf 0 is not a root"
54 );
55
56 Ok(())
57}
58
59#[test]
60fn test_merkle_tree() -> Result<(), MerkleError> {
61 let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
62 let store = MerkleStore::from(&mtree);
63
64 // STORE LEAVES ARE CORRECT -------------------------------------------------------------------
65 // checks the leaves in the store corresponds to the expected values
66 assert_eq!(
67 store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 0)),
68 Ok(VALUES4[0]),
69 "node 0 must be in the tree"
70 );
71 assert_eq!(
72 store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 1)),
73 Ok(VALUES4[1]),
74 "node 1 must be in the tree"
75 );
76 assert_eq!(
77 store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 2)),
78 Ok(VALUES4[2]),
79 "node 2 must be in the tree"
80 );
81 assert_eq!(
82 store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)),
83 Ok(VALUES4[3]),
84 "node 3 must be in the tree"
85 );
86
87 // STORE LEAVES MATCH TREE --------------------------------------------------------------------
88 // sanity check the values returned by the store and the tree
89 assert_eq!(
90 mtree.get_node(NodeIndex::make(mtree.depth(), 0)),
91 store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 0)),
92 "node 0 must be the same for both MerkleTree and MerkleStore"
93 );
94 assert_eq!(
95 mtree.get_node(NodeIndex::make(mtree.depth(), 1)),
96 store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 1)),
97 "node 1 must be the same for both MerkleTree and MerkleStore"
98 );
99 assert_eq!(
100 mtree.get_node(NodeIndex::make(mtree.depth(), 2)),
101 store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 2)),
102 "node 2 must be the same for both MerkleTree and MerkleStore"
103 );
104 assert_eq!(
105 mtree.get_node(NodeIndex::make(mtree.depth(), 3)),
106 store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)),
107 "node 3 must be the same for both MerkleTree and MerkleStore"
108 );
109
110 // STORE MERKLE PATH MATCHES ==============================================================
111 // assert the merkle path returned by the store is the same as the one in the tree
112 let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 0)).unwrap();
113 assert_eq!(
114 VALUES4[0], result.value,
115 "Value for merkle path at index 0 must match leaf value"
116 );
117 assert_eq!(
118 mtree.get_path(NodeIndex::make(mtree.depth(), 0)),
119 Ok(result.path),
120 "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
121 );
122
123 let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 1)).unwrap();
124 assert_eq!(
125 VALUES4[1], result.value,
126 "Value for merkle path at index 0 must match leaf value"
127 );
128 assert_eq!(
129 mtree.get_path(NodeIndex::make(mtree.depth(), 1)),
130 Ok(result.path),
131 "merkle path for index 1 must be the same for the MerkleTree and MerkleStore"
132 );
133
134 let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 2)).unwrap();
135 assert_eq!(
136 VALUES4[2], result.value,
137 "Value for merkle path at index 0 must match leaf value"
138 );
139 assert_eq!(
140 mtree.get_path(NodeIndex::make(mtree.depth(), 2)),
141 Ok(result.path),
142 "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
143 );
144
145 let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 3)).unwrap();
146 assert_eq!(
147 VALUES4[3], result.value,
148 "Value for merkle path at index 0 must match leaf value"
149 );
150 assert_eq!(
151 mtree.get_path(NodeIndex::make(mtree.depth(), 3)),
152 Ok(result.path),
153 "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
154 );
155
156 Ok(())
157}
158
159#[test]
160fn test_empty_roots() {
161 let store = MerkleStore::default();
162 let mut root = RpoDigest::default();
163
164 for depth in 0..255 {
165 root = Rpo256::merge(&[root; 2]);
166 assert!(
167 store.get_node(root, NodeIndex::make(0, 0)).is_ok(),
168 "The root of the empty tree of depth {depth} must be registered"
169 );
170 }
171}
172
173#[test]
174fn test_leaf_paths_for_empty_trees() -> Result<(), MerkleError> {
175 let store = MerkleStore::default();
176
177 // Starts at 1 because leafs are not included in the store.
178 // Ends at 64 because it is not possible to represent an index of a depth greater than 64,
179 // because a u64 is used to index the leaf.
180 seq!(DEPTH in 1_u8..64_u8 {
181 let smt = SimpleSmt::<DEPTH>::new()?;
182
183 let index = NodeIndex::make(DEPTH, 0);
184 let store_path = store.get_path(smt.root(), index)?;
185 let smt_path = smt.open(&LeafIndex::<DEPTH>::new(0)?).path;
186 assert_eq!(
187 store_path.value,
188 RpoDigest::default(),
189 "the leaf of an empty tree is always ZERO"
190 );
191 assert_eq!(
192 store_path.path, smt_path,
193 "the returned merkle path does not match the computed values"
194 );
195 assert_eq!(
196 store_path.path.compute_root(DEPTH.into(), RpoDigest::default()).unwrap(),
197 smt.root(),
198 "computed root from the path must match the empty tree root"
199 );
200
201 });
202
203 Ok(())
204}
205
206#[test]
207fn test_get_invalid_node() {
208 let mtree =
209 MerkleTree::new(digests_to_words(&VALUES4)).expect("creating a merkle tree must work");
210 let store = MerkleStore::from(&mtree);
211 let _ = store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3));
212}
213
214#[test]
215fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
216 let keys2: [u64; 2] = [0, 1];
217 let leaves2: [Word; 2] = [int_to_leaf(1), int_to_leaf(2)];
218 let smt = SimpleSmt::<1>::with_leaves(keys2.into_iter().zip(leaves2)).unwrap();
219 let store = MerkleStore::from(&smt);
220
221 let idx = NodeIndex::make(1, 0);
222 assert_eq!(smt.get_node(idx).unwrap(), leaves2[0].into());
223 assert_eq!(store.get_node(smt.root(), idx).unwrap(), smt.get_node(idx).unwrap());
224
225 let idx = NodeIndex::make(1, 1);
226 assert_eq!(smt.get_node(idx).unwrap(), leaves2[1].into());
227 assert_eq!(store.get_node(smt.root(), idx).unwrap(), smt.get_node(idx).unwrap());
228
229 Ok(())
230}
231
232#[test]
233fn test_sparse_merkle_tree() -> Result<(), MerkleError> {
234 let smt =
235 SimpleSmt::<SMT_MAX_DEPTH>::with_leaves(KEYS4.into_iter().zip(digests_to_words(&VALUES4)))
236 .unwrap();
237
238 let store = MerkleStore::from(&smt);
239
240 // STORE LEAVES ARE CORRECT ==============================================================
241 // checks the leaves in the store corresponds to the expected values
242 assert_eq!(
243 store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 0)),
244 Ok(VALUES4[0]),
245 "node 0 must be in the tree"
246 );
247 assert_eq!(
248 store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 1)),
249 Ok(VALUES4[1]),
250 "node 1 must be in the tree"
251 );
252 assert_eq!(
253 store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 2)),
254 Ok(VALUES4[2]),
255 "node 2 must be in the tree"
256 );
257 assert_eq!(
258 store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 3)),
259 Ok(VALUES4[3]),
260 "node 3 must be in the tree"
261 );
262 assert_eq!(
263 store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 4)),
264 Ok(RpoDigest::default()),
265 "unmodified node 4 must be ZERO"
266 );
267
268 // STORE LEAVES MATCH TREE ===============================================================
269 // sanity check the values returned by the store and the tree
270 assert_eq!(
271 smt.get_node(NodeIndex::make(SMT_MAX_DEPTH, 0)),
272 store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 0)),
273 "node 0 must be the same for both SparseMerkleTree and MerkleStore"
274 );
275 assert_eq!(
276 smt.get_node(NodeIndex::make(SMT_MAX_DEPTH, 1)),
277 store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 1)),
278 "node 1 must be the same for both SparseMerkleTree and MerkleStore"
279 );
280 assert_eq!(
281 smt.get_node(NodeIndex::make(SMT_MAX_DEPTH, 2)),
282 store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 2)),
283 "node 2 must be the same for both SparseMerkleTree and MerkleStore"
284 );
285 assert_eq!(
286 smt.get_node(NodeIndex::make(SMT_MAX_DEPTH, 3)),
287 store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 3)),
288 "node 3 must be the same for both SparseMerkleTree and MerkleStore"
289 );
290 assert_eq!(
291 smt.get_node(NodeIndex::make(SMT_MAX_DEPTH, 4)),
292 store.get_node(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 4)),
293 "node 4 must be the same for both SparseMerkleTree and MerkleStore"
294 );
295
296 // STORE MERKLE PATH MATCHES ==============================================================
297 // assert the merkle path returned by the store is the same as the one in the tree
298 let result = store.get_path(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 0)).unwrap();
299 assert_eq!(
300 VALUES4[0], result.value,
301 "Value for merkle path at index 0 must match leaf value"
302 );
303 assert_eq!(
304 smt.open(&LeafIndex::<SMT_MAX_DEPTH>::new(0).unwrap()).path,
305 result.path,
306 "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
307 );
308
309 let result = store.get_path(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 1)).unwrap();
310 assert_eq!(
311 VALUES4[1], result.value,
312 "Value for merkle path at index 1 must match leaf value"
313 );
314 assert_eq!(
315 smt.open(&LeafIndex::<SMT_MAX_DEPTH>::new(1).unwrap()).path,
316 result.path,
317 "merkle path for index 1 must be the same for the MerkleTree and MerkleStore"
318 );
319
320 let result = store.get_path(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 2)).unwrap();
321 assert_eq!(
322 VALUES4[2], result.value,
323 "Value for merkle path at index 2 must match leaf value"
324 );
325 assert_eq!(
326 smt.open(&LeafIndex::<SMT_MAX_DEPTH>::new(2).unwrap()).path,
327 result.path,
328 "merkle path for index 2 must be the same for the MerkleTree and MerkleStore"
329 );
330
331 let result = store.get_path(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 3)).unwrap();
332 assert_eq!(
333 VALUES4[3], result.value,
334 "Value for merkle path at index 3 must match leaf value"
335 );
336 assert_eq!(
337 smt.open(&LeafIndex::<SMT_MAX_DEPTH>::new(3).unwrap()).path,
338 result.path,
339 "merkle path for index 3 must be the same for the MerkleTree and MerkleStore"
340 );
341
342 let result = store.get_path(smt.root(), NodeIndex::make(SMT_MAX_DEPTH, 4)).unwrap();
343 assert_eq!(
344 RpoDigest::default(),
345 result.value,
346 "Value for merkle path at index 4 must match leaf value"
347 );
348 assert_eq!(
349 smt.open(&LeafIndex::<SMT_MAX_DEPTH>::new(4).unwrap()).path,
350 result.path,
351 "merkle path for index 4 must be the same for the MerkleTree and MerkleStore"
352 );
353
354 Ok(())
355}
356
357#[test]
358fn test_add_merkle_paths() -> Result<(), MerkleError> {
359 let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
360
361 let i0 = 0;
362 let p0 = mtree.get_path(NodeIndex::make(2, i0)).unwrap();
363
364 let i1 = 1;
365 let p1 = mtree.get_path(NodeIndex::make(2, i1)).unwrap();
366
367 let i2 = 2;
368 let p2 = mtree.get_path(NodeIndex::make(2, i2)).unwrap();
369
370 let i3 = 3;
371 let p3 = mtree.get_path(NodeIndex::make(2, i3)).unwrap();
372
373 let paths = [
374 (i0, VALUES4[i0 as usize], p0),
375 (i1, VALUES4[i1 as usize], p1),
376 (i2, VALUES4[i2 as usize], p2),
377 (i3, VALUES4[i3 as usize], p3),
378 ];
379
380 let mut store = MerkleStore::default();
381 store.add_merkle_paths(paths.clone()).expect("the valid paths must work");
382
383 let pmt = PartialMerkleTree::with_paths(paths).unwrap();
384
385 // STORE LEAVES ARE CORRECT ==============================================================
386 // checks the leaves in the store corresponds to the expected values
387 assert_eq!(
388 store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 0)),
389 Ok(VALUES4[0]),
390 "node 0 must be in the pmt"
391 );
392 assert_eq!(
393 store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 1)),
394 Ok(VALUES4[1]),
395 "node 1 must be in the pmt"
396 );
397 assert_eq!(
398 store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 2)),
399 Ok(VALUES4[2]),
400 "node 2 must be in the pmt"
401 );
402 assert_eq!(
403 store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 3)),
404 Ok(VALUES4[3]),
405 "node 3 must be in the pmt"
406 );
407
408 // STORE LEAVES MATCH PMT ================================================================
409 // sanity check the values returned by the store and the pmt
410 assert_eq!(
411 pmt.get_node(NodeIndex::make(pmt.max_depth(), 0)),
412 store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 0)),
413 "node 0 must be the same for both PartialMerkleTree and MerkleStore"
414 );
415 assert_eq!(
416 pmt.get_node(NodeIndex::make(pmt.max_depth(), 1)),
417 store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 1)),
418 "node 1 must be the same for both PartialMerkleTree and MerkleStore"
419 );
420 assert_eq!(
421 pmt.get_node(NodeIndex::make(pmt.max_depth(), 2)),
422 store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 2)),
423 "node 2 must be the same for both PartialMerkleTree and MerkleStore"
424 );
425 assert_eq!(
426 pmt.get_node(NodeIndex::make(pmt.max_depth(), 3)),
427 store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 3)),
428 "node 3 must be the same for both PartialMerkleTree and MerkleStore"
429 );
430
431 // STORE MERKLE PATH MATCHES ==============================================================
432 // assert the merkle path returned by the store is the same as the one in the pmt
433 let result = store.get_path(pmt.root(), NodeIndex::make(pmt.max_depth(), 0)).unwrap();
434 assert_eq!(
435 VALUES4[0], result.value,
436 "Value for merkle path at index 0 must match leaf value"
437 );
438 assert_eq!(
439 pmt.get_path(NodeIndex::make(pmt.max_depth(), 0)),
440 Ok(result.path),
441 "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
442 );
443
444 let result = store.get_path(pmt.root(), NodeIndex::make(pmt.max_depth(), 1)).unwrap();
445 assert_eq!(
446 VALUES4[1], result.value,
447 "Value for merkle path at index 0 must match leaf value"
448 );
449 assert_eq!(
450 pmt.get_path(NodeIndex::make(pmt.max_depth(), 1)),
451 Ok(result.path),
452 "merkle path for index 1 must be the same for the MerkleTree and MerkleStore"
453 );
454
455 let result = store.get_path(pmt.root(), NodeIndex::make(pmt.max_depth(), 2)).unwrap();
456 assert_eq!(
457 VALUES4[2], result.value,
458 "Value for merkle path at index 0 must match leaf value"
459 );
460 assert_eq!(
461 pmt.get_path(NodeIndex::make(pmt.max_depth(), 2)),
462 Ok(result.path),
463 "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
464 );
465
466 let result = store.get_path(pmt.root(), NodeIndex::make(pmt.max_depth(), 3)).unwrap();
467 assert_eq!(
468 VALUES4[3], result.value,
469 "Value for merkle path at index 0 must match leaf value"
470 );
471 assert_eq!(
472 pmt.get_path(NodeIndex::make(pmt.max_depth(), 3)),
473 Ok(result.path),
474 "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
475 );
476
477 Ok(())
478}
479
480#[test]
481fn wont_open_to_different_depth_root() {
482 let empty = EmptySubtreeRoots::empty_hashes(64);
483 let a = [ONE; 4];
484 let b = [Felt::new(2); 4];
485
486 // Compute the root for a different depth. We cherry-pick this specific depth to prevent a
487 // regression to a bug in the past that allowed the user to fetch a node at a depth lower than
488 // the inserted path of a Merkle tree.
489 let mut root = Rpo256::merge(&[a.into(), b.into()]);
490 for depth in (1..=63).rev() {
491 root = Rpo256::merge(&[root, empty[depth]]);
492 }
493
494 // For this example, the depth of the Merkle tree is 1, as we have only two leaves. Here we
495 // attempt to fetch a node on the maximum depth, and it should fail because the root shouldn't
496 // exist for the set.
497 let mtree = MerkleTree::new(vec![a, b]).unwrap();
498 let store = MerkleStore::from(&mtree);
499 let index = NodeIndex::root();
500 let err = store.get_node(root, index).err().unwrap();
501 assert_eq!(err, MerkleError::RootNotInStore(root));
502}
503
504#[test]
505fn store_path_opens_from_leaf() {
506 let a = [ONE; 4];
507 let b = [Felt::new(2); 4];
508 let c = [Felt::new(3); 4];
509 let d = [Felt::new(4); 4];
510 let e = [Felt::new(5); 4];
511 let f = [Felt::new(6); 4];
512 let g = [Felt::new(7); 4];
513 let h = [Felt::new(8); 4];
514
515 let i = Rpo256::merge(&[a.into(), b.into()]);
516 let j = Rpo256::merge(&[c.into(), d.into()]);
517 let k = Rpo256::merge(&[e.into(), f.into()]);
518 let l = Rpo256::merge(&[g.into(), h.into()]);
519
520 let m = Rpo256::merge(&[i, j]);
521 let n = Rpo256::merge(&[k, l]);
522
523 let root = Rpo256::merge(&[m, n]);
524
525 let mtree = MerkleTree::new(vec![a, b, c, d, e, f, g, h]).unwrap();
526 let store = MerkleStore::from(&mtree);
527 let path = store.get_path(root, NodeIndex::make(3, 1)).unwrap().path;
528
529 let expected = MerklePath::new([a.into(), j, n].to_vec());
530 assert_eq!(path, expected);
531}
532
533#[test]
534fn test_set_node() -> Result<(), MerkleError> {
535 let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
536 let mut store = MerkleStore::from(&mtree);
537 let value = int_to_node(42);
538 let index = NodeIndex::make(mtree.depth(), 0);
539 let new_root = store.set_node(mtree.root(), index, value)?.root;
540 assert_eq!(store.get_node(new_root, index), Ok(value), "Value must have changed");
541
542 Ok(())
543}
544
545#[test]
546fn test_constructors() -> Result<(), MerkleError> {
547 let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
548 let store = MerkleStore::from(&mtree);
549
550 let depth = mtree.depth();
551 let leaves = 2u64.pow(depth.into());
552 for index in 0..leaves {
553 let index = NodeIndex::make(depth, index);
554 let value_path = store.get_path(mtree.root(), index)?;
555 assert_eq!(mtree.get_path(index)?, value_path.path);
556 }
557
558 const DEPTH: u8 = 32;
559 let smt =
560 SimpleSmt::<DEPTH>::with_leaves(KEYS4.into_iter().zip(digests_to_words(&VALUES4))).unwrap();
561 let store = MerkleStore::from(&smt);
562
563 for key in KEYS4 {
564 let index = NodeIndex::make(DEPTH, key);
565 let value_path = store.get_path(smt.root(), index)?;
566 assert_eq!(smt.open(&LeafIndex::<DEPTH>::new(key).unwrap()).path, value_path.path);
567 }
568
569 let d = 2;
570 let paths = [
571 (0, VALUES4[0], mtree.get_path(NodeIndex::make(d, 0)).unwrap()),
572 (1, VALUES4[1], mtree.get_path(NodeIndex::make(d, 1)).unwrap()),
573 (2, VALUES4[2], mtree.get_path(NodeIndex::make(d, 2)).unwrap()),
574 (3, VALUES4[3], mtree.get_path(NodeIndex::make(d, 3)).unwrap()),
575 ];
576
577 let mut store1 = MerkleStore::default();
578 store1.add_merkle_paths(paths.clone())?;
579
580 let mut store2 = MerkleStore::default();
581 store2.add_merkle_path(0, VALUES4[0], mtree.get_path(NodeIndex::make(d, 0))?)?;
582 store2.add_merkle_path(1, VALUES4[1], mtree.get_path(NodeIndex::make(d, 1))?)?;
583 store2.add_merkle_path(2, VALUES4[2], mtree.get_path(NodeIndex::make(d, 2))?)?;
584 store2.add_merkle_path(3, VALUES4[3], mtree.get_path(NodeIndex::make(d, 3))?)?;
585 let pmt = PartialMerkleTree::with_paths(paths).unwrap();
586
587 for key in [0, 1, 2, 3] {
588 let index = NodeIndex::make(d, key);
589 let value_path1 = store1.get_path(pmt.root(), index)?;
590 let value_path2 = store2.get_path(pmt.root(), index)?;
591 assert_eq!(value_path1, value_path2);
592
593 let index = NodeIndex::make(d, key);
594 assert_eq!(pmt.get_path(index)?, value_path1.path);
595 }
596
597 Ok(())
598}
599
600#[test]
601fn node_path_should_be_truncated_by_midtier_insert() {
602 let key = 0b11010010_11001100_11001100_11001100_11001100_11001100_11001100_11001100_u64;
603
604 let mut store = MerkleStore::new();
605 let root: RpoDigest = EmptySubtreeRoots::empty_hashes(64)[0];
606
607 // insert first node - works as expected
608 let depth = 64;
609 let node = RpoDigest::from([Felt::new(key); WORD_SIZE]);
610 let index = NodeIndex::new(depth, key).unwrap();
611 let root = store.set_node(root, index, node).unwrap().root;
612 let result = store.get_node(root, index).unwrap();
613 let path = store.get_path(root, index).unwrap().path;
614 assert_eq!(node, result);
615 assert_eq!(path.depth(), depth);
616 assert!(path.verify(index.value(), result, &root));
617
618 // flip the first bit of the key and insert the second node on a different depth
619 let key = key ^ (1 << 63);
620 let key = key >> 8;
621 let depth = 56;
622 let node = RpoDigest::from([Felt::new(key); WORD_SIZE]);
623 let index = NodeIndex::new(depth, key).unwrap();
624 let root = store.set_node(root, index, node).unwrap().root;
625 let result = store.get_node(root, index).unwrap();
626 let path = store.get_path(root, index).unwrap().path;
627 assert_eq!(node, result);
628 assert_eq!(path.depth(), depth);
629 assert!(path.verify(index.value(), result, &root));
630
631 // attempt to fetch a path of the second node to depth 64
632 // should fail because the previously inserted node will remove its sub-tree from the set
633 let key = key << 8;
634 let index = NodeIndex::new(64, key).unwrap();
635 assert!(store.get_node(root, index).is_err());
636}
637
638// LEAF TRAVERSAL
639// ================================================================================================
640
641#[test]
642fn get_leaf_depth_works_depth_64() {
643 let mut store = MerkleStore::new();
644 let mut root: RpoDigest = EmptySubtreeRoots::empty_hashes(64)[0];
645 let key = u64::MAX;
646
647 // this will create a rainbow tree and test all opening to depth 64
648 for d in 0..64 {
649 let k = key & (u64::MAX >> d);
650 let node = RpoDigest::from([Felt::new(k); WORD_SIZE]);
651 let index = NodeIndex::new(64, k).unwrap();
652
653 // assert the leaf doesn't exist before the insert. the returned depth should always
654 // increment with the paths count of the set, as they are intersecting one another up to
655 // the first bits of the used key.
656 assert_eq!(d, store.get_leaf_depth(root, 64, k).unwrap());
657
658 // insert and assert the correct depth
659 root = store.set_node(root, index, node).unwrap().root;
660 assert_eq!(64, store.get_leaf_depth(root, 64, k).unwrap());
661 }
662}
663
664#[test]
665fn get_leaf_depth_works_with_incremental_depth() {
666 let mut store = MerkleStore::new();
667 let mut root: RpoDigest = EmptySubtreeRoots::empty_hashes(64)[0];
668
669 // insert some path to the left of the root and assert it
670 let key = 0b01001011_10110110_00001101_01110100_00111011_10101101_00000100_01000001_u64;
671 assert_eq!(0, store.get_leaf_depth(root, 64, key).unwrap());
672 let depth = 64;
673 let index = NodeIndex::new(depth, key).unwrap();
674 let node = RpoDigest::from([Felt::new(key); WORD_SIZE]);
675 root = store.set_node(root, index, node).unwrap().root;
676 assert_eq!(depth, store.get_leaf_depth(root, 64, key).unwrap());
677
678 // flip the key to the right of the root and insert some content on depth 16
679 let key = 0b11001011_10110110_00000000_00000000_00000000_00000000_00000000_00000000_u64;
680 assert_eq!(1, store.get_leaf_depth(root, 64, key).unwrap());
681 let depth = 16;
682 let index = NodeIndex::new(depth, key >> (64 - depth)).unwrap();
683 let node = RpoDigest::from([Felt::new(key); WORD_SIZE]);
684 root = store.set_node(root, index, node).unwrap().root;
685 assert_eq!(depth, store.get_leaf_depth(root, 64, key).unwrap());
686
687 // attempt the sibling of the previous leaf
688 let key = 0b11001011_10110111_00000000_00000000_00000000_00000000_00000000_00000000_u64;
689 assert_eq!(16, store.get_leaf_depth(root, 64, key).unwrap());
690 let index = NodeIndex::new(depth, key >> (64 - depth)).unwrap();
691 let node = RpoDigest::from([Felt::new(key); WORD_SIZE]);
692 root = store.set_node(root, index, node).unwrap().root;
693 assert_eq!(depth, store.get_leaf_depth(root, 64, key).unwrap());
694
695 // move down to the next depth and assert correct behavior
696 let key = 0b11001011_10110100_00000000_00000000_00000000_00000000_00000000_00000000_u64;
697 assert_eq!(15, store.get_leaf_depth(root, 64, key).unwrap());
698 let depth = 17;
699 let index = NodeIndex::new(depth, key >> (64 - depth)).unwrap();
700 let node = RpoDigest::from([Felt::new(key); WORD_SIZE]);
701 root = store.set_node(root, index, node).unwrap().root;
702 assert_eq!(depth, store.get_leaf_depth(root, 64, key).unwrap());
703}
704
705#[test]
706fn get_leaf_depth_works_with_depth_8() {
707 let mut store = MerkleStore::new();
708 let mut root: RpoDigest = EmptySubtreeRoots::empty_hashes(8)[0];
709
710 // insert some random, 8 depth keys. `a` diverges from the first bit
711 let a = 0b01101001_u64;
712 let b = 0b10011001_u64;
713 let c = 0b10010110_u64;
714 let d = 0b11110110_u64;
715
716 for k in [a, b, c, d] {
717 let index = NodeIndex::new(8, k).unwrap();
718 let node = RpoDigest::from([Felt::new(k); WORD_SIZE]);
719 root = store.set_node(root, index, node).unwrap().root;
720 }
721
722 // assert all leaves returns the inserted depth
723 for k in [a, b, c, d] {
724 assert_eq!(8, store.get_leaf_depth(root, 8, k).unwrap());
725 }
726
727 // flip last bit of a and expect it to return the the same depth, but for an empty node
728 assert_eq!(8, store.get_leaf_depth(root, 8, 0b01101000_u64).unwrap());
729
730 // flip fourth bit of a and expect an empty node on depth 4
731 assert_eq!(4, store.get_leaf_depth(root, 8, 0b01111001_u64).unwrap());
732
733 // flip third bit of a and expect an empty node on depth 3
734 assert_eq!(3, store.get_leaf_depth(root, 8, 0b01001001_u64).unwrap());
735
736 // flip second bit of a and expect an empty node on depth 2
737 assert_eq!(2, store.get_leaf_depth(root, 8, 0b00101001_u64).unwrap());
738
739 // flip fourth bit of c and expect an empty node on depth 4
740 assert_eq!(4, store.get_leaf_depth(root, 8, 0b10000110_u64).unwrap());
741
742 // flip second bit of d and expect an empty node on depth 3 as depth 2 conflicts with b and c
743 assert_eq!(3, store.get_leaf_depth(root, 8, 0b10110110_u64).unwrap());
744
745 // duplicate the tree on `a` and assert the depth is short-circuited by such sub-tree
746 let index = NodeIndex::new(8, a).unwrap();
747 root = store.set_node(root, index, root).unwrap().root;
748 assert_eq!(Err(MerkleError::DepthTooBig(9)), store.get_leaf_depth(root, 8, a));
749}
750
751#[test]
752fn find_lone_leaf() {
753 let mut store = MerkleStore::new();
754 let empty = EmptySubtreeRoots::empty_hashes(64);
755 let mut root: RpoDigest = empty[0];
756
757 // insert a single leaf into the store at depth 64
758 let key_a = 0b01010101_10101010_00001111_01110100_00111011_10101101_00000100_01000001_u64;
759 let idx_a = NodeIndex::make(64, key_a);
760 let val_a = RpoDigest::from([ONE, ONE, ONE, ONE]);
761 root = store.set_node(root, idx_a, val_a).unwrap().root;
762
763 // for every ancestor of A, A should be a long leaf
764 for depth in 1..64 {
765 let parent_index = NodeIndex::make(depth, key_a >> (64 - depth));
766 let parent = store.get_node(root, parent_index).unwrap();
767
768 let res = store.find_lone_leaf(parent, parent_index, 64).unwrap();
769 assert_eq!(res, Some((idx_a, val_a)));
770 }
771
772 // insert another leaf into the store such that it has the same 8 bit prefix as A
773 let key_b = 0b01010101_01111010_00001111_01110100_00111011_10101101_00000100_01000001_u64;
774 let idx_b = NodeIndex::make(64, key_b);
775 let val_b = RpoDigest::from([ONE, ONE, ONE, ZERO]);
776 root = store.set_node(root, idx_b, val_b).unwrap().root;
777
778 // for any node which is common between A and B, find_lone_leaf() should return None as the
779 // node has two descendants
780 for depth in 1..9 {
781 let parent_index = NodeIndex::make(depth, key_a >> (64 - depth));
782 let parent = store.get_node(root, parent_index).unwrap();
783
784 let res = store.find_lone_leaf(parent, parent_index, 64).unwrap();
785 assert_eq!(res, None);
786 }
787
788 // for other ancestors of A and B, A and B should be lone leaves respectively
789 for depth in 9..64 {
790 let parent_index = NodeIndex::make(depth, key_a >> (64 - depth));
791 let parent = store.get_node(root, parent_index).unwrap();
792
793 let res = store.find_lone_leaf(parent, parent_index, 64).unwrap();
794 assert_eq!(res, Some((idx_a, val_a)));
795 }
796
797 for depth in 9..64 {
798 let parent_index = NodeIndex::make(depth, key_b >> (64 - depth));
799 let parent = store.get_node(root, parent_index).unwrap();
800
801 let res = store.find_lone_leaf(parent, parent_index, 64).unwrap();
802 assert_eq!(res, Some((idx_b, val_b)));
803 }
804
805 // for any other node, find_lone_leaf() should return None as they have no leaf nodes
806 let parent_index = NodeIndex::make(16, 0b01010101_11111111);
807 let parent = store.get_node(root, parent_index).unwrap();
808 let res = store.find_lone_leaf(parent, parent_index, 64).unwrap();
809 assert_eq!(res, None);
810}
811
812// SUBSET EXTRACTION
813// ================================================================================================
814
815#[test]
816fn mstore_subset() {
817 // add a Merkle tree of depth 3 to the store
818 let mtree = MerkleTree::new(digests_to_words(&VALUES8)).unwrap();
819 let mut store = MerkleStore::default();
820 let empty_store_num_nodes = store.nodes.len();
821 store.extend(mtree.inner_nodes());
822
823 // build 3 subtrees contained within the above Merkle tree; note that subtree2 is a subset
824 // of subtree1
825 let subtree1 = MerkleTree::new(digests_to_words(&VALUES8[..4])).unwrap();
826 let subtree2 = MerkleTree::new(digests_to_words(&VALUES8[2..4])).unwrap();
827 let subtree3 = MerkleTree::new(digests_to_words(&VALUES8[6..])).unwrap();
828
829 // --- extract all 3 subtrees ---------------------------------------------
830
831 let substore = store.subset([subtree1.root(), subtree2.root(), subtree3.root()].iter());
832
833 // number of nodes should increase by 4: 3 nodes form subtree1 and 1 node from subtree3
834 assert_eq!(substore.nodes.len(), empty_store_num_nodes + 4);
835
836 // make sure paths that all subtrees are in the store
837 check_mstore_subtree(&substore, &subtree1);
838 check_mstore_subtree(&substore, &subtree2);
839 check_mstore_subtree(&substore, &subtree3);
840
841 // --- extract subtrees 1 and 3 -------------------------------------------
842 // this should give the same result as above as subtree2 is nested within subtree1
843
844 let substore = store.subset([subtree1.root(), subtree3.root()].iter());
845
846 // number of nodes should increase by 4: 3 nodes form subtree1 and 1 node from subtree3
847 assert_eq!(substore.nodes.len(), empty_store_num_nodes + 4);
848
849 // make sure paths that all subtrees are in the store
850 check_mstore_subtree(&substore, &subtree1);
851 check_mstore_subtree(&substore, &subtree2);
852 check_mstore_subtree(&substore, &subtree3);
853}
854
855fn check_mstore_subtree(store: &MerkleStore, subtree: &MerkleTree) {
856 for (i, value) in subtree.leaves() {
857 let index = NodeIndex::new(subtree.depth(), i).unwrap();
858 let path1 = store.get_path(subtree.root(), index).unwrap();
859 assert_eq!(*path1.value, *value);
860
861 let path2 = subtree.get_path(index).unwrap();
862 assert_eq!(path1.path, path2);
863 }
864}
865
866// SERIALIZATION
867// ================================================================================================
868
869#[cfg(feature = "std")]
870#[test]
871fn test_serialization() -> Result<(), Box<dyn Error>> {
872 let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
873 let store = MerkleStore::from(&mtree);
874 let decoded = MerkleStore::read_from_bytes(&store.to_bytes()).expect("deserialization failed");
875 assert_eq!(store, decoded);
876 Ok(())
877}
878
879// MERKLE RECORDER
880// ================================================================================================
881#[test]
882fn test_recorder() {
883 // instantiate recorder from MerkleTree and SimpleSmt
884 let mtree = MerkleTree::new(digests_to_words(&VALUES4)).unwrap();
885
886 const TREE_DEPTH: u8 = 64;
887 let smtree = SimpleSmt::<TREE_DEPTH>::with_leaves(
888 KEYS8.into_iter().zip(VALUES8.into_iter().map(|x| x.into()).rev()),
889 )
890 .unwrap();
891
892 let mut recorder: RecordingMerkleStore =
893 mtree.inner_nodes().chain(smtree.inner_nodes()).collect();
894
895 // get nodes from both trees and make sure they are correct
896 let index_0 = NodeIndex::new(mtree.depth(), 0).unwrap();
897 let node = recorder.get_node(mtree.root(), index_0).unwrap();
898 assert_eq!(node, mtree.get_node(index_0).unwrap());
899
900 let index_1 = NodeIndex::new(TREE_DEPTH, 1).unwrap();
901 let node = recorder.get_node(smtree.root(), index_1).unwrap();
902 assert_eq!(node, smtree.get_node(index_1).unwrap());
903
904 // insert a value and assert that when we request it next time it is accurate
905 let new_value = [ZERO, ZERO, ONE, ONE].into();
906 let index_2 = NodeIndex::new(TREE_DEPTH, 2).unwrap();
907 let root = recorder.set_node(smtree.root(), index_2, new_value).unwrap().root;
908 assert_eq!(recorder.get_node(root, index_2).unwrap(), new_value);
909
910 // construct the proof
911 let rec_map = recorder.into_inner();
912 let (_, proof) = rec_map.finalize();
913 let merkle_store: MerkleStore = proof.into();
914
915 // make sure the proof contains all nodes from both trees
916 let node = merkle_store.get_node(mtree.root(), index_0).unwrap();
917 assert_eq!(node, mtree.get_node(index_0).unwrap());
918
919 let node = merkle_store.get_node(smtree.root(), index_1).unwrap();
920 assert_eq!(node, smtree.get_node(index_1).unwrap());
921
922 let node = merkle_store.get_node(smtree.root(), index_2).unwrap();
923 assert_eq!(
924 node,
925 smtree.get_leaf(&LeafIndex::<TREE_DEPTH>::try_from(index_2).unwrap()).into()
926 );
927
928 // assert that is doesnt contain nodes that were not recorded
929 let not_recorded_index = NodeIndex::new(TREE_DEPTH, 4).unwrap();
930 assert!(merkle_store.get_node(smtree.root(), not_recorded_index).is_err());
931 assert!(smtree.get_node(not_recorded_index).is_ok());
932}