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