miden-crypto/src/merkle/smt/full/tests.rs
src/merkle/smt/full/tests.rsRaw
1use alloc::vec::Vec;
2
3use super::{Felt, LeafIndex, NodeIndex, Rpo256, RpoDigest, Smt, SmtLeaf, EMPTY_WORD, SMT_DEPTH};
4use crate::{
5 merkle::{smt::SparseMerkleTree, EmptySubtreeRoots, MerkleStore},
6 utils::{Deserializable, Serializable},
7 Word, ONE, WORD_SIZE,
8};
9
10// SMT
11// --------------------------------------------------------------------------------------------
12
13/// This test checks that inserting twice at the same key functions as expected. The test covers
14/// only the case where the key is alone in its leaf
15#[test]
16fn test_smt_insert_at_same_key() {
17 let mut smt = Smt::default();
18 let mut store: MerkleStore = MerkleStore::default();
19
20 assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
21
22 let key_1: RpoDigest = {
23 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
24
25 RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
26 };
27 let key_1_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key_1).into();
28
29 let value_1 = [ONE; WORD_SIZE];
30 let value_2 = [ONE + ONE; WORD_SIZE];
31
32 // Insert value 1 and ensure root is as expected
33 {
34 let leaf_node = build_empty_or_single_leaf_node(key_1, value_1);
35 let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
36
37 let old_value_1 = smt.insert(key_1, value_1);
38 assert_eq!(old_value_1, EMPTY_WORD);
39
40 assert_eq!(smt.root(), tree_root);
41 }
42
43 // Insert value 2 and ensure root is as expected
44 {
45 let leaf_node = build_empty_or_single_leaf_node(key_1, value_2);
46 let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
47
48 let old_value_2 = smt.insert(key_1, value_2);
49 assert_eq!(old_value_2, value_1);
50
51 assert_eq!(smt.root(), tree_root);
52 }
53}
54
55/// This test checks that inserting twice at the same key functions as expected. The test covers
56/// only the case where the leaf type is `SmtLeaf::Multiple`
57#[test]
58fn test_smt_insert_at_same_key_2() {
59 // The most significant u64 used for both keys (to ensure they map to the same leaf)
60 let key_msb: u64 = 42;
61
62 let key_already_present: RpoDigest =
63 RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(key_msb)]);
64 let key_already_present_index: NodeIndex =
65 LeafIndex::<SMT_DEPTH>::from(key_already_present).into();
66 let value_already_present = [ONE + ONE + ONE; WORD_SIZE];
67
68 let mut smt =
69 Smt::with_entries(core::iter::once((key_already_present, value_already_present))).unwrap();
70 let mut store: MerkleStore = {
71 let mut store = MerkleStore::default();
72
73 let leaf_node = build_empty_or_single_leaf_node(key_already_present, value_already_present);
74 store
75 .set_node(*EmptySubtreeRoots::entry(SMT_DEPTH, 0), key_already_present_index, leaf_node)
76 .unwrap();
77 store
78 };
79
80 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(key_msb)]);
81 let key_1_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key_1).into();
82
83 assert_eq!(key_1_index, key_already_present_index);
84
85 let value_1 = [ONE; WORD_SIZE];
86 let value_2 = [ONE + ONE; WORD_SIZE];
87
88 // Insert value 1 and ensure root is as expected
89 {
90 // Note: key_1 comes first because it is smaller
91 let leaf_node = build_multiple_leaf_node(&[
92 (key_1, value_1),
93 (key_already_present, value_already_present),
94 ]);
95 let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
96
97 let old_value_1 = smt.insert(key_1, value_1);
98 assert_eq!(old_value_1, EMPTY_WORD);
99
100 assert_eq!(smt.root(), tree_root);
101 }
102
103 // Insert value 2 and ensure root is as expected
104 {
105 let leaf_node = build_multiple_leaf_node(&[
106 (key_1, value_2),
107 (key_already_present, value_already_present),
108 ]);
109 let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
110
111 let old_value_2 = smt.insert(key_1, value_2);
112 assert_eq!(old_value_2, value_1);
113
114 assert_eq!(smt.root(), tree_root);
115 }
116}
117
118/// This test ensures that the root of the tree is as expected when we add/remove 3 items at 3
119/// different keys. This also tests that the merkle paths produced are as expected.
120#[test]
121fn test_smt_insert_and_remove_multiple_values() {
122 fn insert_values_and_assert_path(
123 smt: &mut Smt,
124 store: &mut MerkleStore,
125 key_values: &[(RpoDigest, Word)],
126 ) {
127 for &(key, value) in key_values {
128 let key_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key).into();
129
130 let leaf_node = build_empty_or_single_leaf_node(key, value);
131 let tree_root = store.set_node(smt.root(), key_index, leaf_node).unwrap().root;
132
133 let _ = smt.insert(key, value);
134
135 assert_eq!(smt.root(), tree_root);
136
137 let expected_path = store.get_path(tree_root, key_index).unwrap();
138 assert_eq!(smt.open(&key).into_parts().0, expected_path.path);
139 }
140 }
141 let mut smt = Smt::default();
142 let mut store: MerkleStore = MerkleStore::default();
143
144 assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
145
146 let key_1: RpoDigest = {
147 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
148
149 RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
150 };
151
152 let key_2: RpoDigest = {
153 let raw = 0b_11111111_11111111_11111111_11111111_11111111_11111111_11111111_11111111_u64;
154
155 RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
156 };
157
158 let key_3: RpoDigest = {
159 let raw = 0b_00000000_00000000_00000000_00000000_00000000_00000000_00000000_00000000_u64;
160
161 RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
162 };
163
164 let value_1 = [ONE; WORD_SIZE];
165 let value_2 = [ONE + ONE; WORD_SIZE];
166 let value_3 = [ONE + ONE + ONE; WORD_SIZE];
167
168 // Insert values in the tree
169 let key_values = [(key_1, value_1), (key_2, value_2), (key_3, value_3)];
170 insert_values_and_assert_path(&mut smt, &mut store, &key_values);
171
172 // Remove values from the tree
173 let key_empty_values = [(key_1, EMPTY_WORD), (key_2, EMPTY_WORD), (key_3, EMPTY_WORD)];
174 insert_values_and_assert_path(&mut smt, &mut store, &key_empty_values);
175
176 let empty_root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
177 assert_eq!(smt.root(), empty_root);
178
179 // an empty tree should have no leaves or inner nodes
180 assert!(smt.leaves.is_empty());
181 assert!(smt.inner_nodes.is_empty());
182}
183
184/// This tests that inserting the empty value does indeed remove the key-value contained at the
185/// leaf. We insert & remove 3 values at the same leaf to ensure that all cases are covered (empty,
186/// single, multiple).
187#[test]
188fn test_smt_removal() {
189 let mut smt = Smt::default();
190
191 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
192
193 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
194 let key_2: RpoDigest =
195 RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
196 let key_3: RpoDigest =
197 RpoDigest::from([3_u32.into(), 3_u32.into(), 3_u32.into(), Felt::new(raw)]);
198
199 let value_1 = [ONE; WORD_SIZE];
200 let value_2 = [2_u32.into(); WORD_SIZE];
201 let value_3: [Felt; 4] = [3_u32.into(); WORD_SIZE];
202
203 // insert key-value 1
204 {
205 let old_value_1 = smt.insert(key_1, value_1);
206 assert_eq!(old_value_1, EMPTY_WORD);
207
208 assert_eq!(smt.get_leaf(&key_1), SmtLeaf::Single((key_1, value_1)));
209 }
210
211 // insert key-value 2
212 {
213 let old_value_2 = smt.insert(key_2, value_2);
214 assert_eq!(old_value_2, EMPTY_WORD);
215
216 assert_eq!(
217 smt.get_leaf(&key_2),
218 SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)])
219 );
220 }
221
222 // insert key-value 3
223 {
224 let old_value_3 = smt.insert(key_3, value_3);
225 assert_eq!(old_value_3, EMPTY_WORD);
226
227 assert_eq!(
228 smt.get_leaf(&key_3),
229 SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2), (key_3, value_3)])
230 );
231 }
232
233 // remove key 3
234 {
235 let old_value_3 = smt.insert(key_3, EMPTY_WORD);
236 assert_eq!(old_value_3, value_3);
237
238 assert_eq!(
239 smt.get_leaf(&key_3),
240 SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)])
241 );
242 }
243
244 // remove key 2
245 {
246 let old_value_2 = smt.insert(key_2, EMPTY_WORD);
247 assert_eq!(old_value_2, value_2);
248
249 assert_eq!(smt.get_leaf(&key_2), SmtLeaf::Single((key_1, value_1)));
250 }
251
252 // remove key 1
253 {
254 let old_value_1 = smt.insert(key_1, EMPTY_WORD);
255 assert_eq!(old_value_1, value_1);
256
257 assert_eq!(smt.get_leaf(&key_1), SmtLeaf::new_empty(key_1.into()));
258 }
259}
260
261/// This tests that we can correctly calculate prospective leaves -- that is, we can construct
262/// correct [`SmtLeaf`] values for a theoretical insertion on a Merkle tree without mutating or
263/// cloning the tree.
264#[test]
265fn test_prospective_hash() {
266 let mut smt = Smt::default();
267
268 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
269
270 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
271 let key_2: RpoDigest =
272 RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
273 // Sort key_3 before key_1, to test non-append insertion.
274 let key_3: RpoDigest =
275 RpoDigest::from([0_u32.into(), 0_u32.into(), 0_u32.into(), Felt::new(raw)]);
276
277 let value_1 = [ONE; WORD_SIZE];
278 let value_2 = [2_u32.into(); WORD_SIZE];
279 let value_3: [Felt; 4] = [3_u32.into(); WORD_SIZE];
280
281 // insert key-value 1
282 {
283 let prospective =
284 smt.construct_prospective_leaf(smt.get_leaf(&key_1), &key_1, &value_1).hash();
285 smt.insert(key_1, value_1);
286
287 let leaf = smt.get_leaf(&key_1);
288 assert_eq!(
289 prospective,
290 leaf.hash(),
291 "prospective hash for leaf {leaf:?} did not match actual hash",
292 );
293 }
294
295 // insert key-value 2
296 {
297 let prospective =
298 smt.construct_prospective_leaf(smt.get_leaf(&key_2), &key_2, &value_2).hash();
299 smt.insert(key_2, value_2);
300
301 let leaf = smt.get_leaf(&key_2);
302 assert_eq!(
303 prospective,
304 leaf.hash(),
305 "prospective hash for leaf {leaf:?} did not match actual hash",
306 );
307 }
308
309 // insert key-value 3
310 {
311 let prospective =
312 smt.construct_prospective_leaf(smt.get_leaf(&key_3), &key_3, &value_3).hash();
313 smt.insert(key_3, value_3);
314
315 let leaf = smt.get_leaf(&key_3);
316 assert_eq!(
317 prospective,
318 leaf.hash(),
319 "prospective hash for leaf {leaf:?} did not match actual hash",
320 );
321 }
322
323 // remove key 3
324 {
325 let old_leaf = smt.get_leaf(&key_3);
326 let old_value_3 = smt.insert(key_3, EMPTY_WORD);
327 assert_eq!(old_value_3, value_3);
328 let prospective_leaf =
329 smt.construct_prospective_leaf(smt.get_leaf(&key_3), &key_3, &old_value_3);
330
331 assert_eq!(
332 old_leaf.hash(),
333 prospective_leaf.hash(),
334 "removing and prospectively re-adding a leaf didn't yield the original leaf:\
335 \n original leaf: {old_leaf:?}\
336 \n prospective leaf: {prospective_leaf:?}",
337 );
338 }
339
340 // remove key 2
341 {
342 let old_leaf = smt.get_leaf(&key_2);
343 let old_value_2 = smt.insert(key_2, EMPTY_WORD);
344 assert_eq!(old_value_2, value_2);
345 let prospective_leaf =
346 smt.construct_prospective_leaf(smt.get_leaf(&key_2), &key_2, &old_value_2);
347
348 assert_eq!(
349 old_leaf.hash(),
350 prospective_leaf.hash(),
351 "removing and prospectively re-adding a leaf didn't yield the original leaf:\
352 \n original leaf: {old_leaf:?}\
353 \n prospective leaf: {prospective_leaf:?}",
354 );
355 }
356
357 // remove key 1
358 {
359 let old_leaf = smt.get_leaf(&key_1);
360 let old_value_1 = smt.insert(key_1, EMPTY_WORD);
361 assert_eq!(old_value_1, value_1);
362 let prospective_leaf =
363 smt.construct_prospective_leaf(smt.get_leaf(&key_1), &key_1, &old_value_1);
364 assert_eq!(
365 old_leaf.hash(),
366 prospective_leaf.hash(),
367 "removing and prospectively re-adding a leaf didn't yield the original leaf:\
368 \n original leaf: {old_leaf:?}\
369 \n prospective leaf: {prospective_leaf:?}",
370 );
371 }
372}
373
374/// This tests that we can perform prospective changes correctly.
375#[test]
376fn test_prospective_insertion() {
377 let mut smt = Smt::default();
378
379 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
380
381 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
382 let key_2: RpoDigest =
383 RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
384 // Sort key_3 before key_1, to test non-append insertion.
385 let key_3: RpoDigest =
386 RpoDigest::from([0_u32.into(), 0_u32.into(), 0_u32.into(), Felt::new(raw)]);
387
388 let value_1 = [ONE; WORD_SIZE];
389 let value_2 = [2_u32.into(); WORD_SIZE];
390 let value_3: [Felt; 4] = [3_u32.into(); WORD_SIZE];
391
392 let root_empty = smt.root();
393
394 let root_1 = {
395 smt.insert(key_1, value_1);
396 smt.root()
397 };
398
399 let root_2 = {
400 smt.insert(key_2, value_2);
401 smt.root()
402 };
403
404 let root_3 = {
405 smt.insert(key_3, value_3);
406 smt.root()
407 };
408
409 // Test incremental updates.
410
411 let mut smt = Smt::default();
412
413 let mutations = smt.compute_mutations(vec![(key_1, value_1)]);
414 assert_eq!(mutations.root(), root_1, "prospective root 1 did not match actual root 1");
415 smt.apply_mutations(mutations).unwrap();
416 assert_eq!(smt.root(), root_1, "mutations before and after apply did not match");
417
418 let mutations = smt.compute_mutations(vec![(key_2, value_2)]);
419 assert_eq!(mutations.root(), root_2, "prospective root 2 did not match actual root 2");
420 let mutations =
421 smt.compute_mutations(vec![(key_3, EMPTY_WORD), (key_2, value_2), (key_3, value_3)]);
422 assert_eq!(mutations.root(), root_3, "mutations before and after apply did not match");
423 smt.apply_mutations(mutations).unwrap();
424
425 // Edge case: multiple values at the same key, where a later pair restores the original value.
426 let mutations = smt.compute_mutations(vec![(key_3, EMPTY_WORD), (key_3, value_3)]);
427 assert_eq!(mutations.root(), root_3);
428 smt.apply_mutations(mutations).unwrap();
429 assert_eq!(smt.root(), root_3);
430
431 // Test batch updates, and that the order doesn't matter.
432 let pairs =
433 vec![(key_3, value_2), (key_2, EMPTY_WORD), (key_1, EMPTY_WORD), (key_3, EMPTY_WORD)];
434 let mutations = smt.compute_mutations(pairs);
435 assert_eq!(
436 mutations.root(),
437 root_empty,
438 "prospective root for batch removal did not match actual root",
439 );
440 smt.apply_mutations(mutations).unwrap();
441 assert_eq!(smt.root(), root_empty, "mutations before and after apply did not match");
442
443 let pairs = vec![(key_3, value_3), (key_1, value_1), (key_2, value_2)];
444 let mutations = smt.compute_mutations(pairs);
445 assert_eq!(mutations.root(), root_3);
446 smt.apply_mutations(mutations).unwrap();
447 assert_eq!(smt.root(), root_3);
448}
449
450/// Tests that 2 key-value pairs stored in the same leaf have the same path
451#[test]
452fn test_smt_path_to_keys_in_same_leaf_are_equal() {
453 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
454
455 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
456 let key_2: RpoDigest =
457 RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
458
459 let value_1 = [ONE; WORD_SIZE];
460 let value_2 = [2_u32.into(); WORD_SIZE];
461
462 let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
463
464 assert_eq!(smt.open(&key_1), smt.open(&key_2));
465}
466
467/// Tests that an empty leaf hashes to the empty word
468#[test]
469fn test_empty_leaf_hash() {
470 let smt = Smt::default();
471
472 let leaf = smt.get_leaf(&RpoDigest::default());
473 assert_eq!(leaf.hash(), EMPTY_WORD.into());
474}
475
476/// Tests that `get_value()` works as expected
477#[test]
478fn test_smt_get_value() {
479 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, ONE]);
480 let key_2: RpoDigest = RpoDigest::from([2_u32, 2_u32, 2_u32, 2_u32]);
481
482 let value_1 = [ONE; WORD_SIZE];
483 let value_2 = [2_u32.into(); WORD_SIZE];
484
485 let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
486
487 let returned_value_1 = smt.get_value(&key_1);
488 let returned_value_2 = smt.get_value(&key_2);
489
490 assert_eq!(value_1, returned_value_1);
491 assert_eq!(value_2, returned_value_2);
492
493 // Check that a key with no inserted value returns the empty word
494 let key_no_value = RpoDigest::from([42_u32, 42_u32, 42_u32, 42_u32]);
495
496 assert_eq!(EMPTY_WORD, smt.get_value(&key_no_value));
497}
498
499/// Tests that `entries()` works as expected
500#[test]
501fn test_smt_entries() {
502 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, ONE]);
503 let key_2: RpoDigest = RpoDigest::from([2_u32, 2_u32, 2_u32, 2_u32]);
504
505 let value_1 = [ONE; WORD_SIZE];
506 let value_2 = [2_u32.into(); WORD_SIZE];
507
508 let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
509
510 let mut entries = smt.entries();
511
512 // Note: for simplicity, we assume the order `(k1,v1), (k2,v2)`. If a new implementation
513 // switches the order, it is OK to modify the order here as well.
514 assert_eq!(&(key_1, value_1), entries.next().unwrap());
515 assert_eq!(&(key_2, value_2), entries.next().unwrap());
516 assert!(entries.next().is_none());
517}
518
519/// Tests that `EMPTY_ROOT` constant generated in the `Smt` equals to the root of the empty tree of
520/// depth 64
521#[test]
522fn test_smt_check_empty_root_constant() {
523 // get the root of the empty tree of depth 64
524 let empty_root_64_depth = EmptySubtreeRoots::empty_hashes(64)[0];
525
526 assert_eq!(empty_root_64_depth, Smt::EMPTY_ROOT);
527}
528
529// SMT LEAF
530// --------------------------------------------------------------------------------------------
531
532#[test]
533fn test_empty_smt_leaf_serialization() {
534 let empty_leaf = SmtLeaf::new_empty(LeafIndex::new_max_depth(42));
535
536 let mut serialized = empty_leaf.to_bytes();
537 // extend buffer with random bytes
538 serialized.extend([1, 2, 3, 4, 5]);
539 let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
540
541 assert_eq!(empty_leaf, deserialized);
542}
543
544#[test]
545fn test_single_smt_leaf_serialization() {
546 let single_leaf = SmtLeaf::new_single(
547 RpoDigest::from([10_u32, 11_u32, 12_u32, 13_u32]),
548 [1_u32.into(), 2_u32.into(), 3_u32.into(), 4_u32.into()],
549 );
550
551 let mut serialized = single_leaf.to_bytes();
552 // extend buffer with random bytes
553 serialized.extend([1, 2, 3, 4, 5]);
554 let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
555
556 assert_eq!(single_leaf, deserialized);
557}
558
559#[test]
560fn test_multiple_smt_leaf_serialization_success() {
561 let multiple_leaf = SmtLeaf::new_multiple(vec![
562 (
563 RpoDigest::from([10_u32, 11_u32, 12_u32, 13_u32]),
564 [1_u32.into(), 2_u32.into(), 3_u32.into(), 4_u32.into()],
565 ),
566 (
567 RpoDigest::from([100_u32, 101_u32, 102_u32, 13_u32]),
568 [11_u32.into(), 12_u32.into(), 13_u32.into(), 14_u32.into()],
569 ),
570 ])
571 .unwrap();
572
573 let mut serialized = multiple_leaf.to_bytes();
574 // extend buffer with random bytes
575 serialized.extend([1, 2, 3, 4, 5]);
576 let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
577
578 assert_eq!(multiple_leaf, deserialized);
579}
580
581// HELPERS
582// --------------------------------------------------------------------------------------------
583
584fn build_empty_or_single_leaf_node(key: RpoDigest, value: Word) -> RpoDigest {
585 if value == EMPTY_WORD {
586 SmtLeaf::new_empty(key.into()).hash()
587 } else {
588 SmtLeaf::Single((key, value)).hash()
589 }
590}
591
592fn build_multiple_leaf_node(kv_pairs: &[(RpoDigest, Word)]) -> RpoDigest {
593 let elements: Vec<Felt> = kv_pairs
594 .iter()
595 .flat_map(|(key, value)| {
596 let key_elements = key.into_iter();
597 let value_elements = (*value).into_iter();
598
599 key_elements.chain(value_elements)
600 })
601 .collect();
602
603 Rpo256::hash_elements(&elements)
604}