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