arnaucube
arnaucube/miden-cryptoPublic Self-hosted repository
miden-crypto/src/merkle/store/tests.rs
src/merkle/store/tests.rsRaw
  1use super::{
  2    DefaultMerkleStore as MerkleStore, EmptySubtreeRoots, MerkleError, MerklePath, NodeIndex,
  3    PartialMerkleTree, RecordingMerkleStore, Rpo256, RpoDigest,
  4};
  5use crate::{
  6    merkle::{digests_to_words, int_to_leaf, int_to_node, MerkleTree, SimpleSmt},
  7    Felt, Word, ONE, WORD_SIZE, ZERO,
  8};
  9
 10#[cfg(feature = "std")]
 11use super::{Deserializable, Serializable};
 12
 13#[cfg(feature = "std")]
 14use std::error::Error;
 15
 16// TEST DATA
 17// ================================================================================================
 18
 19const KEYS4: [u64; 4] = [0, 1, 2, 3];
 20const VALUES4: [RpoDigest; 4] = [int_to_node(1), int_to_node(2), int_to_node(3), int_to_node(4)];
 21
 22const KEYS8: [u64; 8] = [0, 1, 2, 3, 4, 5, 6, 7];
 23const VALUES8: [RpoDigest; 8] = [
 24    int_to_node(1),
 25    int_to_node(2),
 26    int_to_node(3),
 27    int_to_node(4),
 28    int_to_node(5),
 29    int_to_node(6),
 30    int_to_node(7),
 31    int_to_node(8),
 32];
 33
 34// TESTS
 35// ================================================================================================
 36
 37#[test]
 38fn test_root_not_in_store() -> Result<(), MerkleError> {
 39    let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
 40    let store = MerkleStore::from(&mtree);
 41    assert_eq!(
 42        store.get_node(VALUES4[0], NodeIndex::make(mtree.depth(), 0)),
 43        Err(MerkleError::RootNotInStore(VALUES4[0])),
 44        "Leaf 0 is not a root"
 45    );
 46    assert_eq!(
 47        store.get_path(VALUES4[0], NodeIndex::make(mtree.depth(), 0)),
 48        Err(MerkleError::RootNotInStore(VALUES4[0])),
 49        "Leaf 0 is not a root"
 50    );
 51
 52    Ok(())
 53}
 54
 55#[test]
 56fn test_merkle_tree() -> Result<(), MerkleError> {
 57    let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
 58    let store = MerkleStore::from(&mtree);
 59
 60    // STORE LEAVES ARE CORRECT -------------------------------------------------------------------
 61    // checks the leaves in the store corresponds to the expected values
 62    assert_eq!(
 63        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 0)),
 64        Ok(VALUES4[0]),
 65        "node 0 must be in the tree"
 66    );
 67    assert_eq!(
 68        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 1)),
 69        Ok(VALUES4[1]),
 70        "node 1 must be in the tree"
 71    );
 72    assert_eq!(
 73        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 2)),
 74        Ok(VALUES4[2]),
 75        "node 2 must be in the tree"
 76    );
 77    assert_eq!(
 78        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)),
 79        Ok(VALUES4[3]),
 80        "node 3 must be in the tree"
 81    );
 82
 83    // STORE LEAVES MATCH TREE --------------------------------------------------------------------
 84    // sanity check the values returned by the store and the tree
 85    assert_eq!(
 86        mtree.get_node(NodeIndex::make(mtree.depth(), 0)),
 87        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 0)),
 88        "node 0 must be the same for both MerkleTree and MerkleStore"
 89    );
 90    assert_eq!(
 91        mtree.get_node(NodeIndex::make(mtree.depth(), 1)),
 92        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 1)),
 93        "node 1 must be the same for both MerkleTree and MerkleStore"
 94    );
 95    assert_eq!(
 96        mtree.get_node(NodeIndex::make(mtree.depth(), 2)),
 97        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 2)),
 98        "node 2 must be the same for both MerkleTree and MerkleStore"
 99    );
