miden-crypto/src/merkle/tiered_smt/tests.rs
src/merkle/tiered_smt/tests.rsRaw
1use super::{
2 super::{super::ONE, empty_roots::EMPTY_WORD, Felt, MerkleStore, WORD_SIZE, ZERO},
3 EmptySubtreeRoots, InnerNodeInfo, NodeIndex, Rpo256, RpoDigest, TieredSmt, Vec, Word,
4};
5
6// INSERTION TESTS
7// ================================================================================================
8
9#[test]
10fn tsmt_insert_one() {
11 let mut smt = TieredSmt::default();
12 let mut store = MerkleStore::default();
13
14 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
15 let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
16 let value = [ONE; WORD_SIZE];
17
18 // since the tree is empty, the first node will be inserted at depth 16 and the index will be
19 // 16 most significant bits of the key
20 let index = NodeIndex::make(16, raw >> 48);
21 let leaf_node = build_leaf_node(key, value, 16);
22 let tree_root = store.set_node(smt.root(), index, leaf_node).unwrap().root;
23
24 smt.insert(key, value);
25
26 assert_eq!(smt.root(), tree_root);
27
28 // make sure the value was inserted, and the node is at the expected index
29 assert_eq!(smt.get_value(key), value);
30 assert_eq!(smt.get_node(index).unwrap(), leaf_node);
31
32 // make sure the paths we get from the store and the tree match
33 let expected_path = store.get_path(tree_root, index).unwrap();
34 assert_eq!(smt.get_path(index).unwrap(), expected_path.path);
35
36 // make sure inner nodes match
37 let expected_nodes = get_non_empty_nodes(&store);
38 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
39 assert_eq!(actual_nodes.len(), expected_nodes.len());
40 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
41
42 // make sure leaves are returned correctly
43 let mut leaves = smt.upper_leaves();
44 assert_eq!(leaves.next(), Some((leaf_node, key, value)));
45 assert_eq!(leaves.next(), None);
46}
47
48#[test]
49fn tsmt_insert_two_16() {
50 let mut smt = TieredSmt::default();
51 let mut store = MerkleStore::default();
52
53 // --- insert the first value ---------------------------------------------
54 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
55 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
56 let val_a = [ONE; WORD_SIZE];
57 smt.insert(key_a, val_a);
58
59 // --- insert the second value --------------------------------------------
60 // the key for this value has the same 16-bit prefix as the key for the first value,
61 // thus, on insertions, both values should be pushed to depth 32 tier
62 let raw_b = 0b_10101010_10101010_10011111_11111111_10010110_10010011_11100000_00000000_u64;
63 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
64 let val_b = [Felt::new(2); WORD_SIZE];
65 smt.insert(key_b, val_b);
66
67 // --- build Merkle store with equivalent data ----------------------------
68 let mut tree_root = get_init_root();
69 let index_a = NodeIndex::make(32, raw_a >> 32);
70 let leaf_node_a = build_leaf_node(key_a, val_a, 32);
71 tree_root = store.set_node(tree_root, index_a, leaf_node_a).unwrap().root;
72
73 let index_b = NodeIndex::make(32, raw_b >> 32);
74 let leaf_node_b = build_leaf_node(key_b, val_b, 32);
75 tree_root = store.set_node(tree_root, index_b, leaf_node_b).unwrap().root;
76
77 // --- verify that data is consistent between store and tree --------------
78
79 assert_eq!(smt.root(), tree_root);
80
81 assert_eq!(smt.get_value(key_a), val_a);
82 assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
83 let expected_path = store.get_path(tree_root, index_a).unwrap().path;
84 assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
85
86 assert_eq!(smt.get_value(key_b), val_b);
87 assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
88 let expected_path = store.get_path(tree_root, index_b).unwrap().path;
89 assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
90
91 // make sure inner nodes match - the store contains more entries because it keeps track of
92 // all prior state - so, we don't check that the number of inner nodes is the same in both
93 let expected_nodes = get_non_empty_nodes(&store);
94 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
95 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
96
97 // make sure leaves are returned correctly
98 let mut leaves = smt.upper_leaves();
99 assert_eq!(leaves.next(), Some((leaf_node_a, key_a, val_a)));
100 assert_eq!(leaves.next(), Some((leaf_node_b, key_b, val_b)));
101 assert_eq!(leaves.next(), None);
102}
103
104#[test]
105fn tsmt_insert_two_32() {
106 let mut smt = TieredSmt::default();
107 let mut store = MerkleStore::default();
108
109 // --- insert the first value ---------------------------------------------
110 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
111 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
112 let val_a = [ONE; WORD_SIZE];
113 smt.insert(key_a, val_a);
114
115 // --- insert the second value --------------------------------------------
116 // the key for this value has the same 32-bit prefix as the key for the first value,
117 // thus, on insertions, both values should be pushed to depth 48 tier
118 let raw_b = 0b_10101010_10101010_00011111_11111111_00010110_10010011_11100000_00000000_u64;
119 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
120 let val_b = [Felt::new(2); WORD_SIZE];
121 smt.insert(key_b, val_b);
122
123 // --- build Merkle store with equivalent data ----------------------------
124 let mut tree_root = get_init_root();
125 let index_a = NodeIndex::make(48, raw_a >> 16);
126 let leaf_node_a = build_leaf_node(key_a, val_a, 48);
127 tree_root = store.set_node(tree_root, index_a, leaf_node_a).unwrap().root;
128
129 let index_b = NodeIndex::make(48, raw_b >> 16);
130 let leaf_node_b = build_leaf_node(key_b, val_b, 48);
131 tree_root = store.set_node(tree_root, index_b, leaf_node_b).unwrap().root;
132
133 // --- verify that data is consistent between store and tree --------------
134
135 assert_eq!(smt.root(), tree_root);
136
137 assert_eq!(smt.get_value(key_a), val_a);
138 assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
139 let expected_path = store.get_path(tree_root, index_a).unwrap().path;
140 assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
141
142 assert_eq!(smt.get_value(key_b), val_b);
143 assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
144 let expected_path = store.get_path(tree_root, index_b).unwrap().path;
145 assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
146
147 // make sure inner nodes match - the store contains more entries because it keeps track of
148 // all prior state - so, we don't check that the number of inner nodes is the same in both
149 let expected_nodes = get_non_empty_nodes(&store);
150 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
151 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
152}
153
154#[test]
155fn tsmt_insert_three() {
156 let mut smt = TieredSmt::default();
157 let mut store = MerkleStore::default();
158
159 // --- insert the first value ---------------------------------------------
160 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
161 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
162 let val_a = [ONE; WORD_SIZE];
163 smt.insert(key_a, val_a);
164
165 // --- insert the second value --------------------------------------------
166 // the key for this value has the same 16-bit prefix as the key for the first value,
167 // thus, on insertions, both values should be pushed to depth 32 tier
168 let raw_b = 0b_10101010_10101010_10011111_11111111_10010110_10010011_11100000_00000000_u64;
169 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
170 let val_b = [Felt::new(2); WORD_SIZE];
171 smt.insert(key_b, val_b);
172
173 // --- insert the third value ---------------------------------------------
174 // the key for this value has the same 16-bit prefix as the keys for the first two,
175 // values; thus, on insertions, it will be inserted into depth 32 tier, but will not
176 // affect locations of the other two values
177 let raw_c = 0b_10101010_10101010_11011111_11111111_10010110_10010011_11100000_00000000_u64;
178 let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
179 let val_c = [Felt::new(3); WORD_SIZE];
180 smt.insert(key_c, val_c);
181
182 // --- build Merkle store with equivalent data ----------------------------
183 let mut tree_root = get_init_root();
184 let index_a = NodeIndex::make(32, raw_a >> 32);
185 let leaf_node_a = build_leaf_node(key_a, val_a, 32);
186 tree_root = store.set_node(tree_root, index_a, leaf_node_a).unwrap().root;
187
188 let index_b = NodeIndex::make(32, raw_b >> 32);
189 let leaf_node_b = build_leaf_node(key_b, val_b, 32);
190 tree_root = store.set_node(tree_root, index_b, leaf_node_b).unwrap().root;
191
192 let index_c = NodeIndex::make(32, raw_c >> 32);
193 let leaf_node_c = build_leaf_node(key_c, val_c, 32);
194 tree_root = store.set_node(tree_root, index_c, leaf_node_c).unwrap().root;
195
196 // --- verify that data is consistent between store and tree --------------
197
198 assert_eq!(smt.root(), tree_root);
199
200 assert_eq!(smt.get_value(key_a), val_a);
201 assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
202 let expected_path = store.get_path(tree_root, index_a).unwrap().path;
203 assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
204
205 assert_eq!(smt.get_value(key_b), val_b);
206 assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
207 let expected_path = store.get_path(tree_root, index_b).unwrap().path;
208 assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
209
210 assert_eq!(smt.get_value(key_c), val_c);
211 assert_eq!(smt.get_node(index_c).unwrap(), leaf_node_c);
212 let expected_path = store.get_path(tree_root, index_c).unwrap().path;
213 assert_eq!(smt.get_path(index_c).unwrap(), expected_path);
214
215 // make sure inner nodes match - the store contains more entries because it keeps track of
216 // all prior state - so, we don't check that the number of inner nodes is the same in both
217 let expected_nodes = get_non_empty_nodes(&store);
218 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
219 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
220}
221
222// UPDATE TESTS
223// ================================================================================================
224
225#[test]
226fn tsmt_update() {
227 let mut smt = TieredSmt::default();
228 let mut store = MerkleStore::default();
229
230 // --- insert a value into the tree ---------------------------------------
231 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
232 let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
233 let value_a = [ONE; WORD_SIZE];
234 smt.insert(key, value_a);
235
236 // --- update the value ---------------------------------------------------
237 let value_b = [Felt::new(2); WORD_SIZE];
238 smt.insert(key, value_b);
239
240 // --- verify consistency -------------------------------------------------
241 let mut tree_root = get_init_root();
242 let index = NodeIndex::make(16, raw >> 48);
243 let leaf_node = build_leaf_node(key, value_b, 16);
244 tree_root = store.set_node(tree_root, index, leaf_node).unwrap().root;
245
246 assert_eq!(smt.root(), tree_root);
247
248 assert_eq!(smt.get_value(key), value_b);
249 assert_eq!(smt.get_node(index).unwrap(), leaf_node);
250 let expected_path = store.get_path(tree_root, index).unwrap().path;
251 assert_eq!(smt.get_path(index).unwrap(), expected_path);
252
253 // make sure inner nodes match - the store contains more entries because it keeps track of
254 // all prior state - so, we don't check that the number of inner nodes is the same in both
255 let expected_nodes = get_non_empty_nodes(&store);
256 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
257 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
258}
259
260// DELETION TESTS
261// ================================================================================================
262
263#[test]
264fn tsmt_delete_16() {
265 let mut smt = TieredSmt::default();
266
267 // --- insert a value into the tree ---------------------------------------
268 let smt0 = smt.clone();
269 let raw_a = 0b_01010101_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
270 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
271 let value_a = [ONE, ONE, ONE, ONE];
272 smt.insert(key_a, value_a);
273
274 // --- insert another value into the tree ---------------------------------
275 let smt1 = smt.clone();
276 let raw_b = 0b_01011111_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
277 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
278 let value_b = [ONE, ONE, ONE, ZERO];
279 smt.insert(key_b, value_b);
280
281 // --- delete the last inserted value -------------------------------------
282 assert_eq!(smt.insert(key_b, [ZERO; 4]), value_b);
283 assert_eq!(smt, smt1);
284
285 // --- delete the first inserted value ------------------------------------
286 assert_eq!(smt.insert(key_a, [ZERO; 4]), value_a);
287 assert_eq!(smt, smt0);
288}
289
290#[test]
291fn tsmt_delete_32() {
292 let mut smt = TieredSmt::default();
293
294 // --- insert a value into the tree ---------------------------------------
295 let smt0 = smt.clone();
296 let raw_a = 0b_01010101_01101100_01111111_11111111_10010110_10010011_11100000_00000000_u64;
297 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
298 let value_a = [ONE, ONE, ONE, ONE];
299 smt.insert(key_a, value_a);
300
301 // --- insert another with the same 16-bit prefix into the tree -----------
302 let smt1 = smt.clone();
303 let raw_b = 0b_01010101_01101100_00111111_11111111_10010110_10010011_11100000_00000000_u64;
304 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
305 let value_b = [ONE, ONE, ONE, ZERO];
306 smt.insert(key_b, value_b);
307
308 // --- insert the 3rd value with the same 16-bit prefix into the tree -----
309 let smt2 = smt.clone();
310 let raw_c = 0b_01010101_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
311 let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
312 let value_c = [ONE, ONE, ZERO, ZERO];
313 smt.insert(key_c, value_c);
314
315 // --- delete the last inserted value -------------------------------------
316 assert_eq!(smt.insert(key_c, [ZERO; 4]), value_c);
317 assert_eq!(smt, smt2);
318
319 // --- delete the last inserted value -------------------------------------
320 assert_eq!(smt.insert(key_b, [ZERO; 4]), value_b);
321 assert_eq!(smt, smt1);
322
323 // --- delete the first inserted value ------------------------------------
324 assert_eq!(smt.insert(key_a, [ZERO; 4]), value_a);
325 assert_eq!(smt, smt0);
326}
327
328#[test]
329fn tsmt_delete_48_same_32_bit_prefix() {
330 let mut smt = TieredSmt::default();
331
332 // test the case when all values share the same 32-bit prefix
333
334 // --- insert a value into the tree ---------------------------------------
335 let smt0 = smt.clone();
336 let raw_a = 0b_01010101_01010101_11111111_11111111_10010110_10010011_11100000_00000000_u64;
337 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
338 let value_a = [ONE, ONE, ONE, ONE];
339 smt.insert(key_a, value_a);
340
341 // --- insert another with the same 32-bit prefix into the tree -----------
342 let smt1 = smt.clone();
343 let raw_b = 0b_01010101_01010101_11111111_11111111_11010110_10010011_11100000_00000000_u64;
344 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
345 let value_b = [ONE, ONE, ONE, ZERO];
346 smt.insert(key_b, value_b);
347
348 // --- insert the 3rd value with the same 32-bit prefix into the tree -----
349 let smt2 = smt.clone();
350 let raw_c = 0b_01010101_01010101_11111111_11111111_11110110_10010011_11100000_00000000_u64;
351 let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
352 let value_c = [ONE, ONE, ZERO, ZERO];
353 smt.insert(key_c, value_c);
354
355 // --- delete the last inserted value -------------------------------------
356 assert_eq!(smt.insert(key_c, [ZERO; 4]), value_c);
357 assert_eq!(smt, smt2);
358
359 // --- delete the last inserted value -------------------------------------
360 assert_eq!(smt.insert(key_b, [ZERO; 4]), value_b);
361 assert_eq!(smt, smt1);
362
363 // --- delete the first inserted value ------------------------------------
364 assert_eq!(smt.insert(key_a, [ZERO; 4]), value_a);
365 assert_eq!(smt, smt0);
366}
367
368#[test]
369fn tsmt_delete_48_mixed_prefix() {
370 let mut smt = TieredSmt::default();
371
372 // test the case when some values share a 32-bit prefix and others share a 16-bit prefix
373
374 // --- insert a value into the tree ---------------------------------------
375 let smt0 = smt.clone();
376 let raw_a = 0b_01010101_01010101_11111111_11111111_10010110_10010011_11100000_00000000_u64;
377 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
378 let value_a = [ONE, ONE, ONE, ONE];
379 smt.insert(key_a, value_a);
380
381 // --- insert another with the same 16-bit prefix into the tree -----------
382 let smt1 = smt.clone();
383 let raw_b = 0b_01010101_01010101_01111111_11111111_10010110_10010011_11100000_00000000_u64;
384 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
385 let value_b = [ONE, ONE, ONE, ZERO];
386 smt.insert(key_b, value_b);
387
388 // --- insert a value with the same 32-bit prefix as the first value -----
389 let smt2 = smt.clone();
390 let raw_c = 0b_01010101_01010101_11111111_11111111_11010110_10010011_11100000_00000000_u64;
391 let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
392 let value_c = [ONE, ONE, ZERO, ZERO];
393 smt.insert(key_c, value_c);
394
395 // --- insert another value with the same 32-bit prefix as the first value
396 let smt3 = smt.clone();
397 let raw_d = 0b_01010101_01010101_11111111_11111111_11110110_10010011_11100000_00000000_u64;
398 let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_d)]);
399 let value_d = [ONE, ZERO, ZERO, ZERO];
400 smt.insert(key_d, value_d);
401
402 // --- delete the inserted values one-by-one ------------------------------
403 assert_eq!(smt.insert(key_d, [ZERO; 4]), value_d);
404 assert_eq!(smt, smt3);
405
406 assert_eq!(smt.insert(key_c, [ZERO; 4]), value_c);
407 assert_eq!(smt, smt2);
408
409 assert_eq!(smt.insert(key_b, [ZERO; 4]), value_b);
410 assert_eq!(smt, smt1);
411
412 assert_eq!(smt.insert(key_a, [ZERO; 4]), value_a);
413 assert_eq!(smt, smt0);
414}
415
416#[test]
417fn tsmt_delete_64() {
418 let mut smt = TieredSmt::default();
419
420 // test the case when all values share the same 48-bit prefix
421
422 // --- insert a value into the tree ---------------------------------------
423 let smt0 = smt.clone();
424 let raw_a = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000000_u64;
425 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
426 let value_a = [ONE, ONE, ONE, ONE];
427 smt.insert(key_a, value_a);
428
429 // --- insert a value with the same 48-bit prefix into the tree -----------
430 let smt1 = smt.clone();
431 let raw_b = 0b_01010101_01010101_11111111_11111111_10110101_10101010_10111100_00000000_u64;
432 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
433 let value_b = [ONE, ONE, ONE, ZERO];
434 smt.insert(key_b, value_b);
435
436 // --- insert a value with the same 32-bit prefix into the tree -----------
437 let smt2 = smt.clone();
438 let raw_c = 0b_01010101_01010101_11111111_11111111_11111101_10101010_10111100_00000000_u64;
439 let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
440 let value_c = [ONE, ONE, ZERO, ZERO];
441 smt.insert(key_c, value_c);
442
443 let smt3 = smt.clone();
444 let raw_d = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000000_u64;
445 let key_d = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_d)]);
446 let value_d = [ONE, ZERO, ZERO, ZERO];
447 smt.insert(key_d, value_d);
448
449 // --- delete the last inserted value -------------------------------------
450 assert_eq!(smt.insert(key_d, [ZERO; 4]), value_d);
451 assert_eq!(smt, smt3);
452
453 assert_eq!(smt.insert(key_c, [ZERO; 4]), value_c);
454 assert_eq!(smt, smt2);
455
456 assert_eq!(smt.insert(key_b, [ZERO; 4]), value_b);
457 assert_eq!(smt, smt1);
458
459 assert_eq!(smt.insert(key_a, [ZERO; 4]), value_a);
460 assert_eq!(smt, smt0);
461}
462
463// BOTTOM TIER TESTS
464// ================================================================================================
465
466#[test]
467fn tsmt_bottom_tier() {
468 let mut smt = TieredSmt::default();
469 let mut store = MerkleStore::default();
470
471 // common prefix for the keys
472 let prefix = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
473
474 // --- insert the first value ---------------------------------------------
475 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(prefix)]);
476 let val_a = [ONE; WORD_SIZE];
477 smt.insert(key_a, val_a);
478
479 // --- insert the second value --------------------------------------------
480 // this key has the same 64-bit prefix and thus both values should end up in the same
481 // node at depth 64
482 let key_b = RpoDigest::from([ZERO, ONE, ONE, Felt::new(prefix)]);
483 let val_b = [Felt::new(2); WORD_SIZE];
484 smt.insert(key_b, val_b);
485
486 // --- build Merkle store with equivalent data ----------------------------
487 let index = NodeIndex::make(64, prefix);
488 // to build bottom leaf we sort by key starting with the least significant element, thus
489 // key_b is smaller than key_a.
490 let leaf_node = build_bottom_leaf_node(&[key_b, key_a], &[val_b, val_a]);
491 let mut tree_root = get_init_root();
492 tree_root = store.set_node(tree_root, index, leaf_node).unwrap().root;
493
494 // --- verify that data is consistent between store and tree --------------
495
496 assert_eq!(smt.root(), tree_root);
497
498 assert_eq!(smt.get_value(key_a), val_a);
499 assert_eq!(smt.get_value(key_b), val_b);
500
501 assert_eq!(smt.get_node(index).unwrap(), leaf_node);
502 let expected_path = store.get_path(tree_root, index).unwrap().path;
503 assert_eq!(smt.get_path(index).unwrap(), expected_path);
504
505 // make sure inner nodes match - the store contains more entries because it keeps track of
506 // all prior state - so, we don't check that the number of inner nodes is the same in both
507 let expected_nodes = get_non_empty_nodes(&store);
508 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
509 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
510
511 // make sure leaves are returned correctly
512 let smt_clone = smt.clone();
513 let mut leaves = smt_clone.bottom_leaves();
514 assert_eq!(leaves.next(), Some((leaf_node, vec![(key_b, val_b), (key_a, val_a)])));
515 assert_eq!(leaves.next(), None);
516
517 // --- update a leaf at the bottom tier -------------------------------------------------------
518
519 let val_a2 = [Felt::new(3); WORD_SIZE];
520 assert_eq!(smt.insert(key_a, val_a2), val_a);
521
522 let leaf_node = build_bottom_leaf_node(&[key_b, key_a], &[val_b, val_a2]);
523 store.set_node(tree_root, index, leaf_node).unwrap();
524
525 let expected_nodes = get_non_empty_nodes(&store);
526 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
527 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
528
529 let mut leaves = smt.bottom_leaves();
530 assert_eq!(leaves.next(), Some((leaf_node, vec![(key_b, val_b), (key_a, val_a2)])));
531 assert_eq!(leaves.next(), None);
532}
533
534#[test]
535fn tsmt_bottom_tier_two() {
536 let mut smt = TieredSmt::default();
537 let mut store = MerkleStore::default();
538
539 // --- insert the first value ---------------------------------------------
540 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
541 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
542 let val_a = [ONE; WORD_SIZE];
543 smt.insert(key_a, val_a);
544
545 // --- insert the second value --------------------------------------------
546 // the key for this value has the same 48-bit prefix as the key for the first value,
547 // thus, on insertions, both should end up in different nodes at depth 64
548 let raw_b = 0b_10101010_10101010_00011111_11111111_10010110_10010011_01100000_00000000_u64;
549 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
550 let val_b = [Felt::new(2); WORD_SIZE];
551 smt.insert(key_b, val_b);
552
553 // --- build Merkle store with equivalent data ----------------------------
554 let mut tree_root = get_init_root();
555 let index_a = NodeIndex::make(64, raw_a);
556 let leaf_node_a = build_bottom_leaf_node(&[key_a], &[val_a]);
557 tree_root = store.set_node(tree_root, index_a, leaf_node_a).unwrap().root;
558
559 let index_b = NodeIndex::make(64, raw_b);
560 let leaf_node_b = build_bottom_leaf_node(&[key_b], &[val_b]);
561 tree_root = store.set_node(tree_root, index_b, leaf_node_b).unwrap().root;
562
563 // --- verify that data is consistent between store and tree --------------
564
565 assert_eq!(smt.root(), tree_root);
566
567 assert_eq!(smt.get_value(key_a), val_a);
568 assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
569 let expected_path = store.get_path(tree_root, index_a).unwrap().path;
570 assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
571
572 assert_eq!(smt.get_value(key_b), val_b);
573 assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
574 let expected_path = store.get_path(tree_root, index_b).unwrap().path;
575 assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
576
577 // make sure inner nodes match - the store contains more entries because it keeps track of
578 // all prior state - so, we don't check that the number of inner nodes is the same in both
579 let expected_nodes = get_non_empty_nodes(&store);
580 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
581 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
582
583 // make sure leaves are returned correctly
584 let mut leaves = smt.bottom_leaves();
585 assert_eq!(leaves.next(), Some((leaf_node_b, vec![(key_b, val_b)])));
586 assert_eq!(leaves.next(), Some((leaf_node_a, vec![(key_a, val_a)])));
587 assert_eq!(leaves.next(), None);
588}
589
590// GET PROOF TESTS
591// ================================================================================================
592
593#[test]
594fn tsmt_get_proof() {
595 let mut smt = TieredSmt::default();
596
597 // --- insert a value into the tree ---------------------------------------
598 let raw_a = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000000_u64;
599 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
600 let value_a = [ONE, ONE, ONE, ONE];
601 smt.insert(key_a, value_a);
602
603 // --- insert a value with the same 48-bit prefix into the tree -----------
604 let raw_b = 0b_01010101_01010101_11111111_11111111_10110101_10101010_10111100_00000000_u64;
605 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
606 let value_b = [ONE, ONE, ONE, ZERO];
607 smt.insert(key_b, value_b);
608
609 let smt_alt = smt.clone();
610
611 // --- insert a value with the same 32-bit prefix into the tree -----------
612 let raw_c = 0b_01010101_01010101_11111111_11111111_11111101_10101010_10111100_00000000_u64;
613 let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
614 let value_c = [ONE, ONE, ZERO, ZERO];
615 smt.insert(key_c, value_c);
616
617 // --- insert a value with the same 64-bit prefix as A into the tree ------
618 let raw_d = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000000_u64;
619 let key_d = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_d)]);
620 let value_d = [ONE, ZERO, ZERO, ZERO];
621 smt.insert(key_d, value_d);
622
623 // at this point the tree looks as follows:
624 // - A and D are located in the same node at depth 64.
625 // - B is located at depth 64 and shares the same 48-bit prefix with A and D.
626 // - C is located at depth 48 and shares the same 32-bit prefix with A, B, and D.
627
628 // --- generate proof for key A and test that it verifies correctly -------
629 let proof = smt.prove(key_a);
630 assert!(proof.verify_membership(&key_a, &value_a, &smt.root()));
631
632 assert!(!proof.verify_membership(&key_a, &value_b, &smt.root()));
633 assert!(!proof.verify_membership(&key_a, &EMPTY_WORD, &smt.root()));
634 assert!(!proof.verify_membership(&key_b, &value_a, &smt.root()));
635 assert!(!proof.verify_membership(&key_a, &value_a, &smt_alt.root()));
636
637 assert_eq!(proof.get(&key_a), Some(value_a));
638 assert_eq!(proof.get(&key_b), None);
639
640 // since A and D are stored in the same node, we should be able to use the proof to verify
641 // membership of D
642 assert!(proof.verify_membership(&key_d, &value_d, &smt.root()));
643 assert_eq!(proof.get(&key_d), Some(value_d));
644
645 // --- generate proof for key B and test that it verifies correctly -------
646 let proof = smt.prove(key_b);
647 assert!(proof.verify_membership(&key_b, &value_b, &smt.root()));
648
649 assert!(!proof.verify_membership(&key_b, &value_a, &smt.root()));
650 assert!(!proof.verify_membership(&key_b, &EMPTY_WORD, &smt.root()));
651 assert!(!proof.verify_membership(&key_a, &value_b, &smt.root()));
652 assert!(!proof.verify_membership(&key_b, &value_b, &smt_alt.root()));
653
654 assert_eq!(proof.get(&key_b), Some(value_b));
655 assert_eq!(proof.get(&key_a), None);
656
657 // --- generate proof for key C and test that it verifies correctly -------
658 let proof = smt.prove(key_c);
659 assert!(proof.verify_membership(&key_c, &value_c, &smt.root()));
660
661 assert!(!proof.verify_membership(&key_c, &value_a, &smt.root()));
662 assert!(!proof.verify_membership(&key_c, &EMPTY_WORD, &smt.root()));
663 assert!(!proof.verify_membership(&key_a, &value_c, &smt.root()));
664 assert!(!proof.verify_membership(&key_c, &value_c, &smt_alt.root()));
665
666 assert_eq!(proof.get(&key_c), Some(value_c));
667 assert_eq!(proof.get(&key_b), None);
668
669 // --- generate proof for key D and test that it verifies correctly -------
670 let proof = smt.prove(key_d);
671 assert!(proof.verify_membership(&key_d, &value_d, &smt.root()));
672
673 assert!(!proof.verify_membership(&key_d, &value_b, &smt.root()));
674 assert!(!proof.verify_membership(&key_d, &EMPTY_WORD, &smt.root()));
675 assert!(!proof.verify_membership(&key_b, &value_d, &smt.root()));
676 assert!(!proof.verify_membership(&key_d, &value_d, &smt_alt.root()));
677
678 assert_eq!(proof.get(&key_d), Some(value_d));
679 assert_eq!(proof.get(&key_b), None);
680
681 // since A and D are stored in the same node, we should be able to use the proof to verify
682 // membership of A
683 assert!(proof.verify_membership(&key_a, &value_a, &smt.root()));
684 assert_eq!(proof.get(&key_a), Some(value_a));
685
686 // --- generate proof for an empty key at depth 64 ------------------------
687 // this key has the same 48-bit prefix as A but is different from B
688 let raw = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000011_u64;
689 let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
690
691 let proof = smt.prove(key);
692 assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root()));
693
694 assert!(!proof.verify_membership(&key, &value_a, &smt.root()));
695 assert!(!proof.verify_membership(&key, &EMPTY_WORD, &smt_alt.root()));
696
697 assert_eq!(proof.get(&key), Some(EMPTY_WORD));
698 assert_eq!(proof.get(&key_b), None);
699
700 // the same proof should verify against any key with the same 64-bit prefix
701 let key2 = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw)]);
702 assert!(proof.verify_membership(&key2, &EMPTY_WORD, &smt.root()));
703 assert_eq!(proof.get(&key2), Some(EMPTY_WORD));
704
705 // but verifying if against a key with the same 63-bit prefix (or smaller) should fail
706 let raw3 = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000010_u64;
707 let key3 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw3)]);
708 assert!(!proof.verify_membership(&key3, &EMPTY_WORD, &smt.root()));
709 assert_eq!(proof.get(&key3), None);
710
711 // --- generate proof for an empty key at depth 48 ------------------------
712 // this key has the same 32-prefix as A, B, C, and D, but is different from C
713 let raw = 0b_01010101_01010101_11111111_11111111_00110101_10101010_11111100_00000000_u64;
714 let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
715
716 let proof = smt.prove(key);
717 assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root()));
718
719 assert!(!proof.verify_membership(&key, &value_a, &smt.root()));
720 assert!(!proof.verify_membership(&key, &EMPTY_WORD, &smt_alt.root()));
721
722 assert_eq!(proof.get(&key), Some(EMPTY_WORD));
723 assert_eq!(proof.get(&key_b), None);
724
725 // the same proof should verify against any key with the same 48-bit prefix
726 let raw2 = 0b_01010101_01010101_11111111_11111111_00110101_10101010_01111100_00000000_u64;
727 let key2 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw2)]);
728 assert!(proof.verify_membership(&key2, &EMPTY_WORD, &smt.root()));
729 assert_eq!(proof.get(&key2), Some(EMPTY_WORD));
730
731 // but verifying against a key with the same 47-bit prefix (or smaller) should fail
732 let raw3 = 0b_01010101_01010101_11111111_11111111_00110101_10101011_11111100_00000000_u64;
733 let key3 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw3)]);
734 assert!(!proof.verify_membership(&key3, &EMPTY_WORD, &smt.root()));
735 assert_eq!(proof.get(&key3), None);
736}
737
738// ERROR TESTS
739// ================================================================================================
740
741#[test]
742fn tsmt_node_not_available() {
743 let mut smt = TieredSmt::default();
744
745 let raw = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
746 let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
747 let value = [ONE; WORD_SIZE];
748
749 // build an index which is just below the inserted leaf node
750 let index = NodeIndex::make(17, raw >> 47);
751
752 // since we haven't inserted the node yet, we should be able to get node and path to this index
753 assert!(smt.get_node(index).is_ok());
754 assert!(smt.get_path(index).is_ok());
755
756 smt.insert(key, value);
757
758 // but once the node is inserted, everything under it should be unavailable
759 assert!(smt.get_node(index).is_err());
760 assert!(smt.get_path(index).is_err());
761
762 let index = NodeIndex::make(32, raw >> 32);
763 assert!(smt.get_node(index).is_err());
764 assert!(smt.get_path(index).is_err());
765
766 let index = NodeIndex::make(34, raw >> 30);
767 assert!(smt.get_node(index).is_err());
768 assert!(smt.get_path(index).is_err());
769
770 let index = NodeIndex::make(50, raw >> 14);
771 assert!(smt.get_node(index).is_err());
772 assert!(smt.get_path(index).is_err());
773
774 let index = NodeIndex::make(64, raw);
775 assert!(smt.get_node(index).is_err());
776 assert!(smt.get_path(index).is_err());
777}
778
779// HELPER FUNCTIONS
780// ================================================================================================
781
782fn get_init_root() -> RpoDigest {
783 EmptySubtreeRoots::empty_hashes(64)[0]
784}
785
786fn build_leaf_node(key: RpoDigest, value: Word, depth: u8) -> RpoDigest {
787 Rpo256::merge_in_domain(&[key, value.into()], depth.into())
788}
789
790fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest {
791 assert_eq!(keys.len(), values.len());
792
793 let mut elements = Vec::with_capacity(keys.len());
794 for (key, val) in keys.iter().zip(values.iter()) {
795 elements.extend_from_slice(key.as_elements());
796 elements.extend_from_slice(val.as_slice());
797 }
798
799 Rpo256::hash_elements(&elements)
800}
801
802fn get_non_empty_nodes(store: &MerkleStore) -> Vec<InnerNodeInfo> {
803 store
804 .inner_nodes()
805 .filter(|node| !is_empty_subtree(&node.value))
806 .collect::<Vec<_>>()
807}
808
809fn is_empty_subtree(node: &RpoDigest) -> bool {
810 EmptySubtreeRoots::empty_hashes(255).contains(node)
811}