miden-crypto/src/merkle/mmr/tests.rs
src/merkle/mmr/tests.rsRaw
1use super::{
2 super::{InnerNodeInfo, Vec},
3 bit::TrueBitPositionIterator,
4 full::{high_bitmask, leaf_to_corresponding_tree, nodes_in_forest},
5 Mmr, MmrPeaks, Rpo256,
6};
7use crate::{
8 hash::rpo::RpoDigest,
9 merkle::{int_to_node, MerklePath},
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::new();
119 postorder.push(LEAVES[0]);
120 postorder.push(LEAVES[1]);
121 postorder.push(Rpo256::merge(&[LEAVES[0], LEAVES[1]]));
122 postorder.push(LEAVES[2]);
123 postorder.push(LEAVES[3]);
124 postorder.push(Rpo256::merge(&[LEAVES[2], LEAVES[3]]));
125 postorder.push(Rpo256::merge(&[postorder[2], postorder[5]]));
126 postorder.push(LEAVES[4]);
127 postorder.push(LEAVES[5]);
128 postorder.push(Rpo256::merge(&[LEAVES[4], LEAVES[5]]));
129 postorder.push(LEAVES[6]);
130
131 let mut mmr = Mmr::new();
132 assert_eq!(mmr.forest(), 0);
133 assert_eq!(mmr.nodes.len(), 0);
134
135 mmr.add(LEAVES[0]);
136 assert_eq!(mmr.forest(), 1);
137 assert_eq!(mmr.nodes.len(), 1);
138 assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
139
140 let acc = mmr.accumulator();
141 assert_eq!(acc.num_leaves, 1);
142 assert_eq!(acc.peaks, &[postorder[0]]);
143
144 mmr.add(LEAVES[1]);
145 assert_eq!(mmr.forest(), 2);
146 assert_eq!(mmr.nodes.len(), 3);
147 assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
148
149 let acc = mmr.accumulator();
150 assert_eq!(acc.num_leaves, 2);
151 assert_eq!(acc.peaks, &[postorder[2]]);
152
153 mmr.add(LEAVES[2]);
154 assert_eq!(mmr.forest(), 3);
155 assert_eq!(mmr.nodes.len(), 4);
156 assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
157
158 let acc = mmr.accumulator();
159 assert_eq!(acc.num_leaves, 3);
160 assert_eq!(acc.peaks, &[postorder[2], postorder[3]]);
161
162 mmr.add(LEAVES[3]);
163 assert_eq!(mmr.forest(), 4);
164 assert_eq!(mmr.nodes.len(), 7);
165 assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
166
167 let acc = mmr.accumulator();
168 assert_eq!(acc.num_leaves, 4);
169 assert_eq!(acc.peaks, &[postorder[6]]);
170
171 mmr.add(LEAVES[4]);
172 assert_eq!(mmr.forest(), 5);
173 assert_eq!(mmr.nodes.len(), 8);
174 assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
175
176 let acc = mmr.accumulator();
177 assert_eq!(acc.num_leaves, 5);
178 assert_eq!(acc.peaks, &[postorder[6], postorder[7]]);
179
180 mmr.add(LEAVES[5]);
181 assert_eq!(mmr.forest(), 6);
182 assert_eq!(mmr.nodes.len(), 10);
183 assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
184
185 let acc = mmr.accumulator();
186 assert_eq!(acc.num_leaves, 6);
187 assert_eq!(acc.peaks, &[postorder[6], postorder[9]]);
188
189 mmr.add(LEAVES[6]);
190 assert_eq!(mmr.forest(), 7);
191 assert_eq!(mmr.nodes.len(), 11);
192 assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]);
193
194 let acc = mmr.accumulator();
195 assert_eq!(acc.num_leaves, 7);
196 assert_eq!(acc.peaks, &[postorder[6], postorder[9], postorder[10]]);
197}
198
199#[test]
200fn test_mmr_open() {
201 let mmr: Mmr = LEAVES.into();
202 let h01 = Rpo256::merge(&[LEAVES[0], LEAVES[1]]);
203 let h23 = Rpo256::merge(&[LEAVES[2], LEAVES[3]]);
204
205 // node at pos 7 is the root
206 assert!(mmr.open(7).is_err(), "Element 7 is not in the tree, result should be None");
207
208 // node at pos 6 is the root
209 let empty: MerklePath = MerklePath::new(vec![]);
210 let opening = mmr
211 .open(6)
212 .expect("Element 6 is contained in the tree, expected an opening result.");
213 assert_eq!(opening.merkle_path, empty);
214 assert_eq!(opening.forest, mmr.forest);
215 assert_eq!(opening.position, 6);
216 assert!(
217 mmr.accumulator().verify(LEAVES[6], opening),
218 "MmrProof should be valid for the current accumulator."
219 );
220
221 // nodes 4,5 are depth 1
222 let root_to_path = MerklePath::new(vec![LEAVES[4]]);
223 let opening = mmr
224 .open(5)
225 .expect("Element 5 is contained in the tree, expected an opening result.");
226 assert_eq!(opening.merkle_path, root_to_path);
227 assert_eq!(opening.forest, mmr.forest);
228 assert_eq!(opening.position, 5);
229 assert!(
230 mmr.accumulator().verify(LEAVES[5], opening),
231 "MmrProof should be valid for the current accumulator."
232 );
233
234 let root_to_path = MerklePath::new(vec![LEAVES[5]]);
235 let opening = mmr
236 .open(4)
237 .expect("Element 4 is contained in the tree, expected an opening result.");
238 assert_eq!(opening.merkle_path, root_to_path);
239 assert_eq!(opening.forest, mmr.forest);
240 assert_eq!(opening.position, 4);
241 assert!(
242 mmr.accumulator().verify(LEAVES[4], opening),
243 "MmrProof should be valid for the current accumulator."
244 );
245
246 // nodes 0,1,2,3 are detph 2
247 let root_to_path = MerklePath::new(vec![LEAVES[2], h01]);
248 let opening = mmr
249 .open(3)
250 .expect("Element 3 is contained in the tree, expected an opening result.");
251 assert_eq!(opening.merkle_path, root_to_path);
252 assert_eq!(opening.forest, mmr.forest);
253 assert_eq!(opening.position, 3);
254 assert!(
255 mmr.accumulator().verify(LEAVES[3], opening),
256 "MmrProof should be valid for the current accumulator."
257 );
258
259 let root_to_path = MerklePath::new(vec![LEAVES[3], h01]);
260 let opening = mmr
261 .open(2)
262 .expect("Element 2 is contained in the tree, expected an opening result.");
263 assert_eq!(opening.merkle_path, root_to_path);
264 assert_eq!(opening.forest, mmr.forest);
265 assert_eq!(opening.position, 2);
266 assert!(
267 mmr.accumulator().verify(LEAVES[2], opening),
268 "MmrProof should be valid for the current accumulator."
269 );
270
271 let root_to_path = MerklePath::new(vec![LEAVES[0], h23]);
272 let opening = mmr
273 .open(1)
274 .expect("Element 1 is contained in the tree, expected an opening result.");
275 assert_eq!(opening.merkle_path, root_to_path);
276 assert_eq!(opening.forest, mmr.forest);
277 assert_eq!(opening.position, 1);
278 assert!(
279 mmr.accumulator().verify(LEAVES[1], opening),
280 "MmrProof should be valid for the current accumulator."
281 );
282
283 let root_to_path = MerklePath::new(vec![LEAVES[1], h23]);
284 let opening = mmr
285 .open(0)
286 .expect("Element 0 is contained in the tree, expected an opening result.");
287 assert_eq!(opening.merkle_path, root_to_path);
288 assert_eq!(opening.forest, mmr.forest);
289 assert_eq!(opening.position, 0);
290 assert!(
291 mmr.accumulator().verify(LEAVES[0], opening),
292 "MmrProof should be valid for the current accumulator."
293 );
294}
295
296#[test]
297fn test_mmr_get() {
298 let mmr: Mmr = LEAVES.into();
299 assert_eq!(mmr.get(0).unwrap(), LEAVES[0], "value at pos 0 must correspond");
300 assert_eq!(mmr.get(1).unwrap(), LEAVES[1], "value at pos 1 must correspond");
301 assert_eq!(mmr.get(2).unwrap(), LEAVES[2], "value at pos 2 must correspond");
302 assert_eq!(mmr.get(3).unwrap(), LEAVES[3], "value at pos 3 must correspond");
303 assert_eq!(mmr.get(4).unwrap(), LEAVES[4], "value at pos 4 must correspond");
304 assert_eq!(mmr.get(5).unwrap(), LEAVES[5], "value at pos 5 must correspond");
305 assert_eq!(mmr.get(6).unwrap(), LEAVES[6], "value at pos 6 must correspond");
306 assert!(mmr.get(7).is_err());
307}
308
309#[test]
310fn test_mmr_invariants() {
311 let mut mmr = Mmr::new();
312 for v in 1..=1028 {
313 mmr.add(int_to_node(v));
314 let accumulator = mmr.accumulator();
315 assert_eq!(v as usize, mmr.forest(), "MMR leaf count must increase by one on every add");
316 assert_eq!(
317 v as usize, accumulator.num_leaves,
318 "MMR and its accumulator must match leaves count"
319 );
320 assert_eq!(
321 accumulator.num_leaves.count_ones() as usize,
322 accumulator.peaks.len(),
323 "bits on leaves must match the number of peaks"
324 );
325
326 let expected_nodes: usize = TrueBitPositionIterator::new(mmr.forest())
327 .map(|bit_pos| nodes_in_forest(1 << bit_pos))
328 .sum();
329
330 assert_eq!(
331 expected_nodes,
332 mmr.nodes.len(),
333 "the sum of every tree size must be equal to the number of nodes in the MMR (forest: {:b})",
334 mmr.forest(),
335 );
336 }
337}
338
339#[test]
340fn test_bit_position_iterator() {
341 assert_eq!(TrueBitPositionIterator::new(0).count(), 0);
342 assert_eq!(TrueBitPositionIterator::new(0).rev().count(), 0);
343
344 assert_eq!(TrueBitPositionIterator::new(1).collect::<Vec<u32>>(), vec![0]);
345 assert_eq!(TrueBitPositionIterator::new(1).rev().collect::<Vec<u32>>(), vec![0],);
346
347 assert_eq!(TrueBitPositionIterator::new(2).collect::<Vec<u32>>(), vec![1]);
348 assert_eq!(TrueBitPositionIterator::new(2).rev().collect::<Vec<u32>>(), vec![1],);
349
350 assert_eq!(TrueBitPositionIterator::new(3).collect::<Vec<u32>>(), vec![0, 1],);
351 assert_eq!(TrueBitPositionIterator::new(3).rev().collect::<Vec<u32>>(), vec![1, 0],);
352
353 assert_eq!(
354 TrueBitPositionIterator::new(0b11010101).collect::<Vec<u32>>(),
355 vec![0, 2, 4, 6, 7],
356 );
357 assert_eq!(
358 TrueBitPositionIterator::new(0b11010101).rev().collect::<Vec<u32>>(),
359 vec![7, 6, 4, 2, 0],
360 );
361}
362
363#[test]
364fn test_mmr_inner_nodes() {
365 let mmr: Mmr = LEAVES.into();
366 let nodes: Vec<InnerNodeInfo> = mmr.inner_nodes().collect();
367
368 let h01 = Rpo256::merge(&[LEAVES[0], LEAVES[1]]);
369 let h23 = Rpo256::merge(&[LEAVES[2], LEAVES[3]]);
370 let h0123 = Rpo256::merge(&[h01, h23]);
371 let h45 = Rpo256::merge(&[LEAVES[4], LEAVES[5]]);
372 let postorder = vec![
373 InnerNodeInfo {
374 value: h01,
375 left: LEAVES[0],
376 right: LEAVES[1],
377 },
378 InnerNodeInfo {
379 value: h23,
380 left: LEAVES[2],
381 right: LEAVES[3],
382 },
383 InnerNodeInfo { value: h0123, left: h01, right: h23 },
384 InnerNodeInfo {
385 value: h45,
386 left: LEAVES[4],
387 right: LEAVES[5],
388 },
389 ];
390
391 assert_eq!(postorder, nodes);
392}
393
394#[test]
395fn test_mmr_hash_peaks() {
396 let mmr: Mmr = LEAVES.into();
397 let peaks = mmr.accumulator();
398
399 let first_peak = Rpo256::merge(&[
400 Rpo256::merge(&[LEAVES[0], LEAVES[1]]),
401 Rpo256::merge(&[LEAVES[2], LEAVES[3]]),
402 ]);
403 let second_peak = Rpo256::merge(&[LEAVES[4], LEAVES[5]]);
404 let third_peak = LEAVES[6];
405
406 // minimum length is 16
407 let mut expected_peaks = [first_peak, second_peak, third_peak].to_vec();
408 expected_peaks.resize(16, RpoDigest::default());
409 assert_eq!(
410 peaks.hash_peaks(),
411 *Rpo256::hash_elements(&digests_to_elements(&expected_peaks))
412 );
413}
414
415#[test]
416fn test_mmr_peaks_hash_less_than_16() {
417 let mut peaks = Vec::new();
418
419 for i in 0..16 {
420 peaks.push(int_to_node(i));
421 let accumulator = MmrPeaks {
422 num_leaves: (1 << peaks.len()) - 1,
423 peaks: peaks.clone(),
424 };
425
426 // minimum length is 16
427 let mut expected_peaks = peaks.clone();
428 expected_peaks.resize(16, RpoDigest::default());
429 assert_eq!(
430 accumulator.hash_peaks(),
431 *Rpo256::hash_elements(&digests_to_elements(&expected_peaks))
432 );
433 }
434}
435
436#[test]
437fn test_mmr_peaks_hash_odd() {
438 let peaks: Vec<_> = (0..=17).map(int_to_node).collect();
439
440 let accumulator = MmrPeaks {
441 num_leaves: (1 << peaks.len()) - 1,
442 peaks: peaks.clone(),
443 };
444
445 // odd length bigger than 16 is padded to the next even number
446 let mut expected_peaks = peaks;
447 expected_peaks.resize(18, RpoDigest::default());
448 assert_eq!(
449 accumulator.hash_peaks(),
450 *Rpo256::hash_elements(&digests_to_elements(&expected_peaks))
451 );
452}
453
454mod property_tests {
455 use super::leaf_to_corresponding_tree;
456 use proptest::prelude::*;
457
458 proptest! {
459 #[test]
460 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)) {
461 let last_pos = leaves - 1;
462 let lowest_bit = leaves.trailing_zeros();
463
464 assert_eq!(
465 leaf_to_corresponding_tree(last_pos, leaves),
466 Some(lowest_bit),
467 );
468 }
469 }
470
471 proptest! {
472 #[test]
473 fn test_contained_tree_is_always_power_of_two((leaves, pos) in any::<usize>().prop_flat_map(|v| (Just(v), 0..v))) {
474 let tree = leaf_to_corresponding_tree(pos, leaves).expect("pos is smaller than leaves, there should always be a corresponding tree");
475 let mask = 1usize << tree;
476
477 assert!(tree < usize::BITS, "the result must be a bit in usize");
478 assert!(mask & leaves != 0, "the result should be a tree in leaves");
479 }
480 }
481}
482
483// HELPER FUNCTIONS
484// ================================================================================================
485
486fn digests_to_elements(digests: &[RpoDigest]) -> Vec<Felt> {
487 digests.iter().flat_map(Word::from).collect()
488}