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