arnaucube
arnaucube/miden-cryptoPublic Self-hosted repository
miden-crypto/src/merkle/mmr/tests.rs
src/merkle/mmr/tests.rsRaw
  1use super::{
  2    super::{InnerNodeInfo, Rpo256, RpoDigest, Vec},
  3    bit::TrueBitPositionIterator,
  4    full::high_bitmask,
  5    leaf_to_corresponding_tree, nodes_in_forest, Mmr, MmrPeaks, PartialMmr,
  6};
  7use crate::{
  8    merkle::{int_to_node, InOrderIndex, MerklePath, MerkleTree, MmrProof, NodeIndex},
  9    Felt, Word,
 10};
 11
 12#[test]
 13fn test_position_equal_or_higher_than_leafs_is_never_contained() {
 14    let empty_forest = 0;
 15    for pos in 1..1024 {
 16        // pos is index, 0 based
 17        // tree is a length counter, 1 based
 18        // so a valid pos is always smaller, not equal, to tree
 19        assert_eq!(leaf_to_corresponding_tree(pos, pos), None);
 20        assert_eq!(leaf_to_corresponding_tree(pos, pos - 1), None);
 21        // and empty forest has no trees, so no position is valid
 22        assert_eq!(leaf_to_corresponding_tree(pos, empty_forest), None);
 23    }
 24}
 25
 26#[test]
 27fn test_position_zero_is_always_contained_by_the_highest_tree() {
 28    for leaves in 1..1024usize {
 29        let tree = leaves.ilog2();
 30        assert_eq!(leaf_to_corresponding_tree(0, leaves), Some(tree));
 31    }
 32}
 33
 34#[test]
 35fn test_leaf_to_corresponding_tree() {
 36    assert_eq!(leaf_to_corresponding_tree(0, 0b0001), Some(0));
 37    assert_eq!(leaf_to_corresponding_tree(0, 0b0010), Some(1));
 38    assert_eq!(leaf_to_corresponding_tree(0, 0b0011), Some(1));
 39    assert_eq!(leaf_to_corresponding_tree(0, 0b1011), Some(3));
 40
 41    // position one is always owned by the left-most tree
 42    assert_eq!(leaf_to_corresponding_tree(1, 0b0010), Some(1));
 43    assert_eq!(leaf_to_corresponding_tree(1, 0b0011), Some(1));
 44    assert_eq!(leaf_to_corresponding_tree(1, 0b1011), Some(3));
 45
 46    // position two starts as its own root, and then it is merged with the left-most tree
 47    assert_eq!(leaf_to_corresponding_tree(2, 0b0011), Some(0));
 48    assert_eq!(leaf_to_corresponding_tree(2, 0b0100), Some(2));
 49    assert_eq!(leaf_to_corresponding_tree(2, 0b1011), Some(3));
 50
 51    // position tree is merged on the left-most tree
 52    assert_eq!(leaf_to_corresponding_tree(3, 0b0011), None);
 53    assert_eq!(leaf_to_corresponding_tree(3, 0b0100), Some(2));
 54    assert_eq!(leaf_to_corresponding_tree(3, 0b1011), Some(3));
 55
 56    assert_eq!(leaf_to_corresponding_tree(4, 0b0101), Some(0));
 57    assert_eq!(leaf_to_corresponding_tree(4, 0b0110), Some(1));
 58    assert_eq!(leaf_to_corresponding_tree(4, 0b0111), Some(1));
 59    assert_eq!(leaf_to_corresponding_tree(4, 0b1000), Some(3));
 60
 61    assert_eq!(leaf_to_corresponding_tree(12, 0b01101), Some(0));
 62    assert_eq!(leaf_to_corresponding_tree(12, 0b01110), Some(1));
 63    assert_eq!(leaf_to_corresponding_tree(12, 0b01111), Some(1));
 64    assert_eq!(leaf_to_corresponding_tree(12, 0b10000), Some(4));
 65}
 66
 67#[test]
 68fn test_high_bitmask() {
 69    assert_eq!(high_bitmask(0), usize::MAX);
 70    assert_eq!(high_bitmask(1), usize::MAX << 1);
 71    assert_eq!(high_bitmask(usize::BITS - 2), 0b11usize.rotate_right(2));
 72    assert_eq!(high_bitmask(usize::BITS - 1), 0b1usize.rotate_right(1));
 73    assert_eq!(high_bitmask(usize::BITS), 0, "overflow should be handled");
 74}
 75
 76#[test]
 77fn test_nodes_in_forest() {
 78    assert_eq!(nodes_in_forest(0b0000), 0);
 79    assert_eq!(nodes_in_forest(0b0001), 1);
 80    assert_eq!(nodes_in_forest(0b0010), 3);
 81    assert_eq!(nodes_in_forest(0b0011), 4);
 82    assert_eq!(nodes_in_forest(0b0100), 7);
 83    assert_eq!(nodes_in_forest(0b0101), 8);
 84    assert_eq!(nodes_in_forest(0b0110), 10);
 85    assert_eq!(nodes_in_forest(0b0111), 11);
 86    assert_eq!(nodes_in_forest(0b1000), 15);
 87    assert_eq!(nodes_in_forest(0b1001), 16);
 88    assert_eq!(nodes_in_forest(0b1010), 18);
 89    assert_eq!(nodes_in_forest(0b1011), 19);
 90}
 91
 92#[test]
 93fn test_nodes_in_forest_single_bit() {
 94    assert_eq!(nodes_in_forest(2usize.pow(0)), 2usize.pow(1) - 1);
 95    assert_eq!(nodes_in_forest(2usize.pow(1)), 2usize.pow(2) - 1);
 96    assert_eq!(nodes_in_forest(2usize.pow(2)), 2usize.pow(3) - 1);
 97    assert_eq!(nodes_in_forest(2usize.pow(3)), 2usize.pow(4) - 1);
 98
 99    for bit in 0..(usize::BITS - 1) {
100        let size = 2usize.pow(bit + 1) - 1;
101        assert_eq!(nodes_in_forest(1usize << bit), size);
102    }
103}
104
105const LEAVES: [RpoDigest; 7] = [
106    int_to_node(0),
107    int_to_node(1),
108    int_to_node(2),
109    int_to_node(3),
110    int_to_node(4),
111    int_to_node(5),
112    int_to_node(6),
113];
114
115#[test]
116fn test_mmr_simple() {
117    let mut postorder = Vec::new();
118    postorder.push(LEAVES[0]);
119    postorder.push(LEAVES[1]);
120    postorder.push(merge(LEAVES[0], LEAVES[1]));
121    postorder.push(LEAVES[2]);
122    postorder.push(LEAVES[3]);
123    postorder.push(merge(LEAVES[2], LEAVES[3]));
124    postorder.push(merge(postorder[2], postorder[5]));
125    postorder.push(LEAVES[4]);
126    postorder.push(LEAVES[5]);
127    postorder.push(merge(LEAVES[4], LEAVES[5]));
128    postorder.push(LEAVES[6]);
129
130    let mut mmr = Mmr::new();
131    assert_eq!(mmr.forest(), 0);
132    assert_eq!(mmr.nodes.len(), 0);
133
134    mmr.add(LEAVES[0]);
135    assert_eq!(mmr.forest(), 1);
136    assert_eq!(mmr.nodes.len(), 1);
137    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
138
139    let acc = mmr.accumulator();
140    assert_eq!(acc.num_leaves(), 1);
141    assert_eq!(acc.peaks(), &[postorder[0]]);
142
143    mmr.add(LEAVES[1]);
144    assert_eq!(mmr.forest(), 2);
145    assert_eq!(mmr.nodes.len(), 3);
146    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
147
148    let acc = mmr.accumulator();
149    assert_eq!(acc.num_leaves(), 2);
150    assert_eq!(acc.peaks(), &[postorder[2]]);
151
152    mmr.add(LEAVES[2]);
153    assert_eq!(mmr.forest(), 3);
154    assert_eq!(mmr.nodes.len(), 4);
155    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
156
157    let acc = mmr.accumulator();
158    assert_eq!(acc.num_leaves(), 3);
159    assert_eq!(acc.peaks(), &[postorder[2], postorder[3]]);
160
161    mmr.add(LEAVES[3]);
162    assert_eq!(mmr.forest(), 4);
163    assert_eq!(mmr.nodes.len(), 7);
164    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
165
166    let acc = mmr.accumulator();
167    assert_eq!(acc.num_leaves(), 4);
168    assert_eq!(acc.peaks(), &[postorder[6]]);
169
170    mmr.add(LEAVES[4]);
171    assert_eq!(mmr.forest(), 5);
172    assert_eq!(mmr.nodes.len(), 8);
173    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
174
175    let acc = mmr.accumulator();
176    assert_eq!(acc.num_leaves(), 5);
177    assert_eq!(acc.peaks(), &[postorder[6], postorder[7]]);
178
179    mmr.add(LEAVES[5]);
180    assert_eq!(mmr.forest(), 6);
181    assert_eq!(mmr.nodes.len(), 10);
182    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
183
184    let acc = mmr.accumulator();
185    assert_eq!(acc.num_leaves(), 6);
186    assert_eq!(acc.peaks(), &[postorder[6], postorder[9]]);
187
188    mmr.add(LEAVES[6]);
189    assert_eq!(mmr.forest(), 7);
190    assert_eq!(mmr.nodes.len(), 11);
191    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
192
193    let acc = mmr.accumulator();
194    assert_eq!(acc.num_leaves(), 7);
195    assert_eq!(acc.peaks(), &[postorder[6], postorder[9], postorder[10]]);
196}
197
198#[test]
199fn test_mmr_open() {
200    let mmr: Mmr = LEAVES.into();
201    let h01 = merge(LEAVES[0], LEAVES[1]);
202    let h23 = merge(LEAVES[2], LEAVES[3]);
203
204    // node at pos 7 is the root
205    assert!(mmr.open(7).is_err(), "Element 7 is not in the tree, result should be None");
206
207    // node at pos 6 is the root
208    let empty: MerklePath = MerklePath::new(vec![]);
209    let opening = mmr
210        .open(6)
211        .expect("Element 6 is contained in the tree, expected an opening result.");
212    assert_eq!(opening.merkle_path, empty);
213    assert_eq!(opening.forest, mmr.forest);
214    assert_eq!(opening.position, 6);
215    assert!(
216        mmr.accumulator().verify(LEAVES[6], opening),
217        "MmrProof should be valid for the current accumulator."
218    );
219
220    // nodes 4,5 are depth 1
221    let root_to_path = MerklePath::new(vec![LEAVES[4]]);
222    let opening = mmr
223        .open(5)
224        .expect("Element 5 is contained in the tree, expected an opening result.");
225    assert_eq!(opening.merkle_path, root_to_path);
226    assert_eq!(opening.forest, mmr.forest);
227    assert_eq!(opening.position, 5);
228    assert!(
229        mmr.accumulator().verify(LEAVES[5], opening),
230        "MmrProof should be valid for the current accumulator."
231    );
232
233    let root_to_path = MerklePath::new(vec![LEAVES[5]]);
234    let opening = mmr
235        .open(4)
236        .expect("Element 4 is contained in the tree, expected an opening result.");
237    assert_eq!(opening.merkle_path, root_to_path);
238    assert_eq!(opening.forest, mmr.forest);
239    assert_eq!(opening.position, 4);
240    assert!(
241        mmr.accumulator().verify(LEAVES[4], opening),
242        "MmrProof should be valid for the current accumulator."
243    );
244
245    // nodes 0,1,2,3 are detph 2
246    let root_to_path = MerklePath::new(vec![LEAVES[2], h01]);
247    let opening = mmr
248        .open(3)
249        .expect("Element 3 is contained in the tree, expected an opening result.");
250    assert_eq!(opening.merkle_path, root_to_path);
251    assert_eq!(opening.forest, mmr.forest);
252    assert_eq!(opening.position, 3);
253    assert!(
254        mmr.accumulator().verify(LEAVES[3], opening),
255        "MmrProof should be valid for the current accumulator."
256    );
257
258    let root_to_path = MerklePath::new(vec![LEAVES[3], h01]);
259    let opening = mmr
260        .open(2)
261        .expect("Element 2 is contained in the tree, expected an opening result.");
262    assert_eq!(opening.merkle_path, root_to_path);
263    assert_eq!(opening.forest, mmr.forest);
264    assert_eq!(opening.position, 2);
265    assert!(
266        mmr.accumulator().verify(LEAVES[2], opening),
267        "MmrProof should be valid for the current accumulator."
268    );
269
270    let root_to_path = MerklePath::new(vec![LEAVES[0], h23]);
271    let opening = mmr
272        .open(1)
273        .expect("Element 1 is contained in the tree, expected an opening result.");
274    assert_eq!(opening.merkle_path, root_to_path);
275    assert_eq!(opening.forest, mmr.forest);
276    assert_eq!(opening.position, 1);
277    assert!(
278        mmr.accumulator().verify(LEAVES[1], opening),
279        "MmrProof should be valid for the current accumulator."
280    );
281
282    let root_to_path = MerklePath::new(vec![LEAVES[1], h23]);
283    let opening = mmr
284        .open(0)
285        .expect("Element 0 is contained in the tree, expected an opening result.");
286    assert_eq!(opening.merkle_path, root_to_path);
287    assert_eq!(opening.forest, mmr.forest);
288    assert_eq!(opening.position, 0);
289    assert!(
290        mmr.accumulator().verify(LEAVES[0], opening),
291        "MmrProof should be valid for the current accumulator."
292    );
293}
294
295/// Tests the openings of a simple Mmr with a single tree of depth 8.
296#[test]
297fn test_mmr_open_eight() {
298    let leaves = [
299        int_to_node(0),
300        int_to_node(1),
301        int_to_node(2),
302        int_to_node(3),
303        int_to_node(4),
304        int_to_node(5),
305        int_to_node(6),
306        int_to_node(7),
307    ];
308
309    let mtree: MerkleTree = leaves.as_slice().try_into().unwrap();
310    let forest = leaves.len();
311    let mmr: Mmr = leaves.into();
312    let root = mtree.root();
313
314    let position = 0;
315    let proof = mmr.open(position).unwrap();
316    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
317    assert_eq!(proof, MmrProof { forest, position, merkle_path });
318    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
319
320    let position = 1;
321    let proof = mmr.open(position).unwrap();
322    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
323    assert_eq!(proof, MmrProof { forest, position, merkle_path });
324    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
325
326    let position = 2;
327    let proof = mmr.open(position).unwrap();
328    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
329    assert_eq!(proof, MmrProof { forest, position, merkle_path });
330    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
331
332    let position = 3;
333    let proof = mmr.open(position).unwrap();
334    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
335    assert_eq!(proof, MmrProof { forest, position, merkle_path });
336    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
337
338    let position = 4;
339    let proof = mmr.open(position).unwrap();
340    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
341    assert_eq!(proof, MmrProof { forest, position, merkle_path });
342    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
343
344    let position = 5;
345    let proof = mmr.open(position).unwrap();
346    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
347    assert_eq!(proof, MmrProof { forest, position, merkle_path });
348    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
349
350    let position = 6;
351    let proof = mmr.open(position).unwrap();
352    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
353    assert_eq!(proof, MmrProof { forest, position, merkle_path });
354    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
355
356    let position = 7;
357    let proof = mmr.open(position).unwrap();
358    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
359    assert_eq!(proof, MmrProof { forest, position, merkle_path });
360    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
361}
362
363/// Tests the openings of Mmr with a 3 trees of depths 4, 2, and 1.
364#[test]
365fn test_mmr_open_seven() {
366    let mtree1: MerkleTree = LEAVES[..4].try_into().unwrap();
367    let mtree2: MerkleTree = LEAVES[4..6].try_into().unwrap();
368
369    let forest = LEAVES.len();
370    let mmr: Mmr = LEAVES.into();
371
372    let position = 0;
373    let proof = mmr.open(position).unwrap();
374    let merkle_path: MerklePath =
375        mtree1.get_path(NodeIndex::new(2, position as u64).unwrap()).unwrap();
376    assert_eq!(proof, MmrProof { forest, position, merkle_path });
377    assert_eq!(proof.merkle_path.compute_root(0, LEAVES[0]).unwrap(), mtree1.root());
378
379    let position = 1;
380    let proof = mmr.open(position).unwrap();
381    let merkle_path: MerklePath =
382        mtree1.get_path(NodeIndex::new(2, position as u64).unwrap()).unwrap();
383    assert_eq!(proof, MmrProof { forest, position, merkle_path });
384    assert_eq!(proof.merkle_path.compute_root(1, LEAVES[1]).unwrap(), mtree1.root());
385
386    let position = 2;
387    let proof = mmr.open(position).unwrap();
388    let merkle_path: MerklePath =
389        mtree1.get_path(NodeIndex::new(2, position as u64).unwrap()).unwrap();
390    assert_eq!(proof, MmrProof { forest, position, merkle_path });
391    assert_eq!(proof.merkle_path.compute_root(2, LEAVES[2]).unwrap(), mtree1.root());
392
393    let position = 3;
394    let proof = mmr.open(position).unwrap();
395    let merkle_path: MerklePath =
396        mtree1.get_path(NodeIndex::new(2, position as u64).unwrap()).unwrap();
397    assert_eq!(proof, MmrProof { forest, position, merkle_path });
398    assert_eq!(proof.merkle_path.compute_root(3, LEAVES[3]).unwrap(), mtree1.root());
399
400    let position = 4;
401    let proof = mmr.open(position).unwrap();
402    let merkle_path: MerklePath = mtree2.get_path(NodeIndex::new(1, 0u64).unwrap()).unwrap();
403    assert_eq!(proof, MmrProof { forest, position, merkle_path });
404    assert_eq!(proof.merkle_path.compute_root(0, LEAVES[4]).unwrap(), mtree2.root());
405
406    let position = 5;
407    let proof = mmr.open(position).unwrap();
408    let merkle_path: MerklePath = mtree2.get_path(NodeIndex::new(1, 1u64).unwrap()).unwrap();
409    assert_eq!(proof, MmrProof { forest, position, merkle_path });
410    assert_eq!(proof.merkle_path.compute_root(1, LEAVES[5]).unwrap(), mtree2.root());
411
412    let position = 6;
413    let proof = mmr.open(position).unwrap();
414    let merkle_path: MerklePath = [].as_ref().into();
415    assert_eq!(proof, MmrProof { forest, position, merkle_path });
416    assert_eq!(proof.merkle_path.compute_root(0, LEAVES[6]).unwrap(), LEAVES[6]);
417}
418
419#[test]
420fn test_mmr_get() {
421    let mmr: Mmr = LEAVES.into();
422    assert_eq!(mmr.get(0).unwrap(), LEAVES[0], "value at pos 0 must correspond");
423    assert_eq!(mmr.get(1).unwrap(), LEAVES[1], "value at pos 1 must correspond");
424    assert_eq!(mmr.get(2).unwrap(), LEAVES[2], "value at pos 2 must correspond");
425    assert_eq!(mmr.get(3).unwrap(), LEAVES[3], "value at pos 3 must correspond");
426    assert_eq!(mmr.get(4).unwrap(), LEAVES[4], "value at pos 4 must correspond");
427    assert_eq!(mmr.get(5).unwrap(), LEAVES[5], "value at pos 5 must correspond");
428    assert_eq!(mmr.get(6).unwrap(), LEAVES[6], "value at pos 6 must correspond");
429    assert!(mmr.get(7).is_err());
430}
431
432#[test]
433fn test_mmr_invariants() {
434    let mut mmr = Mmr::new();
435    for v in 1..=1028 {
436        mmr.add(int_to_node(v));
437        let accumulator = mmr.accumulator();
438        assert_eq!(v as usize, mmr.forest(), "MMR leaf count must increase by one on every add");
439        assert_eq!(
440            v as usize,
441            accumulator.num_leaves(),
442            "MMR and its accumulator must match leaves count"
443        );
444        assert_eq!(
445            accumulator.num_leaves().count_ones() as usize,
446            accumulator.peaks().len(),
447            "bits on leaves must match the number of peaks"
448        );
449
450        let expected_nodes: usize = TrueBitPositionIterator::new(mmr.forest())
451            .map(|bit_pos| nodes_in_forest(1 << bit_pos))
452            .sum();
453
454        assert_eq!(
455            expected_nodes,
456            mmr.nodes.len(),
457            "the sum of every tree size must be equal to the number of nodes in the MMR (forest: {:b})",
458            mmr.forest(),
459        );
460    }
461}
462
463#[test]
464fn test_bit_position_iterator() {
465    assert_eq!(TrueBitPositionIterator::new(0).count(), 0);
466    assert_eq!(TrueBitPositionIterator::new(0).rev().count(), 0);
467
468    assert_eq!(TrueBitPositionIterator::new(1).collect::<Vec<u32>>(), vec![0]);
469    assert_eq!(TrueBitPositionIterator::new(1).rev().collect::<Vec<u32>>(), vec![0],);
470
471    assert_eq!(TrueBitPositionIterator::new(2).collect::<Vec<u32>>(), vec![1]);
472    assert_eq!(TrueBitPositionIterator::new(2).rev().collect::<Vec<u32>>(), vec![1],);
473
474    assert_eq!(TrueBitPositionIterator::new(3).collect::<Vec<u32>>(), vec![0, 1],);
475    assert_eq!(TrueBitPositionIterator::new(3).rev().collect::<Vec<u32>>(), vec![1, 0],);
476
477    assert_eq!(
478        TrueBitPositionIterator::new(0b11010101).collect::<Vec<u32>>(),
479        vec![0, 2, 4, 6, 7],
480    );
481    assert_eq!(
482        TrueBitPositionIterator::new(0b11010101).rev().collect::<Vec<u32>>(),
483        vec![7, 6, 4, 2, 0],
484    );
485}
486
487#[test]
488fn test_mmr_inner_nodes() {
489    let mmr: Mmr = LEAVES.into();
490    let nodes: Vec<InnerNodeInfo> = mmr.inner_nodes().collect();
491
492    let h01 = Rpo256::merge(&[LEAVES[0], LEAVES[1]]);
493    let h23 = Rpo256::merge(&[LEAVES[2], LEAVES[3]]);
494    let h0123 = Rpo256::merge(&[h01, h23]);
495    let h45 = Rpo256::merge(&[LEAVES[4], LEAVES[5]]);
496    let postorder = vec![
497        InnerNodeInfo {
498            value: h01,
499            left: LEAVES[0],
500            right: LEAVES[1],
501        },
502        InnerNodeInfo {
503            value: h23,
504            left: LEAVES[2],
505            right: LEAVES[3],
506        },
507        InnerNodeInfo { value: h0123, left: h01, right: h23 },
508        InnerNodeInfo {
509            value: h45,
510            left: LEAVES[4],
511            right: LEAVES[5],
512        },
513    ];
514
515    assert_eq!(postorder, nodes);
516}
517
518#[test]
519fn test_mmr_hash_peaks() {
520    let mmr: Mmr = LEAVES.into();
521    let peaks = mmr.accumulator();
522
523    let first_peak = Rpo256::merge(&[
524        Rpo256::merge(&[LEAVES[0], LEAVES[1]]),
525        Rpo256::merge(&[LEAVES[2], LEAVES[3]]),
526    ]);
527    let second_peak = Rpo256::merge(&[LEAVES[4], LEAVES[5]]);
528    let third_peak = LEAVES[6];
529
530    // minimum length is 16
531    let mut expected_peaks = [first_peak, second_peak, third_peak].to_vec();
532    expected_peaks.resize(16, RpoDigest::default());
533    assert_eq!(
534        peaks.hash_peaks(),
535        *Rpo256::hash_elements(&digests_to_elements(&expected_peaks))
536    );
537}
538
539#[test]
540fn test_mmr_peaks_hash_less_than_16() {
541    let mut peaks = Vec::new();
542
543    for i in 0..16 {
544        peaks.push(int_to_node(i));
545
546        let num_leaves = (1 << peaks.len()) - 1;
547        let accumulator = MmrPeaks::new(num_leaves, peaks.clone()).unwrap();
548
549        // minimum length is 16
550        let mut expected_peaks = peaks.clone();
551        expected_peaks.resize(16, RpoDigest::default());
552        assert_eq!(
553            accumulator.hash_peaks(),
554            *Rpo256::hash_elements(&digests_to_elements(&expected_peaks))
555        );
556    }
557}
558
559#[test]
560fn test_mmr_peaks_hash_odd() {
561    let peaks: Vec<_> = (0..=17).map(int_to_node).collect();
562
563    let num_leaves = (1 << peaks.len()) - 1;
564    let accumulator = MmrPeaks::new(num_leaves, peaks.clone()).unwrap();
565
566    // odd length bigger than 16 is padded to the next even number
567    let mut expected_peaks = peaks;
568    expected_peaks.resize(18, RpoDigest::default());
569    assert_eq!(
570        accumulator.hash_peaks(),
571        *Rpo256::hash_elements(&digests_to_elements(&expected_peaks))
572    );
573}
574
575#[test]
576fn test_mmr_updates() {
577    let mmr: Mmr = LEAVES.into();
578    let acc = mmr.accumulator();
579
580    // original_forest can't have more elements
581    assert!(
582        mmr.get_delta(LEAVES.len() + 1).is_err(),
583        "Can not provide updates for a newer Mmr"
584    );
585
586    // if the number of elements is the same there is no change
587    assert!(
588        mmr.get_delta(LEAVES.len()).unwrap().data.is_empty(),
589        "There are no updates for the same Mmr version"
590    );
591
592    // missing the last element added, which is itself a tree peak
593    assert_eq!(mmr.get_delta(6).unwrap().data, vec![acc.peaks()[2]], "one peak");
594
595    // missing the sibling to complete the tree of depth 2, and the last element
596    assert_eq!(
597        mmr.get_delta(5).unwrap().data,
598        vec![LEAVES[5], acc.peaks()[2]],
599        "one sibling, one peak"
600    );
601
602    // missing the whole last two trees, only send the peaks
603    assert_eq!(
604        mmr.get_delta(4).unwrap().data,
605        vec![acc.peaks()[1], acc.peaks()[2]],
606        "two peaks"
607    );
608
609    // missing the sibling to complete the first tree, and the two last trees
610    assert_eq!(
611        mmr.get_delta(3).unwrap().data,
612        vec![LEAVES[3], acc.peaks()[1], acc.peaks()[2]],
613        "one sibling, two peaks"
614    );
615
616    // missing half of the first tree, only send the computed element (not the leaves), and the new
617    // peaks
618    assert_eq!(
619        mmr.get_delta(2).unwrap().data,
620        vec![mmr.nodes[5], acc.peaks()[1], acc.peaks()[2]],
621        "one sibling, two peaks"
622    );
623
624    assert_eq!(
625        mmr.get_delta(1).unwrap().data,
626        vec![LEAVES[1], mmr.nodes[5], acc.peaks()[1], acc.peaks()[2]],
627        "one sibling, two peaks"
628    );
629
630    assert_eq!(&mmr.get_delta(0).unwrap().data, acc.peaks(), "all peaks");
631}
632
633#[test]
634fn test_partial_mmr_simple() {
635    let mmr: Mmr = LEAVES.into();
636    let acc = mmr.accumulator();
637    let mut partial: PartialMmr = acc.clone().into();
638
639    // check initial state of the partial mmr
640    assert_eq!(partial.peaks(), acc.peaks());
641    assert_eq!(partial.forest(), acc.num_leaves());
642    assert_eq!(partial.forest(), LEAVES.len());
643    assert_eq!(partial.peaks().len(), 3);
644    assert_eq!(partial.nodes.len(), 0);
645
646    // check state after adding tracking one element
647    let proof1 = mmr.open(0).unwrap();
648    let el1 = mmr.get(proof1.position).unwrap();
649    partial.add(proof1.position, el1, &proof1.merkle_path).unwrap();
650
651    // check the number of nodes increased by the number of nodes in the proof
652    assert_eq!(partial.nodes.len(), proof1.merkle_path.len());
653    // check the values match
654    let idx = InOrderIndex::from_leaf_pos(proof1.position);
655    assert_eq!(partial.nodes[&idx.sibling()], proof1.merkle_path[0]);
656    let idx = idx.parent();
657    assert_eq!(partial.nodes[&idx.sibling()], proof1.merkle_path[1]);
658
659    let proof2 = mmr.open(1).unwrap();
660    let el2 = mmr.get(proof2.position).unwrap();
661    partial.add(proof2.position, el2, &proof2.merkle_path).unwrap();
662
663    // check the number of nodes increased by a single element (the one that is not shared)
664    assert_eq!(partial.nodes.len(), 3);
665    // check the values match
666    let idx = InOrderIndex::from_leaf_pos(proof2.position);
667    assert_eq!(partial.nodes[&idx.sibling()], proof2.merkle_path[0]);
668    let idx = idx.parent();
669    assert_eq!(partial.nodes[&idx.sibling()], proof2.merkle_path[1]);
670}
671
672#[test]
673fn test_partial_mmr_update_single() {
674    let mut full = Mmr::new();
675    let zero = int_to_node(0);
676    full.add(zero);
677    let mut partial: PartialMmr = full.accumulator().into();
678
679    let proof = full.open(0).unwrap();
680    partial.add(proof.position, zero, &proof.merkle_path).unwrap();
681
682    for i in 1..100 {
683        let node = int_to_node(i);
684        full.add(node);
685        let delta = full.get_delta(partial.forest()).unwrap();
686        partial.apply(delta).unwrap();
687
688        assert_eq!(partial.forest(), full.forest());
689        assert_eq!(partial.peaks(), full.accumulator().peaks());
690
691        let proof1 = full.open(i as usize).unwrap();
692        partial.add(proof1.position, node, &proof1.merkle_path).unwrap();
693        let proof2 = partial.open(proof1.position).unwrap().unwrap();
694        assert_eq!(proof1.merkle_path, proof2.merkle_path);
695    }
696}
697
698#[test]
699fn test_mmr_add_invalid_odd_leaf() {
700    let mmr: Mmr = LEAVES.into();
701    let acc = mmr.accumulator();
702    let mut partial: PartialMmr = acc.clone().into();
703
704    let empty = MerklePath::new(Vec::new());
705
706    // None of the other leaves should work
707    for node in LEAVES.iter().cloned().rev().skip(1) {
708        let result = partial.add(LEAVES.len() - 1, node, &empty);
709        assert!(result.is_err());
710    }
711
712    let result = partial.add(LEAVES.len() - 1, LEAVES[6], &empty);
713    assert!(result.is_ok());
714}
715
716mod property_tests {
717    use super::leaf_to_corresponding_tree;
718    use proptest::prelude::*;
719
720    proptest! {
721        #[test]
722        fn test_last_position_is_always_contained_in_the_last_tree(leaves in any::<usize>().prop_filter("cant have an empty tree", |v| *v != 0)) {
723            let last_pos = leaves - 1;
724            let lowest_bit = leaves.trailing_zeros();
725
726            assert_eq!(
727                leaf_to_corresponding_tree(last_pos, leaves),
728                Some(lowest_bit),
729            );
730        }
731    }
732
733    proptest! {
734        #[test]
735        fn test_contained_tree_is_always_power_of_two((leaves, pos) in any::<usize>().prop_flat_map(|v| (Just(v), 0..v))) {
736            let tree_bit = leaf_to_corresponding_tree(pos, leaves).expect("pos is smaller than leaves, there should always be a corresponding tree");
737            let mask = 1usize << tree_bit;
738
739            assert!(tree_bit < usize::BITS, "the result must be a bit in usize");
740            assert!(mask & leaves != 0, "the result should be a tree in leaves");
741        }
742    }
743}
744
745// HELPER FUNCTIONS
746// ================================================================================================
747
748fn digests_to_elements(digests: &[RpoDigest]) -> Vec<Felt> {
749    digests.iter().flat_map(Word::from).collect()
750}
751
752// short hand for the rpo hash, used to make test code more concise and easy to read
753fn merge(l: RpoDigest, r: RpoDigest) -> RpoDigest {
754    Rpo256::merge(&[l, r])
755}