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::{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/// Tests that 2 key-value pairs stored in the same leaf have the same path
262#[test]
263fn test_smt_path_to_keys_in_same_leaf_are_equal() {
264 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
265
266 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
267 let key_2: RpoDigest =
268 RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
269
270 let value_1 = [ONE; WORD_SIZE];
271 let value_2 = [2_u32.into(); WORD_SIZE];
272
273 let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
274
275 assert_eq!(smt.open(&key_1), smt.open(&key_2));
276}
277
278/// Tests that an empty leaf hashes to the empty word
279#[test]
280fn test_empty_leaf_hash() {
281 let smt = Smt::default();
282
283 let leaf = smt.get_leaf(&RpoDigest::default());
284 assert_eq!(leaf.hash(), EMPTY_WORD.into());
285}
286
287/// Tests that `get_value()` works as expected
288#[test]
289fn test_smt_get_value() {
290 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, ONE]);
291 let key_2: RpoDigest = RpoDigest::from([2_u32, 2_u32, 2_u32, 2_u32]);
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 = RpoDigest::from([42_u32, 42_u32, 42_u32, 42_u32]);
306
307 assert_eq!(EMPTY_WORD, smt.get_value(&key_no_value));
308}
309
310/// Tests that `entries()` works as expected
311#[test]
312fn test_smt_entries() {
313 let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, ONE]);
314 let key_2: RpoDigest = RpoDigest::from([2_u32, 2_u32, 2_u32, 2_u32]);
315
316 let value_1 = [ONE; WORD_SIZE];
317 let value_2 = [2_u32.into(); WORD_SIZE];
318
319 let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
320
321 let mut entries = smt.entries();
322
323 // Note: for simplicity, we assume the order `(k1,v1), (k2,v2)`. If a new implementation
324 // switches the order, it is OK to modify the order here as well.
325 assert_eq!(&(key_1, value_1), entries.next().unwrap());
326 assert_eq!(&(key_2, value_2), entries.next().unwrap());
327 assert!(entries.next().is_none());
328}
329
330// SMT LEAF
331// --------------------------------------------------------------------------------------------
332
333#[test]
334fn test_empty_smt_leaf_serialization() {
335 let empty_leaf = SmtLeaf::new_empty(LeafIndex::new_max_depth(42));
336
337 let mut serialized = empty_leaf.to_bytes();
338 // extend buffer with random bytes
339 serialized.extend([1, 2, 3, 4, 5]);
340 let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
341
342 assert_eq!(empty_leaf, deserialized);
343}
344
345#[test]
346fn test_single_smt_leaf_serialization() {
347 let single_leaf = SmtLeaf::new_single(
348 RpoDigest::from([10_u32, 11_u32, 12_u32, 13_u32]),
349 [1_u32.into(), 2_u32.into(), 3_u32.into(), 4_u32.into()],
350 );
351
352 let mut serialized = single_leaf.to_bytes();
353 // extend buffer with random bytes
354 serialized.extend([1, 2, 3, 4, 5]);
355 let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
356
357 assert_eq!(single_leaf, deserialized);
358}
359
360#[test]
361fn test_multiple_smt_leaf_serialization_success() {
362 let multiple_leaf = SmtLeaf::new_multiple(vec![
363 (
364 RpoDigest::from([10_u32, 11_u32, 12_u32, 13_u32]),
365 [1_u32.into(), 2_u32.into(), 3_u32.into(), 4_u32.into()],
366 ),
367 (
368 RpoDigest::from([100_u32, 101_u32, 102_u32, 13_u32]),
369 [11_u32.into(), 12_u32.into(), 13_u32.into(), 14_u32.into()],
370 ),
371 ])
372 .unwrap();
373
374 let mut serialized = multiple_leaf.to_bytes();
375 // extend buffer with random bytes
376 serialized.extend([1, 2, 3, 4, 5]);
377 let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
378
379 assert_eq!(multiple_leaf, deserialized);
380}
381
382// HELPERS
383// --------------------------------------------------------------------------------------------
384
385fn build_empty_or_single_leaf_node(key: RpoDigest, value: Word) -> RpoDigest {
386 if value == EMPTY_WORD {
387 SmtLeaf::new_empty(key.into()).hash()
388 } else {
389 SmtLeaf::Single((key, value)).hash()
390 }
391}
392
393fn build_multiple_leaf_node(kv_pairs: &[(RpoDigest, Word)]) -> RpoDigest {
394 let elements: Vec<Felt> = kv_pairs
395 .iter()
396 .flat_map(|(key, value)| {
397 let key_elements = key.into_iter();
398 let value_elements = (*value).into_iter();
399
400 key_elements.chain(value_elements)
401 })
402 .collect();
403
404 Rpo256::hash_elements(&elements)
405}