arnaucube
arnaucube/miden-cryptoPublic Self-hosted repository
miden-crypto/src/merkle/mmr/tests.rs
src/merkle/mmr/tests.rsRaw
  1use super::bit::TrueBitPositionIterator;
  2use super::full::{high_bitmask, leaf_to_corresponding_tree, nodes_in_forest};
  3use super::{super::Vec, Mmr, Rpo256, Word};
  4use crate::merkle::{int_to_node, MerklePath};
  5
  6#[test]
  7fn test_position_equal_or_higher_than_leafs_is_never_contained() {
  8    let empty_forest = 0;
  9    for pos in 1..1024 {
 10        // pos is index, 0 based
 11        // tree is a length counter, 1 based
 12        // so a valid pos is always smaller, not equal, to tree
 13        assert_eq!(leaf_to_corresponding_tree(pos, pos), None);
 14        assert_eq!(leaf_to_corresponding_tree(pos, pos - 1), None);
 15        // and empty forest has no trees, so no position is valid
 16        assert_eq!(leaf_to_corresponding_tree(pos, empty_forest), None);
 17    }
 18}
 19
 20#[test]
 21fn test_position_zero_is_always_contained_by_the_highest_tree() {
 22    for leaves in 1..1024usize {
 23        let tree = leaves.ilog2();
 24        assert_eq!(leaf_to_corresponding_tree(0, leaves), Some(tree));
 25    }
 26}
 27
 28#[test]
 29fn test_leaf_to_corresponding_tree() {
 30    assert_eq!(leaf_to_corresponding_tree(0, 0b0001), Some(0));
 31    assert_eq!(leaf_to_corresponding_tree(0, 0b0010), Some(1));
 32    assert_eq!(leaf_to_corresponding_tree(0, 0b0011), Some(1));
 33    assert_eq!(leaf_to_corresponding_tree(0, 0b1011), Some(3));
 34
 35    // position one is always owned by the left-most tree
 36    assert_eq!(leaf_to_corresponding_tree(1, 0b0010), Some(1));
 37    assert_eq!(leaf_to_corresponding_tree(1, 0b0011), Some(1));
 38    assert_eq!(leaf_to_corresponding_tree(1, 0b1011), Some(3));
 39
 40    // position two starts as its own root, and then it is merged with the left-most tree
 41    assert_eq!(leaf_to_corresponding_tree(2, 0b0011), Some(0));
 42    assert_eq!(leaf_to_corresponding_tree(2, 0b0100), Some(2));
 43    assert_eq!(leaf_to_corresponding_tree(2, 0b1011), Some(3));
 44
 45    // position tree is merged on the left-most tree
 46    assert_eq!(leaf_to_corresponding_tree(3, 0b0011), None);
 47    assert_eq!(leaf_to_corresponding_tree(3, 0b0100), Some(2));
 48    assert_eq!(leaf_to_corresponding_tree(3, 0b1011), Some(3));
 49
 50    assert_eq!(leaf_to_corresponding_tree(4, 0b0101), Some(0));
 51    assert_eq!(leaf_to_corresponding_tree(4, 0b0110), Some(1));
 52    assert_eq!(leaf_to_corresponding_tree(4, 0b0111), Some(1));
 53    assert_eq!(leaf_to_corresponding_tree(4, 0b1000), Some(3));
 54
 55    assert_eq!(leaf_to_corresponding_tree(12, 0b01101), Some(0));
 56    assert_eq!(leaf_to_corresponding_tree(12, 0b01110), Some(1));
 57    assert_eq!(leaf_to_corresponding_tree(12, 0b01111), Some(1));
 58    assert_eq!(leaf_to_corresponding_tree(12, 0b10000), Some(4));
 59}
 60
 61#[test]
 62fn test_high_bitmask() {
 63    assert_eq!(high_bitmask(0), usize::MAX);
 64    assert_eq!(high_bitmask(1), usize::MAX << 1);
 65    assert_eq!(high_bitmask(usize::BITS - 2), 0b11usize.rotate_right(2));
 66    assert_eq!(high_bitmask(usize::BITS - 1), 0b1usize.rotate_right(1));
 67    assert_eq!(high_bitmask(usize::BITS), 0, "overflow should be handled");
 68}
 69
 70#[test]
 71fn test_nodes_in_forest() {
 72    assert_eq!(nodes_in_forest(0b0000), 0);
 73    assert_eq!(nodes_in_forest(0b0001), 1);
 74    assert_eq!(nodes_in_forest(0b0010), 3);
 75    assert_eq!(nodes_in_forest(0b0011), 4);
 76    assert_eq!(nodes_in_forest(0b0100), 7);
 77    assert_eq!(nodes_in_forest(0b0101), 8);
 78    assert_eq!(nodes_in_forest(0b0110), 10);
 79    assert_eq!(nodes_in_forest(0b0111), 11);
 80    assert_eq!(nodes_in_forest(0b1000), 15);
 81    assert_eq!(nodes_in_forest(0b1001), 16);
 82    assert_eq!(nodes_in_forest(0b1010), 18);
 83    assert_eq!(nodes_in_forest(0b1011), 19);
 84}
 85
 86#[test]
 87fn test_nodes_in_forest_single_bit() {
 88    assert_eq!(nodes_in_forest(2usize.pow(0)), 2usize.pow(1) - 1);
 89    assert_eq!(nodes_in_forest(2usize.pow(1)), 2usize.pow(2) - 1);
 90    assert_eq!(nodes_in_forest(2usize.pow(2)), 2usize.pow(3) - 1);
 91    assert_eq!(nodes_in_forest(2usize.pow(3)), 2usize.pow(4) - 1);
 92
 93    for bit in 0..(usize::BITS - 1) {
 94        let size = 2usize.pow(bit + 1) - 1;
 95        assert_eq!(nodes_in_forest(1usize << bit), size);
 96    }
 97}
 98
 99const LEAVES: [Word; 7] = [
100    int_to_node(0),
101    int_to_node(1),
102    int_to_node(2),
103    int_to_node(3),
104    int_to_node(4),
105    int_to_node(5),
106    int_to_node(6),
107];
108
109#[test]
110fn test_mmr_simple() {
111    let mut postorder = Vec::new();
112    postorder.push(LEAVES[0]);
113    postorder.push(LEAVES[1]);
114    postorder.push(*Rpo256::hash_elements(&[LEAVES[0], LEAVES[1]].concat()));
115    postorder.push(LEAVES[2]);
116    postorder.push(LEAVES[3]);
117    postorder.push(*Rpo256::hash_elements(&[LEAVES[2], LEAVES[3]].concat()));
118    postorder.push(*Rpo256::hash_elements(
119        &[postorder[2], postorder[5]].concat(),
120    ));
121    postorder.push(LEAVES[4]);
122    postorder.push(LEAVES[5]);
123    postorder.push(*Rpo256::hash_elements(&[LEAVES[4], LEAVES[5]].concat()));
124    postorder.push(LEAVES[6]);
125
126    let mut mmr = Mmr::new();
127    assert_eq!(mmr.forest(), 0);
128    assert_eq!(mmr.nodes.len(), 0);
129
130    mmr.add(LEAVES[0]);
131    assert_eq!(mmr.forest(), 1);
132    assert_eq!(mmr.nodes.len(), 1);
133    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
134
135    let acc = mmr.accumulator();
136    assert_eq!(acc.num_leaves, 1);
137    assert_eq!(acc.peaks, &[postorder[0]]);
138
139    mmr.add(LEAVES[1]);
140    assert_eq!(mmr.forest(), 2);
141    assert_eq!(mmr.nodes.len(), 3);
142    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
143
144    let acc = mmr.accumulator();
145    assert_eq!(acc.num_leaves, 2);
146    assert_eq!(acc.peaks, &[postorder[2]]);
147
148    mmr.add(LEAVES[2]);
149    assert_eq!(mmr.forest(), 3);
150    assert_eq!(mmr.nodes.len(), 4);
151    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
152
153    let acc = mmr.accumulator();
154    assert_eq!(acc.num_leaves, 3);
155    assert_eq!(acc.peaks, &[postorder[2], postorder[3]]);
156
157    mmr.add(LEAVES[3]);
158    assert_eq!(mmr.forest(), 4);
159    assert_eq!(mmr.nodes.len(), 7);
160    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
161
162    let acc = mmr.accumulator();
163    assert_eq!(acc.num_leaves, 4);
164    assert_eq!(acc.peaks, &[postorder[6]]);
165
166    mmr.add(LEAVES[4]);
167    assert_eq!(mmr.forest(), 5);
168    assert_eq!(mmr.nodes.len(), 8);
169    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
170
171    let acc = mmr.accumulator();
172    assert_eq!(acc.num_leaves, 5);
173    assert_eq!(acc.peaks, &[postorder[6], postorder[7]]);
174
175    mmr.add(LEAVES[5]);
176    assert_eq!(mmr.forest(), 6);
177    assert_eq!(mmr.nodes.len(), 10);
178    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
179
180    let acc = mmr.accumulator();
181    assert_eq!(acc.num_leaves, 6);
182    assert_eq!(acc.peaks, &[postorder[6], postorder[9]]);
183
184    mmr.add(LEAVES[6]);
185    assert_eq!(mmr.forest(), 7);
186    assert_eq!(mmr.nodes.len(), 11);
187    assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
188
189    let acc = mmr.accumulator();
190    assert_eq!(acc.num_leaves, 7);
191    assert_eq!(acc.peaks, &[postorder[6], postorder[9], postorder[10]]);
192}
193
194#[test]
195fn test_mmr_open() {
196    let mmr: Mmr = LEAVES.into();
197    let h01: Word = Rpo256::hash_elements(&LEAVES[0..2].concat()).into();
198    let h23: Word = Rpo256::hash_elements(&LEAVES[2..4].concat()).into();
199
200    // node at pos 7 is the root
201    assert!(
202        mmr.open(7).is_err(),
203        "Element 7 is not in the tree, result should be None"
204    );
205
206    // node at pos 6 is the root
207    let empty: MerklePath = MerklePath::new(vec![]);
208    let opening = mmr
209        .open(6)
210        .expect("Element 6 is contained in the tree, expected an opening result.");
211    assert_eq!(opening.merkle_path, empty);
212    assert_eq!(opening.forest, mmr.forest);
213    assert_eq!(opening.position, 6);
214    assert!(
215        mmr.accumulator().verify(LEAVES[6], opening),
216        "MmrProof should be valid for the current accumulator."
217    );
218
219    // nodes 4,5 are detph 1
220    let root_to_path = MerklePath::new(vec![LEAVES[4]]);
221    let opening = mmr
222        .open(5)
223        .expect("Element 5 is contained in the tree, expected an opening result.");
224    assert_eq!(opening.merkle_path, root_to_path);
225    assert_eq!(opening.forest, mmr.forest);
226    assert_eq!(opening.position, 5);
227    assert!(
228        mmr.accumulator().verify(LEAVES[5], opening),
229        "MmrProof should be valid for the current accumulator."
230    );
231
232    let root_to_path = MerklePath::new(vec![LEAVES[5]]);
233    let opening = mmr
234        .open(4)
235        .expect("Element 4 is contained in the tree, expected an opening result.");
236    assert_eq!(opening.merkle_path, root_to_path);
237    assert_eq!(opening.forest, mmr.forest);
238    assert_eq!(opening.position, 4);
239    assert!(
240        mmr.accumulator().verify(LEAVES[4], opening),
241        "MmrProof should be valid for the current accumulator."
242    );
243
244    // nodes 0,1,2,3 are detph 2
245    let root_to_path = MerklePath::new(vec![LEAVES[2], h01]);
246    let opening = mmr
247        .open(3)
248        .expect("Element 3 is contained in the tree, expected an opening result.");
249    assert_eq!(opening.merkle_path, root_to_path);
250    assert_eq!(opening.forest, mmr.forest);
251    assert_eq!(opening.position, 3);
252    assert!(
253        mmr.accumulator().verify(LEAVES[3], opening),
254        "MmrProof should be valid for the current accumulator."
255    );
256
257    let root_to_path = MerklePath::new(vec![LEAVES[3], h01]);
258    let opening = mmr
259        .open(2)
260        .expect("Element 2 is contained in the tree, expected an opening result.");
261    assert_eq!(opening.merkle_path, root_to_path);
262    assert_eq!(opening.forest, mmr.forest);
263    assert_eq!(opening.position, 2);
264    assert!(
265        mmr.accumulator().verify(LEAVES[2], opening),
266        "MmrProof should be valid for the current accumulator."
267    );
268
269    let root_to_path = MerklePath::new(vec![LEAVES[0], h23]);
270    let opening = mmr
271        .open(1)
272        .expect("Element 1 is contained in the tree, expected an opening result.");
273    assert_eq!(opening.merkle_path, root_to_path);
274    assert_eq!(opening.forest, mmr.forest);
275    assert_eq!(opening.position, 1);
276    assert!(
277        mmr.accumulator().verify(LEAVES[1], opening),
278        "MmrProof should be valid for the current accumulator."
279    );
280
281    let root_to_path = MerklePath::new(vec![LEAVES[1], h23]);
282    let opening = mmr
283        .open(0)
284        .expect("Element 0 is contained in the tree, expected an opening result.");
285    assert_eq!(opening.merkle_path, root_to_path);
286    assert_eq!(opening.forest, mmr.forest);
287    assert_eq!(opening.position, 0);
288    assert!(
289        mmr.accumulator().verify(LEAVES[0], opening),
290        "MmrProof should be valid for the current accumulator."
291    );
292}
293
294#[test]
295fn test_mmr_get() {
296    let mmr: Mmr = LEAVES.into();
297    assert_eq!(
298        mmr.get(0).unwrap(),
299        LEAVES[0],
300        "value at pos 0 must correspond"
301    );
302    assert_eq!(
303        mmr.get(1).unwrap(),
304        LEAVES[1],
305        "value at pos 1 must correspond"
306    );
307    assert_eq!(
308        mmr.get(2).unwrap(),
309        LEAVES[2],
310        "value at pos 2 must correspond"
311    );
312    assert_eq!(
313        mmr.get(3).unwrap(),
314        LEAVES[3],
315        "value at pos 3 must correspond"
316    );
317    assert_eq!(
318        mmr.get(4).unwrap(),
319        LEAVES[4],
320        "value at pos 4 must correspond"
321    );
322    assert_eq!(
323        mmr.get(5).unwrap(),
324        LEAVES[5],
325        "value at pos 5 must correspond"
326    );
327    assert_eq!(
328        mmr.get(6).unwrap(),
329        LEAVES[6],
330        "value at pos 6 must correspond"
331    );
332    assert!(mmr.get(7).is_err());
333}
334
335#[test]
336fn test_mmr_invariants() {
337    let mut mmr = Mmr::new();
338    for v in 1..=1028 {
339        mmr.add(int_to_node(v));
340        let accumulator = mmr.accumulator();
341        assert_eq!(
342            v as usize,
343            mmr.forest(),
344            "MMR leaf count must increase by one on every add"
345        );
346        assert_eq!(
347            v as usize, accumulator.num_leaves,
348            "MMR and its accumulator must match leaves count"
349        );
350        assert_eq!(
351            accumulator.num_leaves.count_ones() as usize,
352            accumulator.peaks.len(),
353            "bits on leaves must match the number of peaks"
354        );
355
356        let expected_nodes: usize = TrueBitPositionIterator::new(mmr.forest())
357            .map(|bit_pos| nodes_in_forest(1 << bit_pos))
358            .sum();
359
360        assert_eq!(
361            expected_nodes,
362            mmr.nodes.len(),
363            "the sum of every tree size must be equal to the number of nodes in the MMR (forest: {:b})",
364            mmr.forest(),
365        );
366    }
367}
368
369#[test]
370fn test_bit_position_iterator() {
371    assert_eq!(TrueBitPositionIterator::new(0).count(), 0);
372    assert_eq!(TrueBitPositionIterator::new(0).rev().count(), 0);
373
374    assert_eq!(
375        TrueBitPositionIterator::new(1).collect::<Vec<u32>>(),
376        vec![0]
377    );
378    assert_eq!(
379        TrueBitPositionIterator::new(1).rev().collect::<Vec<u32>>(),
380        vec![0],
381    );
382
383    assert_eq!(
384        TrueBitPositionIterator::new(2).collect::<Vec<u32>>(),
385        vec![1]
386    );
387    assert_eq!(
388        TrueBitPositionIterator::new(2).rev().collect::<Vec<u32>>(),
389        vec![1],
390    );
391
392    assert_eq!(
393        TrueBitPositionIterator::new(3).collect::<Vec<u32>>(),
394        vec![0, 1],
395    );
396    assert_eq!(
397        TrueBitPositionIterator::new(3).rev().collect::<Vec<u32>>(),
398        vec![1, 0],
399    );
400
401    assert_eq!(
402        TrueBitPositionIterator::new(0b11010101).collect::<Vec<u32>>(),
403        vec![0, 2, 4, 6, 7],
404    );
405    assert_eq!(
406        TrueBitPositionIterator::new(0b11010101)
407            .rev()
408            .collect::<Vec<u32>>(),
409        vec![7, 6, 4, 2, 0],
410    );
411}
412
413mod property_tests {
414    use super::leaf_to_corresponding_tree;
415    use proptest::prelude::*;
416
417    proptest! {
418        #[test]
419        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)) {
420            let last_pos = leaves - 1;
421            let lowest_bit = leaves.trailing_zeros();
422
423            assert_eq!(
424                leaf_to_corresponding_tree(last_pos, leaves),
425                Some(lowest_bit),
426            );
427        }
428    }
429
430    proptest! {
431        #[test]
432        fn test_contained_tree_is_always_power_of_two((leaves, pos) in any::<usize>().prop_flat_map(|v| (Just(v), 0..v))) {
433            let tree = leaf_to_corresponding_tree(pos, leaves).expect("pos is smaller than leaves, there should always be a corresponding tree");
434            let mask = 1usize << tree;
435
436            assert!(tree < usize::BITS, "the result must be a bit in usize");
437            assert!(mask & leaves != 0, "the result should be a tree in leaves");
438        }
439    }
440}