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.peaks(mmr.forest()).unwrap();
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.peaks(mmr.forest()).unwrap();
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.peaks(mmr.forest()).unwrap();
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.peaks(mmr.forest()).unwrap();
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.peaks(mmr.forest()).unwrap();
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.peaks(mmr.forest()).unwrap();
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.peaks(mmr.forest()).unwrap();
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!(
206        mmr.open(7, mmr.forest()).is_err(),
207        "Element 7 is not in the tree, result should be None"
208    );
209
210    // node at pos 6 is the root
211    let empty: MerklePath = MerklePath::new(vec![]);
212    let opening = mmr
213        .open(6, mmr.forest())
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    assert!(
219        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[6], opening),
220        "MmrProof should be valid for the current accumulator."
221    );
222
223    // nodes 4,5 are depth 1
224    let root_to_path = MerklePath::new(vec![LEAVES[4]]);
225    let opening = mmr
226        .open(5, mmr.forest())
227        .expect("Element 5 is contained in the tree, expected an opening result.");
228    assert_eq!(opening.merkle_path, root_to_path);
229    assert_eq!(opening.forest, mmr.forest);
230    assert_eq!(opening.position, 5);
231    assert!(
232        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[5], opening),
233        "MmrProof should be valid for the current accumulator."
234    );
235
236    let root_to_path = MerklePath::new(vec![LEAVES[5]]);
237    let opening = mmr
238        .open(4, mmr.forest())
239        .expect("Element 4 is contained in the tree, expected an opening result.");
240    assert_eq!(opening.merkle_path, root_to_path);
241    assert_eq!(opening.forest, mmr.forest);
242    assert_eq!(opening.position, 4);
243    assert!(
244        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[4], opening),
245        "MmrProof should be valid for the current accumulator."
246    );
247
248    // nodes 0,1,2,3 are detph 2
249    let root_to_path = MerklePath::new(vec![LEAVES[2], h01]);
250    let opening = mmr
251        .open(3, mmr.forest())
252        .expect("Element 3 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, 3);
256    assert!(
257        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[3], opening),
258        "MmrProof should be valid for the current accumulator."
259    );
260
261    let root_to_path = MerklePath::new(vec![LEAVES[3], h01]);
262    let opening = mmr
263        .open(2, mmr.forest())
264        .expect("Element 2 is contained in the tree, expected an opening result.");
265    assert_eq!(opening.merkle_path, root_to_path);
266    assert_eq!(opening.forest, mmr.forest);
267    assert_eq!(opening.position, 2);
268    assert!(
269        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[2], opening),
270        "MmrProof should be valid for the current accumulator."
271    );
272
273    let root_to_path = MerklePath::new(vec![LEAVES[0], h23]);
274    let opening = mmr
275        .open(1, mmr.forest())
276        .expect("Element 1 is contained in the tree, expected an opening result.");
277    assert_eq!(opening.merkle_path, root_to_path);
278    assert_eq!(opening.forest, mmr.forest);
279    assert_eq!(opening.position, 1);
280    assert!(
281        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[1], opening),
282        "MmrProof should be valid for the current accumulator."
283    );
284
285    let root_to_path = MerklePath::new(vec![LEAVES[1], h23]);
286    let opening = mmr
287        .open(0, mmr.forest())
288        .expect("Element 0 is contained in the tree, expected an opening result.");
289    assert_eq!(opening.merkle_path, root_to_path);
290    assert_eq!(opening.forest, mmr.forest);
291    assert_eq!(opening.position, 0);
292    assert!(
293        mmr.peaks(mmr.forest()).unwrap().verify(LEAVES[0], opening),
294        "MmrProof should be valid for the current accumulator."
295    );
296}
297
298#[test]
299fn test_mmr_open_older_version() {
300    let mmr: Mmr = LEAVES.into();
301
302    fn is_even(v: &usize) -> bool {
303        v & 1 == 0
304    }
305
306    // merkle path of a node is empty if there are no elements to pair with it
307    for pos in (0..mmr.forest()).filter(is_even) {
308        let forest = pos + 1;
309        let proof = mmr.open(pos, forest).unwrap();
310        assert_eq!(proof.forest, forest);
311        assert_eq!(proof.merkle_path.nodes(), []);
312        assert_eq!(proof.position, pos);
313    }
314
315    // openings match that of a merkle tree
316    let mtree: MerkleTree = LEAVES[..4].try_into().unwrap();
317    for forest in 4..=LEAVES.len() {
318        for pos in 0..4 {
319            let idx = NodeIndex::new(2, pos).unwrap();
320            let path = mtree.get_path(idx).unwrap();
321            let proof = mmr.open(pos as usize, forest).unwrap();
322            assert_eq!(path, proof.merkle_path);
323        }
324    }
325    let mtree: MerkleTree = LEAVES[4..6].try_into().unwrap();
326    for forest in 6..=LEAVES.len() {
327        for pos in 0..2 {
328            let idx = NodeIndex::new(1, pos).unwrap();
329            let path = mtree.get_path(idx).unwrap();
330            // account for the bigger tree with 4 elements
331            let mmr_pos = (pos + 4) as usize;
332            let proof = mmr.open(mmr_pos, forest).unwrap();
333            assert_eq!(path, proof.merkle_path);
334        }
335    }
336}
337
338/// Tests the openings of a simple Mmr with a single tree of depth 8.
339#[test]
340fn test_mmr_open_eight() {
341    let leaves = [
342        int_to_node(0),
343        int_to_node(1),
344        int_to_node(2),
345        int_to_node(3),
346        int_to_node(4),
347        int_to_node(5),
348        int_to_node(6),
349        int_to_node(7),
350    ];
351
352    let mtree: MerkleTree = leaves.as_slice().try_into().unwrap();
353    let forest = leaves.len();
354    let mmr: Mmr = leaves.into();
355    let root = mtree.root();
356
357    let position = 0;
358    let proof = mmr.open(position, mmr.forest()).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    let position = 1;
364    let proof = mmr.open(position, mmr.forest()).unwrap();
365    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
366    assert_eq!(proof, MmrProof { forest, position, merkle_path });
367    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
368
369    let position = 2;
370    let proof = mmr.open(position, mmr.forest()).unwrap();
371    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
372    assert_eq!(proof, MmrProof { forest, position, merkle_path });
373    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
374
375    let position = 3;
376    let proof = mmr.open(position, mmr.forest()).unwrap();
377    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
378    assert_eq!(proof, MmrProof { forest, position, merkle_path });
379    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
380
381    let position = 4;
382    let proof = mmr.open(position, mmr.forest()).unwrap();
383    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
384    assert_eq!(proof, MmrProof { forest, position, merkle_path });
385    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
386
387    let position = 5;
388    let proof = mmr.open(position, mmr.forest()).unwrap();
389    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
390    assert_eq!(proof, MmrProof { forest, position, merkle_path });
391    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
392
393    let position = 6;
394    let proof = mmr.open(position, mmr.forest()).unwrap();
395    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
396    assert_eq!(proof, MmrProof { forest, position, merkle_path });
397    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
398
399    let position = 7;
400    let proof = mmr.open(position, mmr.forest()).unwrap();
401    let merkle_path = mtree.get_path(NodeIndex::new(3, position as u64).unwrap()).unwrap();
402    assert_eq!(proof, MmrProof { forest, position, merkle_path });
403    assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root);
404}
405
406/// Tests the openings of Mmr with a 3 trees of depths 4, 2, and 1.
407#[test]
408fn test_mmr_open_seven() {
409    let mtree1: MerkleTree = LEAVES[..4].try_into().unwrap();
410    let mtree2: MerkleTree = LEAVES[4..6].try_into().unwrap();
411
412    let forest = LEAVES.len();
413    let mmr: Mmr = LEAVES.into();
414
415    let position = 0;
416    let proof = mmr.open(position, mmr.forest()).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(0, LEAVES[0]).unwrap(), mtree1.root());
421
422    let position = 1;
423    let proof = mmr.open(position, mmr.forest()).unwrap();
424    let merkle_path: MerklePath =
425        mtree1.get_path(NodeIndex::new(2, position as u64).unwrap()).unwrap();
426    assert_eq!(proof, MmrProof { forest, position, merkle_path });
427    assert_eq!(proof.merkle_path.compute_root(1, LEAVES[1]).unwrap(), mtree1.root());
428
429    let position = 2;
430    let proof = mmr.open(position, mmr.forest()).unwrap();
431    let merkle_path: MerklePath =
432        mtree1.get_path(NodeIndex::new(2, position as u64).unwrap()).unwrap();
433    assert_eq!(proof, MmrProof { forest, position, merkle_path });
434    assert_eq!(proof.merkle_path.compute_root(2, LEAVES[2]).unwrap(), mtree1.root());
435
436    let position = 3;
437    let proof = mmr.open(position, mmr.forest()).unwrap();
438    let merkle_path: MerklePath =
439        mtree1.get_path(NodeIndex::new(2, position as u64).unwrap()).unwrap();
440    assert_eq!(proof, MmrProof { forest, position, merkle_path });
441    assert_eq!(proof.merkle_path.compute_root(3, LEAVES[3]).unwrap(), mtree1.root());
442
443    let position = 4;
444    let proof = mmr.open(position, mmr.forest()).unwrap();
445    let merkle_path: MerklePath = mtree2.get_path(NodeIndex::new(1, 0u64).unwrap()).unwrap();
446    assert_eq!(proof, MmrProof { forest, position, merkle_path });
447    assert_eq!(proof.merkle_path.compute_root(0, LEAVES[4]).unwrap(), mtree2.root());
448
449    let position = 5;
450    let proof = mmr.open(position, mmr.forest()).unwrap();
451    let merkle_path: MerklePath = mtree2.get_path(NodeIndex::new(1, 1u64).unwrap()).unwrap();
452    assert_eq!(proof, MmrProof { forest, position, merkle_path });
453    assert_eq!(proof.merkle_path.compute_root(1, LEAVES[5]).unwrap(), mtree2.root());
454
455    let position = 6;
456    let proof = mmr.open(position, mmr.forest()).unwrap();
457    let merkle_path: MerklePath = [].as_ref().into();
458    assert_eq!(proof, MmrProof { forest, position, merkle_path });
459    assert_eq!(proof.merkle_path.compute_root(0, LEAVES[6]).unwrap(), LEAVES[6]);
460}
461
462#[test]
463fn test_mmr_get() {
464    let mmr: Mmr = LEAVES.into();
465    assert_eq!(mmr.get(0).unwrap(), LEAVES[0], "value at pos 0 must correspond");
466    assert_eq!(mmr.get(1).unwrap(), LEAVES[1], "value at pos 1 must correspond");
467    assert_eq!(mmr.get(2).unwrap(), LEAVES[2], "value at pos 2 must correspond");
468    assert_eq!(mmr.get(3).unwrap(), LEAVES[3], "value at pos 3 must correspond");
469    assert_eq!(mmr.get(4).unwrap(), LEAVES[4], "value at pos 4 must correspond");
470    assert_eq!(mmr.get(5).unwrap(), LEAVES[5], "value at pos 5 must correspond");
471    assert_eq!(mmr.get(6).unwrap(), LEAVES[6], "value at pos 6 must correspond");
472    assert!(mmr.get(7).is_err());
473}
474
475#[test]
476fn test_mmr_invariants() {
477    let mut mmr = Mmr::new();
478    for v in 1..=1028 {
479        mmr.add(int_to_node(v));
480        let accumulator = mmr.peaks(mmr.forest()).unwrap();
481        assert_eq!(v as usize, mmr.forest(), "MMR leaf count must increase by one on every add");
482        assert_eq!(
483            v as usize,
484            accumulator.num_leaves(),
485            "MMR and its accumulator must match leaves count"
486        );
487        assert_eq!(
488            accumulator.num_leaves().count_ones() as usize,
489            accumulator.peaks().len(),
490            "bits on leaves must match the number of peaks"
491        );
492
493        let expected_nodes: usize = TrueBitPositionIterator::new(mmr.forest())
494            .map(|bit_pos| nodes_in_forest(1 << bit_pos))
495            .sum();
496
497        assert_eq!(
498            expected_nodes,
499            mmr.nodes.len(),
500            "the sum of every tree size must be equal to the number of nodes in the MMR (forest: {:b})",
501            mmr.forest(),
502        );
503    }
504}
505
506#[test]
507fn test_bit_position_iterator() {
508    assert_eq!(TrueBitPositionIterator::new(0).count(), 0);
509    assert_eq!(TrueBitPositionIterator::new(0).rev().count(), 0);
510
511    assert_eq!(TrueBitPositionIterator::new(1).collect::<Vec<u32>>(), vec![0]);
512    assert_eq!(TrueBitPositionIterator::new(1).rev().collect::<Vec<u32>>(), vec![0],);
513
514    assert_eq!(TrueBitPositionIterator::new(2).collect::<Vec<u32>>(), vec![1]);
515    assert_eq!(TrueBitPositionIterator::new(2).rev().collect::<Vec<u32>>(), vec![1],);
516
517    assert_eq!(TrueBitPositionIterator::new(3).collect::<Vec<u32>>(), vec![0, 1],);
518    assert_eq!(TrueBitPositionIterator::new(3).rev().collect::<Vec<u32>>(), vec![1, 0],);
519
520    assert_eq!(
521        TrueBitPositionIterator::new(0b11010101).collect::<Vec<u32>>(),
522        vec![0, 2, 4, 6, 7],
523    );
524    assert_eq!(
525        TrueBitPositionIterator::new(0b11010101).rev().collect::<Vec<u32>>(),
526        vec![7, 6, 4, 2, 0],
527    );
528}
529
530#[test]
531fn test_mmr_inner_nodes() {
532    let mmr: Mmr = LEAVES.into();
533    let nodes: Vec<InnerNodeInfo> = mmr.inner_nodes().collect();
534
535    let h01 = Rpo256::merge(&[LEAVES[0], LEAVES[1]]);
536    let h23 = Rpo256::merge(&[LEAVES[2], LEAVES[3]]);
537    let h0123 = Rpo256::merge(&[h01, h23]);
538    let h45 = Rpo256::merge(&[LEAVES[4], LEAVES[5]]);
539    let postorder = vec![
540        InnerNodeInfo {
541            value: h01,
542            left: LEAVES[0],
543            right: LEAVES[1],
544        },
545        InnerNodeInfo {
546            value: h23,
547            left: LEAVES[2],
548            right: LEAVES[3],
549        },
550        InnerNodeInfo { value: h0123, left: h01, right: h23 },
551        InnerNodeInfo {
552            value: h45,
553            left: LEAVES[4],
554            right: LEAVES[5],
555        },
556    ];
557
558    assert_eq!(postorder, nodes);
559}
560
561#[test]
562fn test_mmr_peaks() {
563    let mmr: Mmr = LEAVES.into();
564
565    let forest = 0b0001;
566    let acc = mmr.peaks(forest).unwrap();
567    assert_eq!(acc.num_leaves(), forest);
568    assert_eq!(acc.peaks(), &[mmr.nodes[0]]);
569
570    let forest = 0b0010;
571    let acc = mmr.peaks(forest).unwrap();
572    assert_eq!(acc.num_leaves(), forest);
573    assert_eq!(acc.peaks(), &[mmr.nodes[2]]);
574
575    let forest = 0b0011;
576    let acc = mmr.peaks(forest).unwrap();
577    assert_eq!(acc.num_leaves(), forest);
578    assert_eq!(acc.peaks(), &[mmr.nodes[2], mmr.nodes[3]]);
579
580    let forest = 0b0100;
581    let acc = mmr.peaks(forest).unwrap();
582    assert_eq!(acc.num_leaves(), forest);
583    assert_eq!(acc.peaks(), &[mmr.nodes[6]]);
584
585    let forest = 0b0101;
586    let acc = mmr.peaks(forest).unwrap();
587    assert_eq!(acc.num_leaves(), forest);
588    assert_eq!(acc.peaks(), &[mmr.nodes[6], mmr.nodes[7]]);
589
590    let forest = 0b0110;
591    let acc = mmr.peaks(forest).unwrap();
592    assert_eq!(acc.num_leaves(), forest);
593    assert_eq!(acc.peaks(), &[mmr.nodes[6], mmr.nodes[9]]);
594
595    let forest = 0b0111;
596    let acc = mmr.peaks(forest).unwrap();
597    assert_eq!(acc.num_leaves(), forest);
598    assert_eq!(acc.peaks(), &[mmr.nodes[6], mmr.nodes[9], mmr.nodes[10]]);
599}
600
601#[test]
602fn test_mmr_hash_peaks() {
603    let mmr: Mmr = LEAVES.into();
604    let peaks = mmr.peaks(mmr.forest()).unwrap();
605
606    let first_peak = Rpo256::merge(&[
607        Rpo256::merge(&[LEAVES[0], LEAVES[1]]),
608        Rpo256::merge(&[LEAVES[2], LEAVES[3]]),
609    ]);
610    let second_peak = Rpo256::merge(&[LEAVES[4], LEAVES[5]]);
611    let third_peak = LEAVES[6];
612
613    // minimum length is 16
614    let mut expected_peaks = [first_peak, second_peak, third_peak].to_vec();
615    expected_peaks.resize(16, RpoDigest::default());
616    assert_eq!(
617        peaks.hash_peaks(),
618        *Rpo256::hash_elements(&digests_to_elements(&expected_peaks))
619    );
620}
621
622#[test]
623fn test_mmr_peaks_hash_less_than_16() {
624    let mut peaks = Vec::new();
625
626    for i in 0..16 {
627        peaks.push(int_to_node(i));
628
629        let num_leaves = (1 << peaks.len()) - 1;
630        let accumulator = MmrPeaks::new(num_leaves, peaks.clone()).unwrap();
631
632        // minimum length is 16
633        let mut expected_peaks = peaks.clone();
634        expected_peaks.resize(16, RpoDigest::default());
635        assert_eq!(
636            accumulator.hash_peaks(),
637            *Rpo256::hash_elements(&digests_to_elements(&expected_peaks))
638        );
639    }
640}
641
642#[test]
643fn test_mmr_peaks_hash_odd() {
644    let peaks: Vec<_> = (0..=17).map(int_to_node).collect();
645
646    let num_leaves = (1 << peaks.len()) - 1;
647    let accumulator = MmrPeaks::new(num_leaves, peaks.clone()).unwrap();
648
649    // odd length bigger than 16 is padded to the next even number
650    let mut expected_peaks = peaks;
651    expected_peaks.resize(18, RpoDigest::default());
652    assert_eq!(
653        accumulator.hash_peaks(),
654        *Rpo256::hash_elements(&digests_to_elements(&expected_peaks))
655    );
656}
657
658#[test]
659fn test_mmr_delta() {
660    let mmr: Mmr = LEAVES.into();
661    let acc = mmr.peaks(mmr.forest()).unwrap();
662
663    // original_forest can't have more elements
664    assert!(
665        mmr.get_delta(LEAVES.len() + 1, mmr.forest()).is_err(),
666        "Can not provide updates for a newer Mmr"
667    );
668
669    // if the number of elements is the same there is no change
670    assert!(
671        mmr.get_delta(LEAVES.len(), mmr.forest()).unwrap().data.is_empty(),
672        "There are no updates for the same Mmr version"
673    );
674
675    // missing the last element added, which is itself a tree peak
676    assert_eq!(mmr.get_delta(6, mmr.forest()).unwrap().data, vec![acc.peaks()[2]], "one peak");
677
678    // missing the sibling to complete the tree of depth 2, and the last element
679    assert_eq!(
680        mmr.get_delta(5, mmr.forest()).unwrap().data,
681        vec![LEAVES[5], acc.peaks()[2]],
682        "one sibling, one peak"
683    );
684
685    // missing the whole last two trees, only send the peaks
686    assert_eq!(
687        mmr.get_delta(4, mmr.forest()).unwrap().data,
688        vec![acc.peaks()[1], acc.peaks()[2]],
689        "two peaks"
690    );
691
692    // missing the sibling to complete the first tree, and the two last trees
693    assert_eq!(
694        mmr.get_delta(3, mmr.forest()).unwrap().data,
695        vec![LEAVES[3], acc.peaks()[1], acc.peaks()[2]],
696        "one sibling, two peaks"
697    );
698
699    // missing half of the first tree, only send the computed element (not the leaves), and the new
700    // peaks
701    assert_eq!(
702        mmr.get_delta(2, mmr.forest()).unwrap().data,
703        vec![mmr.nodes[5], acc.peaks()[1], acc.peaks()[2]],
704        "one sibling, two peaks"
705    );
706
707    assert_eq!(
708        mmr.get_delta(1, mmr.forest()).unwrap().data,
709        vec![LEAVES[1], mmr.nodes[5], acc.peaks()[1], acc.peaks()[2]],
710        "one sibling, two peaks"
711    );
712
713    assert_eq!(&mmr.get_delta(0, mmr.forest()).unwrap().data, acc.peaks(), "all peaks");
714}
715
716#[test]
717fn test_mmr_delta_old_forest() {
718    let mmr: Mmr = LEAVES.into();
719
720    // from_forest must be smaller-or-equal to to_forest
721    for version in 1..=mmr.forest() {
722        assert!(mmr.get_delta(version + 1, version).is_err());
723    }
724
725    // when from_forest and to_forest are equal, there are no updates
726    for version in 1..=mmr.forest() {
727        let delta = mmr.get_delta(version, version).unwrap();
728        assert!(delta.data.is_empty());
729        assert_eq!(delta.forest, version);
730    }
731
732    // test update which merges the odd peak to the right
733    for count in 0..(mmr.forest() / 2) {
734        // *2 because every iteration tests a pair
735        // +1 because the Mmr is 1-indexed
736        let from_forest = (count * 2) + 1;
737        let to_forest = (count * 2) + 2;
738        let delta = mmr.get_delta(from_forest, to_forest).unwrap();
739
740        // *2 because every iteration tests a pair
741        // +1 because sibling is the odd element
742        let sibling = (count * 2) + 1;
743        assert_eq!(delta.data, [LEAVES[sibling]]);
744        assert_eq!(delta.forest, to_forest);
745    }
746
747    let version = 4;
748    let delta = mmr.get_delta(1, version).unwrap();
749    assert_eq!(delta.data, [mmr.nodes[1], mmr.nodes[5]]);
750    assert_eq!(delta.forest, version);
751
752    let version = 5;
753    let delta = mmr.get_delta(1, version).unwrap();
754    assert_eq!(delta.data, [mmr.nodes[1], mmr.nodes[5], mmr.nodes[7]]);
755    assert_eq!(delta.forest, version);
756}
757
758#[test]
759fn test_partial_mmr_simple() {
760    let mmr: Mmr = LEAVES.into();
761    let acc = mmr.peaks(mmr.forest()).unwrap();
762    let mut partial: PartialMmr = acc.clone().into();
763
764    // check initial state of the partial mmr
765    assert_eq!(partial.peaks(), acc.peaks());
766    assert_eq!(partial.forest(), acc.num_leaves());
767    assert_eq!(partial.forest(), LEAVES.len());
768    assert_eq!(partial.peaks().len(), 3);
769    assert_eq!(partial.nodes.len(), 0);
770
771    // check state after adding tracking one element
772    let proof1 = mmr.open(0, mmr.forest()).unwrap();
773    let el1 = mmr.get(proof1.position).unwrap();
774    partial.add(proof1.position, el1, &proof1.merkle_path).unwrap();
775
776    // check the number of nodes increased by the number of nodes in the proof
777    assert_eq!(partial.nodes.len(), proof1.merkle_path.len());
778    // check the values match
779    let idx = InOrderIndex::from_leaf_pos(proof1.position);
780    assert_eq!(partial.nodes[&idx.sibling()], proof1.merkle_path[0]);
781    let idx = idx.parent();
782    assert_eq!(partial.nodes[&idx.sibling()], proof1.merkle_path[1]);
783
784    let proof2 = mmr.open(1, mmr.forest()).unwrap();
785    let el2 = mmr.get(proof2.position).unwrap();
786    partial.add(proof2.position, el2, &proof2.merkle_path).unwrap();
787
788    // check the number of nodes increased by a single element (the one that is not shared)
789    assert_eq!(partial.nodes.len(), 3);
790    // check the values match
791    let idx = InOrderIndex::from_leaf_pos(proof2.position);
792    assert_eq!(partial.nodes[&idx.sibling()], proof2.merkle_path[0]);
793    let idx = idx.parent();
794    assert_eq!(partial.nodes[&idx.sibling()], proof2.merkle_path[1]);
795}
796
797#[test]
798fn test_partial_mmr_update_single() {
799    let mut full = Mmr::new();
800    let zero = int_to_node(0);
801    full.add(zero);
802    let mut partial: PartialMmr = full.peaks(full.forest()).unwrap().into();
803
804    let proof = full.open(0, full.forest()).unwrap();
805    partial.add(proof.position, zero, &proof.merkle_path).unwrap();
806
807    for i in 1..100 {
808        let node = int_to_node(i);
809        full.add(node);
810        let delta = full.get_delta(partial.forest(), full.forest()).unwrap();
811        partial.apply(delta).unwrap();
812
813        assert_eq!(partial.forest(), full.forest());
814        assert_eq!(partial.peaks(), full.peaks(full.forest()).unwrap().peaks());
815
816        let proof1 = full.open(i as usize, full.forest()).unwrap();
817        partial.add(proof1.position, node, &proof1.merkle_path).unwrap();
818        let proof2 = partial.open(proof1.position).unwrap().unwrap();
819        assert_eq!(proof1.merkle_path, proof2.merkle_path);
820    }
821}
822
823#[test]
824fn test_mmr_add_invalid_odd_leaf() {
825    let mmr: Mmr = LEAVES.into();
826    let acc = mmr.peaks(mmr.forest()).unwrap();
827    let mut partial: PartialMmr = acc.clone().into();
828
829    let empty = MerklePath::new(Vec::new());
830
831    // None of the other leaves should work
832    for node in LEAVES.iter().cloned().rev().skip(1) {
833        let result = partial.add(LEAVES.len() - 1, node, &empty);
834        assert!(result.is_err());
835    }
836
837    let result = partial.add(LEAVES.len() - 1, LEAVES[6], &empty);
838    assert!(result.is_ok());
839}
840
841mod property_tests {
842    use super::leaf_to_corresponding_tree;
843    use proptest::prelude::*;
844
845    proptest! {
846        #[test]
847        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)) {
848            let last_pos = leaves - 1;
849            let lowest_bit = leaves.trailing_zeros();
850
851            assert_eq!(
852                leaf_to_corresponding_tree(last_pos, leaves),
853                Some(lowest_bit),
854            );
855        }
856    }
857
858    proptest! {
859        #[test]
860        fn test_contained_tree_is_always_power_of_two((leaves, pos) in any::<usize>().prop_flat_map(|v| (Just(v), 0..v))) {
861            let tree_bit = leaf_to_corresponding_tree(pos, leaves).expect("pos is smaller than leaves, there should always be a corresponding tree");
862            let mask = 1usize << tree_bit;
863
864            assert!(tree_bit < usize::BITS, "the result must be a bit in usize");
865            assert!(mask & leaves != 0, "the result should be a tree in leaves");
866        }
867    }
868}
869
870// HELPER FUNCTIONS
871// ================================================================================================
872
873fn digests_to_elements(digests: &[RpoDigest]) -> Vec<Felt> {
874    digests.iter().flat_map(Word::from).collect()
875}
876
877// short hand for the rpo hash, used to make test code more concise and easy to read
878fn merge(l: RpoDigest, r: RpoDigest) -> RpoDigest {
879    Rpo256::merge(&[l, r])
880}