100    assert_eq!(
101        mtree.get_node(NodeIndex::make(mtree.depth(), 3)),
102        store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3)),
103        "node 3 must be the same for both MerkleTree and MerkleStore"
104    );
105
106    // STORE MERKLE PATH MATCHS ==============================================================
107    // assert the merkle path returned by the store is the same as the one in the tree
108    let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 0)).unwrap();
109    assert_eq!(
110        VALUES4[0], result.value,
111        "Value for merkle path at index 0 must match leaf value"
112    );
113    assert_eq!(
114        mtree.get_path(NodeIndex::make(mtree.depth(), 0)),
115        Ok(result.path),
116        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
117    );
118
119    let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 1)).unwrap();
120    assert_eq!(
121        VALUES4[1], result.value,
122        "Value for merkle path at index 0 must match leaf value"
123    );
124    assert_eq!(
125        mtree.get_path(NodeIndex::make(mtree.depth(), 1)),
126        Ok(result.path),
127        "merkle path for index 1 must be the same for the MerkleTree and MerkleStore"
128    );
129
130    let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 2)).unwrap();
131    assert_eq!(
132        VALUES4[2], result.value,
133        "Value for merkle path at index 0 must match leaf value"
134    );
135    assert_eq!(
136        mtree.get_path(NodeIndex::make(mtree.depth(), 2)),
137        Ok(result.path),
138        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
139    );
140
141    let result = store.get_path(mtree.root(), NodeIndex::make(mtree.depth(), 3)).unwrap();
142    assert_eq!(
143        VALUES4[3], result.value,
144        "Value for merkle path at index 0 must match leaf value"
145    );
146    assert_eq!(
147        mtree.get_path(NodeIndex::make(mtree.depth(), 3)),
148        Ok(result.path),
149        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
150    );
151
152    Ok(())
153}
154
155#[test]
156fn test_empty_roots() {
157    let store = MerkleStore::default();
158    let mut root = RpoDigest::default();
159
160    for depth in 0..255 {
161        root = Rpo256::merge(&[root; 2]);
162        assert!(
163            store.get_node(root, NodeIndex::make(0, 0)).is_ok(),
164            "The root of the empty tree of depth {depth} must be registered"
165        );
166    }
167}
168
169#[test]
170fn test_leaf_paths_for_empty_trees() -> Result<(), MerkleError> {
171    let store = MerkleStore::default();
172
173    // Starts at 1 because leafs are not included in the store.
174    // Ends at 64 because it is not possible to represent an index of a depth greater than 64,
175    // because a u64 is used to index the leaf.
176    for depth in 1..64 {
177        let smt = SimpleSmt::new(depth)?;
178
179        let index = NodeIndex::make(depth, 0);
180        let store_path = store.get_path(smt.root(), index)?;
181        let smt_path = smt.get_path(index)?;
182        assert_eq!(
183            store_path.value,
184            RpoDigest::default(),
185            "the leaf of an empty tree is always ZERO"
186        );
187        assert_eq!(
188            store_path.path, smt_path,
189            "the returned merkle path does not match the computed values"
190        );
191        assert_eq!(
192            store_path.path.compute_root(depth.into(), RpoDigest::default()).unwrap(),
193            smt.root(),
194            "computed root from the path must match the empty tree root"
195        );
196    }
197
198    Ok(())
199}
200
201#[test]
202fn test_get_invalid_node() {
203    let mtree =
204        MerkleTree::new(digests_to_words(&VALUES4)).expect("creating a merkle tree must work");
205    let store = MerkleStore::from(&mtree);
206    let _ = store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3));
207}
208
209#[test]
210fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
211    let keys2: [u64; 2] = [0, 1];
212    let leaves2: [Word; 2] = [int_to_leaf(1), int_to_leaf(2)];
213    let smt = SimpleSmt::with_leaves(1, keys2.into_iter().zip(leaves2.into_iter())).unwrap();
214    let store = MerkleStore::from(&smt);
215
216    let idx = NodeIndex::make(1, 0);
217    assert_eq!(smt.get_node(idx).unwrap(), leaves2[0].into());
218    assert_eq!(store.get_node(smt.root(), idx).unwrap(), smt.get_node(idx).unwrap());
219
220    let idx = NodeIndex::make(1, 1);
221    assert_eq!(smt.get_node(idx).unwrap(), leaves2[1].into());
222    assert_eq!(store.get_node(smt.root(), idx).unwrap(), smt.get_node(idx).unwrap());
223
224    Ok(())
225}
226
227#[test]
228fn test_sparse_merkle_tree() -> Result<(), MerkleError> {
229    let smt = SimpleSmt::with_leaves(
230        SimpleSmt::MAX_DEPTH,
231        KEYS4.into_iter().zip(digests_to_words(&VALUES4).into_iter()),
232    )
233    .unwrap();
234
235    let store = MerkleStore::from(&smt);
236
237    // STORE LEAVES ARE CORRECT ==============================================================
238    // checks the leaves in the store corresponds to the expected values
239    assert_eq!(
240        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 0)),
241        Ok(VALUES4[0]),
242        "node 0 must be in the tree"
243    );
244    assert_eq!(
245        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 1)),
246        Ok(VALUES4[1]),
247        "node 1 must be in the tree"
248    );
249    assert_eq!(
250        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 2)),
251        Ok(VALUES4[2]),
252        "node 2 must be in the tree"
253    );
254    assert_eq!(
255        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 3)),
256        Ok(VALUES4[3]),
257        "node 3 must be in the tree"
258    );
259    assert_eq!(
260        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 4)),
261        Ok(RpoDigest::default()),
262        "unmodified node 4 must be ZERO"
263    );
264
265    // STORE LEAVES MATCH TREE ===============================================================
266    // sanity check the values returned by the store and the tree
267    assert_eq!(
268        smt.get_node(NodeIndex::make(smt.depth(), 0)),
269        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 0)),
270        "node 0 must be the same for both SparseMerkleTree and MerkleStore"
271    );
272    assert_eq!(
273        smt.get_node(NodeIndex::make(smt.depth(), 1)),
274        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 1)),
275        "node 1 must be the same for both SparseMerkleTree and MerkleStore"
276    );
277    assert_eq!(
278        smt.get_node(NodeIndex::make(smt.depth(), 2)),
279        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 2)),
280        "node 2 must be the same for both SparseMerkleTree and MerkleStore"
281    );
282    assert_eq!(
283        smt.get_node(NodeIndex::make(smt.depth(), 3)),
284        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 3)),
285        "node 3 must be the same for both SparseMerkleTree and MerkleStore"
286    );
287    assert_eq!(
288        smt.get_node(NodeIndex::make(smt.depth(), 4)),
289        store.get_node(smt.root(), NodeIndex::make(smt.depth(), 4)),
290        "node 4 must be the same for both SparseMerkleTree and MerkleStore"
291    );
292
293    // STORE MERKLE PATH MATCHS ==============================================================
294    // assert the merkle path returned by the store is the same as the one in the tree
295    let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 0)).unwrap();
296    assert_eq!(
297        VALUES4[0], result.value,
298        "Value for merkle path at index 0 must match leaf value"
299    );
300    assert_eq!(
301        smt.get_path(NodeIndex::make(smt.depth(), 0)),
302        Ok(result.path),
303        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
304    );
305
306    let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 1)).unwrap();
307    assert_eq!(
308        VALUES4[1], result.value,
309        "Value for merkle path at index 1 must match leaf value"
310    );
311    assert_eq!(
312        smt.get_path(NodeIndex::make(smt.depth(), 1)),
313        Ok(result.path),
314        "merkle path for index 1 must be the same for the MerkleTree and MerkleStore"
315    );
316
317    let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 2)).unwrap();
318    assert_eq!(
319        VALUES4[2], result.value,
320        "Value for merkle path at index 2 must match leaf value"
321    );
322    assert_eq!(
323        smt.get_path(NodeIndex::make(smt.depth(), 2)),
324        Ok(result.path),
325        "merkle path for index 2 must be the same for the MerkleTree and MerkleStore"
326    );
327
328    let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 3)).unwrap();
329    assert_eq!(
330        VALUES4[3], result.value,
331        "Value for merkle path at index 3 must match leaf value"
332    );
333    assert_eq!(
334        smt.get_path(NodeIndex::make(smt.depth(), 3)),
335        Ok(result.path),
336        "merkle path for index 3 must be the same for the MerkleTree and MerkleStore"
337    );
338
339    let result = store.get_path(smt.root(), NodeIndex::make(smt.depth(), 4)).unwrap();
340    assert_eq!(
341        RpoDigest::default(),
342        result.value,
343        "Value for merkle path at index 4 must match leaf value"
344    );
345    assert_eq!(
346        smt.get_path(NodeIndex::make(smt.depth(), 4)),
347        Ok(result.path),
348        "merkle path for index 4 must be the same for the MerkleTree and MerkleStore"
349    );
350
351    Ok(())
352}
353
354#[test]
355fn test_add_merkle_paths() -> Result<(), MerkleError> {
356    let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
357
358    let i0 = 0;
359    let p0 = mtree.get_path(NodeIndex::make(2, i0)).unwrap();
360
361    let i1 = 1;
362    let p1 = mtree.get_path(NodeIndex::make(2, i1)).unwrap();
363
364    let i2 = 2;
365    let p2 = mtree.get_path(NodeIndex::make(2, i2)).unwrap();
366
367    let i3 = 3;
368    let p3 = mtree.get_path(NodeIndex::make(2, i3)).unwrap();
369
370    let paths = [
371        (i0, VALUES4[i0 as usize], p0),
372        (i1, VALUES4[i1 as usize], p1),
373        (i2, VALUES4[i2 as usize], p2),
374        (i3, VALUES4[i3 as usize], p3),
375    ];
376
377    let mut store = MerkleStore::default();
378    store.add_merkle_paths(paths.clone()).expect("the valid paths must work");
379
380    let pmt = PartialMerkleTree::with_paths(paths).unwrap();
381
382    // STORE LEAVES ARE CORRECT ==============================================================
383    // checks the leaves in the store corresponds to the expected values
384    assert_eq!(
385        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 0)),
386        Ok(VALUES4[0]),
387        "node 0 must be in the pmt"
388    );
389    assert_eq!(
390        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 1)),
391        Ok(VALUES4[1]),
392        "node 1 must be in the pmt"
393    );
394    assert_eq!(
395        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 2)),
396        Ok(VALUES4[2]),
397        "node 2 must be in the pmt"
398    );
399    assert_eq!(
400        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 3)),
401        Ok(VALUES4[3]),
402        "node 3 must be in the pmt"
403    );
404
405    // STORE LEAVES MATCH PMT ================================================================
406    // sanity check the values returned by the store and the pmt
407    assert_eq!(
408        pmt.get_node(NodeIndex::make(pmt.max_depth(), 0)),
409        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 0)),
410        "node 0 must be the same for both PartialMerkleTree and MerkleStore"
411    );
412    assert_eq!(
413        pmt.get_node(NodeIndex::make(pmt.max_depth(), 1)),
414        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 1)),
415        "node 1 must be the same for both PartialMerkleTree and MerkleStore"
416    );
417    assert_eq!(
418        pmt.get_node(NodeIndex::make(pmt.max_depth(), 2)),
419        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 2)),
420        "node 2 must be the same for both PartialMerkleTree and MerkleStore"
421    );
422    assert_eq!(
423        pmt.get_node(NodeIndex::make(pmt.max_depth(), 3)),
424        store.get_node(pmt.root(), NodeIndex::make(pmt.max_depth(), 3)),
425        "node 3 must be the same for both PartialMerkleTree and MerkleStore"
426    );
427
428    // STORE MERKLE PATH MATCHS ==============================================================
429    // assert the merkle path returned by the store is the same as the one in the pmt
430    let result = store.get_path(pmt.root(), NodeIndex::make(pmt.max_depth(), 0)).unwrap();
431    assert_eq!(
432        VALUES4[0], result.value,
433        "Value for merkle path at index 0 must match leaf value"
434    );
435    assert_eq!(
436        pmt.get_path(NodeIndex::make(pmt.max_depth(), 0)),
437        Ok(result.path),
438        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
439    );
440
441    let result = store.get_path(pmt.root(), NodeIndex::make(pmt.max_depth(), 1)).unwrap();
442    assert_eq!(
443        VALUES4[1], result.value,
444        "Value for merkle path at index 0 must match leaf value"
445    );
446    assert_eq!(
447        pmt.get_path(NodeIndex::make(pmt.max_depth(), 1)),
448        Ok(result.path),
449        "merkle path for index 1 must be the same for the MerkleTree and MerkleStore"
450    );
451
452    let result = store.get_path(pmt.root(), NodeIndex::make(pmt.max_depth(), 2)).unwrap();
453    assert_eq!(
454        VALUES4[2], result.value,
455        "Value for merkle path at index 0 must match leaf value"
456    );
457    assert_eq!(
458        pmt.get_path(NodeIndex::make(pmt.max_depth(), 2)),
459        Ok(result.path),
460        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
461    );
462
463    let result = store.get_path(pmt.root(), NodeIndex::make(pmt.max_depth(), 3)).unwrap();
464    assert_eq!(
465        VALUES4[3], result.value,
466        "Value for merkle path at index 0 must match leaf value"
467    );
468    assert_eq!(
469        pmt.get_path(NodeIndex::make(pmt.max_depth(), 3)),
470        Ok(result.path),
471        "merkle path for index 0 must be the same for the MerkleTree and MerkleStore"
472    );
473
474    Ok(())
475}
476
477#[test]
478fn wont_open_to_different_depth_root() {
479    let empty = EmptySubtreeRoots::empty_hashes(64);
480    let a = [ONE; 4];
481    let b = [Felt::new(2); 4];
482
483    // Compute the root for a different depth. We cherry-pick this specific depth to prevent a
484    // regression to a bug in the past that allowed the user to fetch a node at a depth lower than
485    // the inserted path of a Merkle tree.
486    let mut root = Rpo256::merge(&[a.into(), b.into()]);
487    for depth in (1..=63).rev() {
488        root = Rpo256::merge(&[root, empty[depth]]);
489    }
490
491    // For this example, the depth of the Merkle tree is 1, as we have only two leaves. Here we
492    // attempt to fetch a node on the maximum depth, and it should fail because the root shouldn't
493    // exist for the set.
494    let mtree = MerkleTree::new(vec![a, b]).unwrap();
495    let store = MerkleStore::from(&mtree);
496    let index = NodeIndex::root();
497    let err = store.get_node(root, index).err().unwrap();
498    assert_eq!(err, MerkleError::RootNotInStore(root));
499}
500
501#[test]
502fn store_path_opens_from_leaf() {
503    let a = [ONE; 4];
504    let b = [Felt::new(2); 4];
505    let c = [Felt::new(3); 4];
506    let d = [Felt::new(4); 4];
507    let e = [Felt::new(5); 4];
508    let f = [Felt::new(6); 4];
509    let g = [Felt::new(7); 4];
510    let h = [Felt::new(8); 4];
511
512    let i = Rpo256::merge(&[a.into(), b.into()]);
513    let j = Rpo256::merge(&[c.into(), d.into()]);
514    let k = Rpo256::merge(&[e.into(), f.into()]);
515    let l = Rpo256::merge(&[g.into(), h.into()]);
516
517    let m = Rpo256::merge(&[i, j]);
518    let n = Rpo256::merge(&[k, l]);
519
520    let root = Rpo256::merge(&[m, n]);
521
522    let mtree = MerkleTree::new(vec![a, b, c, d, e, f, g, h]).unwrap();
523    let store = MerkleStore::from(&mtree);
524    let path = store.get_path(root, NodeIndex::make(3, 1)).unwrap().path;
525
526    let expected = MerklePath::new([a.into(), j, n].to_vec());
527    assert_eq!(path, expected);
528}
529
530#[test]
531fn test_set_node() -> Result<(), MerkleError> {
532    let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
533    let mut store = MerkleStore::from(&mtree);
534    let value = int_to_node(42);
535    let index = NodeIndex::make(mtree.depth(), 0);
536    let new_root = store.set_node(mtree.root(), index, value)?.root;
537    assert_eq!(store.get_node(new_root, index), Ok(value), "Value must have changed");
538
539    Ok(())
540}
541
542#[test]
543fn test_constructors() -> Result<(), MerkleError> {
544    let mtree = MerkleTree::new(digests_to_words(&VALUES4))?;
545    let store = MerkleStore::from(&mtree);
546
547    let depth = mtree.depth();
548    let leaves = 2u64.pow(depth.into());
549    for index in 0..leaves {
550        let index = NodeIndex::make(depth, index);
551        let value_path = store.get_path(mtree.root(), index)?;
552        assert_eq!(mtree.get_path(index)?, value_path.path);
553    }
554
555    let depth = 32;
556    let smt = SimpleSmt::with_leaves(
557        depth,
558        KEYS4.into_iter().zip(digests_to_words(&VALUES4).into_iter()),
559    )
560    .unwrap();
561    let store = MerkleStore::from(&smt);
562    let depth = smt.depth();
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.get_path(index)?, 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 insersecting 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    let smtree = SimpleSmt::with_leaves(
887        64,
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(smtree.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(smtree.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!(node, smtree.get_leaf(index_2.value()).unwrap().into());
924
925    // assert that is doesnt contain nodes that were not recorded
926    let not_recorded_index = NodeIndex::new(smtree.depth(), 4).unwrap();
927    assert!(merkle_store.get_node(smtree.root(), not_recorded_index).is_err());
928    assert!(smtree.get_node(not_recorded_index).is_ok());
929}