miden-crypto/src/merkle/smt/full/tests.rs
src/merkle/smt/full/tests.rsRaw
1use super::{Felt, LeafIndex, NodeIndex, Rpo256, RpoDigest, Smt, SmtLeaf, EMPTY_WORD, SMT_DEPTH};
2use crate::{
3 merkle::{EmptySubtreeRoots, MerkleStore},
4 utils::{Deserializable, Serializable},
5 Word, ONE, WORD_SIZE,
6};
7use alloc::vec::Vec;
8
9// SMT
10// --------------------------------------------------------------------------------------------
11
12/// This test checks that inserting twice at the same key functions as expected. The test covers
13/// only the case where the key is alone in its leaf
14#[test]
15fn test_smt_insert_at_same_key() {
16 let mut smt = Smt::default();
17 let mut store: MerkleStore = MerkleStore::default();
18
19 assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
20
21 let key_1: RpoDigest = {
22 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
23
24 RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
25 };
26 let key_1_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key_1).into();
27
28 let value_1 = [ONE; WORD_SIZE];
29 let value_2 = [ONE + ONE; WORD_SIZE];
30
31 // Insert value 1 and ensure root is as expected
32 {
33 let leaf_node = build_empty_or_single_leaf_node(key_1, value_1);
34 let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
35
36 let old_value_1 = smt.insert(key_1, value_1);
37 assert_eq!(old_value_1, EMPTY_WORD);
38
39 assert_eq!(smt.root(), tree_root);
40 }
41
42 // Insert value 2 and ensure root is as expected
43 {
44 let leaf_node = build_empty_or_single_leaf_node(key_1, value_2);
45 let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
46
47 let old_value_2 = smt.insert(key_1, value_2);
48 assert_eq!(old_value_2, value_1);
49
50 assert_eq!(smt.root(), tree_root);
51 }
52}
53
54/// This test checks that inserting twice at the same key functions as expected. The test covers
55/// only the case where the leaf type is `SmtLeaf::Multiple`
56#[test]
57fn test_smt_insert_at_same_key_2() {
58 // The most significant u64 used for both keys (to ensure they map to the same leaf)
59 let key_msb: u64 = 42;
60
61 let key_already_present: RpoDigest =
62 RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(key_msb)]);
63 let key_already_present_index: NodeIndex =
64 LeafIndex::<SMT_DEPTH>::from(key_already_present).into();
65 let value_already_present = [ONE + ONE + ONE; WORD_SIZE];
66
67 let mut smt =
68 Smt::with_entries(core::iter::once((key_already_present, value_already_present))).unwrap();
69 let mut store: MerkleStore = {
70 let mut store = MerkleStore::default();
71
72 let leaf_node = build_empty_or_single_leaf_node(key_already_present, value_already_present);
73 store
74 .set_node(*EmptySubtreeRoots::entry(SMT_DEPTH, 0), key_already_present_index, leaf_node)
75 .unwrap();
76 store
77 };
78
79 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(key_msb)]);
80 let key_1_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key_1).into();
81
82 assert_eq!(key_1_index, key_already_present_index);
83
84 let value_1 = [ONE; WORD_SIZE];
85 let value_2 = [ONE + ONE; WORD_SIZE];
86
87 // Insert value 1 and ensure root is as expected
88 {
89 // Note: key_1 comes first because it is smaller
90 let leaf_node = build_multiple_leaf_node(&[
91 (key_1, value_1),
92 (key_already_present, value_already_present),
93 ]);
94 let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
95
96 let old_value_1 = smt.insert(key_1, value_1);
97 assert_eq!(old_value_1, EMPTY_WORD);
98
99 assert_eq!(smt.root(), tree_root);
100 }
101
102 // Insert value 2 and ensure root is as expected
103 {
104 let leaf_node = build_multiple_leaf_node(&[
105 (key_1, value_2),
106 (key_already_present, value_already_present),
107 ]);
108 let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
109
110 let old_value_2 = smt.insert(key_1, value_2);
111 assert_eq!(old_value_2, value_1);
112
113 assert_eq!(smt.root(), tree_root);
114 }
115}
116
117/// This test ensures that the root of the tree is as expected when we add/remove 3 items at 3
118/// different keys. This also tests that the merkle paths produced are as expected.
119#[test]
120fn test_smt_insert_and_remove_multiple_values() {
121 fn insert_values_and_assert_path(
122 smt: &mut Smt,
123 store: &mut MerkleStore,
124 key_values: &[(RpoDigest, Word)],
125 ) {
126 for &(key, value) in key_values {
127 let key_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key).into();
128
129 let leaf_node = build_empty_or_single_leaf_node(key, value);
130 let tree_root = store.set_node(smt.root(), key_index, leaf_node).unwrap().root;
131
132 let _ = smt.insert(key, value);
133
134 assert_eq!(smt.root(), tree_root);
135
136 let expected_path = store.get_path(tree_root, key_index).unwrap();
137 assert_eq!(smt.open(&key).into_parts().0, expected_path.path);
138 }
139 }
140 let mut smt = Smt::default();
141 let mut store: MerkleStore = MerkleStore::default();
142
143 assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
144
145 let key_1: RpoDigest = {
146 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
147
148 RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
149 };
150
151 let key_2: RpoDigest = {
152 let raw = 0b_11111111_11111111_11111111_11111111_11111111_11111111_11111111_11111111_u64;
153
154 RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
155 };
156
157 let key_3: RpoDigest = {
158 let raw = 0b_00000000_00000000_00000000_00000000_00000000_00000000_00000000_00000000_u64;
159
160 RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
161 };
162
163 let value_1 = [ONE; WORD_SIZE];
164 let value_2 = [ONE + ONE; WORD_SIZE];
165 let value_3 = [ONE + ONE + ONE; WORD_SIZE];
166
167 // Insert values in the tree
168 let key_values = [(key_1, value_1), (key_2, value_2), (key_3, value_3)];
169 insert_values_and_assert_path(&mut smt, &mut store, &key_values);
170
171 // Remove values from the tree
172 let key_empty_values = [(key_1, EMPTY_WORD), (key_2, EMPTY_WORD), (key_3, EMPTY_WORD)];
173 insert_values_and_assert_path(&mut smt, &mut store, &key_empty_values);
174
175 let empty_root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
176 assert_eq!(smt.root(), empty_root);
177
178 // an empty tree should have no leaves or inner nodes
179 assert!(smt.leaves.is_empty());
180 assert!(smt.inner_nodes.is_empty());
181}
182
183/// This tests that inserting the empty value does indeed remove the key-value contained at the
184/// leaf. We insert & remove 3 values at the same leaf to ensure that all cases are covered (empty,
185/// single, multiple).
186#[test]
187fn test_smt_removal() {
188 let mut smt = Smt::default();
189
190 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
191
192 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
193 let key_2: RpoDigest =
194 RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
195 let key_3: RpoDigest =
196 RpoDigest::from([3_u32.into(), 3_u32.into(), 3_u32.into(), Felt::new(raw)]);
197
198 let value_1 = [ONE; WORD_SIZE];
199 let value_2 = [2_u32.into(); WORD_SIZE];
200 let value_3: [Felt; 4] = [3_u32.into(); WORD_SIZE];
201
202 // insert key-value 1
203 {
204 let old_value_1 = smt.insert(key_1, value_1);
205 assert_eq!(old_value_1, EMPTY_WORD);
206
207 assert_eq!(smt.get_leaf(&key_1), SmtLeaf::Single((key_1, value_1)));
208 }
209
210 // insert key-value 2
211 {
212 let old_value_2 = smt.insert(key_2, value_2);
213 assert_eq!(old_value_2, EMPTY_WORD);
214
215 assert_eq!(
216 smt.get_leaf(&key_2),
217 SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)])
218 );
219 }
220
221 // insert key-value 3
222 {
223 let old_value_3 = smt.insert(key_3, value_3);
224 assert_eq!(old_value_3, EMPTY_WORD);
225
226 assert_eq!(
227 smt.get_leaf(&key_3),
228 SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2), (key_3, value_3)])
229 );
230 }
231
232 // remove key 3
233 {
234 let old_value_3 = smt.insert(key_3, EMPTY_WORD);
235 assert_eq!(old_value_3, value_3);
236
237 assert_eq!(
238 smt.get_leaf(&key_3),
239 SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)])
240 );
241 }
242
243 // remove key 2
244 {
245 let old_value_2 = smt.insert(key_2, EMPTY_WORD);
246 assert_eq!(old_value_2, value_2);
247
248 assert_eq!(smt.get_leaf(&key_2), SmtLeaf::Single((key_1, value_1)));
249 }
250
251 // remove key 1
252 {
253 let old_value_1 = smt.insert(key_1, EMPTY_WORD);
254 assert_eq!(old_value_1, value_1);
255
256 assert_eq!(smt.get_leaf(&key_1), SmtLeaf::new_empty(key_1.into()));
257 }
258}
259
260/// Tests that 2 key-value pairs stored in the same leaf have the same path
261#[test]
262fn test_smt_path_to_keys_in_same_leaf_are_equal() {
263 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
264
265 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
266 let key_2: RpoDigest =
267 RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
268
269 let value_1 = [ONE; WORD_SIZE];
270 let value_2 = [2_u32.into(); WORD_SIZE];
271
272 let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
273
274 assert_eq!(smt.open(&key_1), smt.open(&key_2));
275}
276
277/// Tests that an empty leaf hashes to the empty word
278#[test]
279fn test_empty_leaf_hash() {
280 let smt = Smt::default();
281
282 let leaf = smt.get_leaf(&RpoDigest::default());
283 assert_eq!(leaf.hash(), EMPTY_WORD.into());
284}
285
286/// Tests that `get_value()` works as expected
287#[test]
288fn test_smt_get_value() {
289 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, ONE]);
290 let key_2: RpoDigest =
291 RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), 2_u32.into()]);
292
293 let value_1 = [ONE; WORD_SIZE];
294 let value_2 = [2_u32.into(); WORD_SIZE];
295
296 let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
297
298 let returned_value_1 = smt.get_value(&key_1);
299 let returned_value_2 = smt.get_value(&key_2);
300
301 assert_eq!(value_1, returned_value_1);
302 assert_eq!(value_2, returned_value_2);
303
304 // Check that a key with no inserted value returns the empty word
305 let key_no_value =
306 RpoDigest::from([42_u32.into(), 42_u32.into(), 42_u32.into(), 42_u32.into()]);
307
308 assert_eq!(EMPTY_WORD, smt.get_value(&key_no_value));
309}
310
311/// Tests that `entries()` works as expected
312#[test]
313fn test_smt_entries() {
314 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, ONE]);
315 let key_2: RpoDigest =
316 RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), 2_u32.into()]);
317
318 let value_1 = [ONE; WORD_SIZE];
319 let value_2 = [2_u32.into(); WORD_SIZE];
320
321 let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
322
323 let mut entries = smt.entries();
324
325 // Note: for simplicity, we assume the order `(k1,v1), (k2,v2)`. If a new implementation
326 // switches the order, it is OK to modify the order here as well.
327 assert_eq!(&(key_1, value_1), entries.next().unwrap());
328 assert_eq!(&(key_2, value_2), entries.next().unwrap());
329 assert!(entries.next().is_none());
330}
331
332// SMT LEAF
333// --------------------------------------------------------------------------------------------
334
335#[test]
336fn test_empty_smt_leaf_serialization() {
337 let empty_leaf = SmtLeaf::new_empty(LeafIndex::new_max_depth(42));
338
339 let mut serialized = empty_leaf.to_bytes();
340 // extend buffer with random bytes
341 serialized.extend([1, 2, 3, 4, 5]);
342 let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
343
344 assert_eq!(empty_leaf, deserialized);
345}
346
347#[test]
348fn test_single_smt_leaf_serialization() {
349 let single_leaf = SmtLeaf::new_single(
350 RpoDigest::from([10_u32.into(), 11_u32.into(), 12_u32.into(), 13_u32.into()]),
351 [1_u32.into(), 2_u32.into(), 3_u32.into(), 4_u32.into()],
352 );
353
354 let mut serialized = single_leaf.to_bytes();
355 // extend buffer with random bytes
356 serialized.extend([1, 2, 3, 4, 5]);
357 let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
358
359 assert_eq!(single_leaf, deserialized);
360}
361
362#[test]
363fn test_multiple_smt_leaf_serialization_success() {
364 let multiple_leaf = SmtLeaf::new_multiple(vec![
365 (
366 RpoDigest::from([10_u32.into(), 11_u32.into(), 12_u32.into(), 13_u32.into()]),
367 [1_u32.into(), 2_u32.into(), 3_u32.into(), 4_u32.into()],
368 ),
369 (
370 RpoDigest::from([100_u32.into(), 101_u32.into(), 102_u32.into(), 13_u32.into()]),
371 [11_u32.into(), 12_u32.into(), 13_u32.into(), 14_u32.into()],
372 ),
373 ])
374 .unwrap();
375
376 let mut serialized = multiple_leaf.to_bytes();
377 // extend buffer with random bytes
378 serialized.extend([1, 2, 3, 4, 5]);
379 let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
380
381 assert_eq!(multiple_leaf, deserialized);
382}
383
384// HELPERS
385// --------------------------------------------------------------------------------------------
386
387fn build_empty_or_single_leaf_node(key: RpoDigest, value: Word) -> RpoDigest {
388 if value == EMPTY_WORD {
389 SmtLeaf::new_empty(key.into()).hash()
390 } else {
391 SmtLeaf::Single((key, value)).hash()
392 }
393}
394
395fn build_multiple_leaf_node(kv_pairs: &[(RpoDigest, Word)]) -> RpoDigest {
396 let elements: Vec<Felt> = kv_pairs
397 .iter()
398 .flat_map(|(key, value)| {
399 let key_elements = key.into_iter();
400 let value_elements = (*value).into_iter();
401
402 key_elements.chain(value_elements)
403 })
404 .collect();
405
406 Rpo256::hash_elements(&elements)
407}