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