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