arnaucube
arnaucube/miden-cryptoPublic Self-hosted repository
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#[test]
464fn tsmt_delete_64_leaf_promotion() {
465    let mut smt = TieredSmt::default();
466
467    // --- delete from bottom tier (no promotion to upper tiers) --------------
468
469    // insert a value into the tree
470    let raw_a = 0b_01010101_01010101_11111111_11111111_10101010_10101010_11111111_00000000_u64;
471    let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
472    let value_a = [ONE, ONE, ONE, ONE];
473    smt.insert(key_a, value_a);
474
475    // insert another value with a key having the same 64-bit prefix
476    let key_b = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]);
477    let value_b = [ONE, ONE, ONE, ZERO];
478    smt.insert(key_b, value_b);
479
480    // insert a value with a key which shared the same 48-bit prefix
481    let raw_c = 0b_01010101_01010101_11111111_11111111_10101010_10101010_00111111_00000000_u64;
482    let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
483    let value_c = [ONE, ONE, ZERO, ZERO];
484    smt.insert(key_c, value_c);
485
486    // delete entry A and compare to the tree which was built from B and C
487    smt.insert(key_a, EMPTY_WORD);
488
489    let mut expected_smt = TieredSmt::default();
490    expected_smt.insert(key_b, value_b);
491    expected_smt.insert(key_c, value_c);
492    assert_eq!(smt, expected_smt);
493
494    // entries B and C should stay at depth 64
495    assert_eq!(smt.nodes.get_leaf_index(&key_b).0.depth(), 64);
496    assert_eq!(smt.nodes.get_leaf_index(&key_c).0.depth(), 64);
497
498    // --- delete from bottom tier (promotion to depth 48) --------------------
499
500    let mut smt = TieredSmt::default();
501    smt.insert(key_a, value_a);
502    smt.insert(key_b, value_b);
503
504    // insert a value with a key which shared the same 32-bit prefix
505    let raw_c = 0b_01010101_01010101_11111111_11111111_11101010_10101010_11111111_00000000_u64;
506    let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
507    smt.insert(key_c, value_c);
508
509    // delete entry A and compare to the tree which was built from B and C
510    smt.insert(key_a, EMPTY_WORD);
511
512    let mut expected_smt = TieredSmt::default();
513    expected_smt.insert(key_b, value_b);
514    expected_smt.insert(key_c, value_c);
515    assert_eq!(smt, expected_smt);
516
517    // entry B moves to depth 48, entry C stays at depth 48
518    assert_eq!(smt.nodes.get_leaf_index(&key_b).0.depth(), 48);
519    assert_eq!(smt.nodes.get_leaf_index(&key_c).0.depth(), 48);
520
521    // --- delete from bottom tier (promotion to depth 32) --------------------
522
523    let mut smt = TieredSmt::default();
524    smt.insert(key_a, value_a);
525    smt.insert(key_b, value_b);
526
527    // insert a value with a key which shared the same 16-bit prefix
528    let raw_c = 0b_01010101_01010101_01111111_11111111_10101010_10101010_11111111_00000000_u64;
529    let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
530    smt.insert(key_c, value_c);
531
532    // delete entry A and compare to the tree which was built from B and C
533    smt.insert(key_a, EMPTY_WORD);
534
535    let mut expected_smt = TieredSmt::default();
536    expected_smt.insert(key_b, value_b);
537    expected_smt.insert(key_c, value_c);
538    assert_eq!(smt, expected_smt);
539
540    // entry B moves to depth 32, entry C stays at depth 32
541    assert_eq!(smt.nodes.get_leaf_index(&key_b).0.depth(), 32);
542    assert_eq!(smt.nodes.get_leaf_index(&key_c).0.depth(), 32);
543
544    // --- delete from bottom tier (promotion to depth 16) --------------------
545
546    let mut smt = TieredSmt::default();
547    smt.insert(key_a, value_a);
548    smt.insert(key_b, value_b);
549
550    // insert a value with a key which shared prefix < 16 bits
551    let raw_c = 0b_01010101_01010100_11111111_11111111_10101010_10101010_11111111_00000000_u64;
552    let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
553    smt.insert(key_c, value_c);
554
555    // delete entry A and compare to the tree which was built from B and C
556    smt.insert(key_a, EMPTY_WORD);
557
558    let mut expected_smt = TieredSmt::default();
559    expected_smt.insert(key_b, value_b);
560    expected_smt.insert(key_c, value_c);
561    assert_eq!(smt, expected_smt);
562
563    // entry B moves to depth 16, entry C stays at depth 16
564    assert_eq!(smt.nodes.get_leaf_index(&key_b).0.depth(), 16);
565    assert_eq!(smt.nodes.get_leaf_index(&key_c).0.depth(), 16);
566}
567
568#[test]
569fn test_order_sensitivity() {
570    let raw = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000001_u64;
571    let value = [ONE; WORD_SIZE];
572
573    let key_1 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
574    let key_2 = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw)]);
575
576    let mut smt_1 = TieredSmt::default();
577
578    smt_1.insert(key_1, value);
579    smt_1.insert(key_2, value);
580    smt_1.insert(key_2, [ZERO; WORD_SIZE]);
581
582    let mut smt_2 = TieredSmt::default();
583    smt_2.insert(key_1, value);
584
585    assert_eq!(smt_1.root(), smt_2.root());
586}
587
588// BOTTOM TIER TESTS
589// ================================================================================================
590
591#[test]
592fn tsmt_bottom_tier() {
593    let mut smt = TieredSmt::default();
594    let mut store = MerkleStore::default();
595
596    // common prefix for the keys
597    let prefix = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
598
599    // --- insert the first value ---------------------------------------------
600    let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(prefix)]);
601    let val_a = [ONE; WORD_SIZE];
602    smt.insert(key_a, val_a);
603
604    // --- insert the second value --------------------------------------------
605    // this key has the same 64-bit prefix and thus both values should end up in the same
606    // node at depth 64
607    let key_b = RpoDigest::from([ZERO, ONE, ONE, Felt::new(prefix)]);
608    let val_b = [Felt::new(2); WORD_SIZE];
609    smt.insert(key_b, val_b);
610
611    // --- build Merkle store with equivalent data ----------------------------
612    let index = NodeIndex::make(64, prefix);
613    // to build bottom leaf we sort by key starting with the least significant element, thus
614    // key_b is smaller than key_a.
615    let leaf_node = build_bottom_leaf_node(&[key_b, key_a], &[val_b, val_a]);
616    let mut tree_root = get_init_root();
617    tree_root = store.set_node(tree_root, index, leaf_node).unwrap().root;
618
619    // --- verify that data is consistent between store and tree --------------
620
621    assert_eq!(smt.root(), tree_root);
622
623    assert_eq!(smt.get_value(key_a), val_a);
624    assert_eq!(smt.get_value(key_b), val_b);
625
626    assert_eq!(smt.get_node(index).unwrap(), leaf_node);
627    let expected_path = store.get_path(tree_root, index).unwrap().path;
628    assert_eq!(smt.get_path(index).unwrap(), expected_path);
629
630    // make sure inner nodes match - the store contains more entries because it keeps track of
631    // all prior state - so, we don't check that the number of inner nodes is the same in both
632    let expected_nodes = get_non_empty_nodes(&store);
633    let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
634    actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
635
636    // make sure leaves are returned correctly
637    let smt_clone = smt.clone();
638    let mut leaves = smt_clone.bottom_leaves();
639    assert_eq!(leaves.next(), Some((leaf_node, vec![(key_b, val_b), (key_a, val_a)])));
640    assert_eq!(leaves.next(), None);
641
642    // --- update a leaf at the bottom tier -------------------------------------------------------
643
644    let val_a2 = [Felt::new(3); WORD_SIZE];
645    assert_eq!(smt.insert(key_a, val_a2), val_a);
646
647    let leaf_node = build_bottom_leaf_node(&[key_b, key_a], &[val_b, val_a2]);
648    store.set_node(tree_root, index, leaf_node).unwrap();
649
650    let expected_nodes = get_non_empty_nodes(&store);
651    let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
652    actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
653
654    let mut leaves = smt.bottom_leaves();
655    assert_eq!(leaves.next(), Some((leaf_node, vec![(key_b, val_b), (key_a, val_a2)])));
656    assert_eq!(leaves.next(), None);
657}
658
659#[test]
660fn tsmt_bottom_tier_two() {
661    let mut smt = TieredSmt::default();
662    let mut store = MerkleStore::default();
663
664    // --- insert the first value ---------------------------------------------
665    let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
666    let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
667    let val_a = [ONE; WORD_SIZE];
668    smt.insert(key_a, val_a);
669
670    // --- insert the second value --------------------------------------------
671    // the key for this value has the same 48-bit prefix as the key for the first value,
672    // thus, on insertions, both should end up in different nodes at depth 64
673    let raw_b = 0b_10101010_10101010_00011111_11111111_10010110_10010011_01100000_00000000_u64;
674    let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
675    let val_b = [Felt::new(2); WORD_SIZE];
676    smt.insert(key_b, val_b);
677
678    // --- build Merkle store with equivalent data ----------------------------
679    let mut tree_root = get_init_root();
680    let index_a = NodeIndex::make(64, raw_a);
681    let leaf_node_a = build_bottom_leaf_node(&[key_a], &[val_a]);
682    tree_root = store.set_node(tree_root, index_a, leaf_node_a).unwrap().root;
683
684    let index_b = NodeIndex::make(64, raw_b);
685    let leaf_node_b = build_bottom_leaf_node(&[key_b], &[val_b]);
686    tree_root = store.set_node(tree_root, index_b, leaf_node_b).unwrap().root;
687
688    // --- verify that data is consistent between store and tree --------------
689
690    assert_eq!(smt.root(), tree_root);
691
692    assert_eq!(smt.get_value(key_a), val_a);
693    assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
694    let expected_path = store.get_path(tree_root, index_a).unwrap().path;
695    assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
696
697    assert_eq!(smt.get_value(key_b), val_b);
698    assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
699    let expected_path = store.get_path(tree_root, index_b).unwrap().path;
700    assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
701
702    // make sure inner nodes match - the store contains more entries because it keeps track of
703    // all prior state - so, we don't check that the number of inner nodes is the same in both
704    let expected_nodes = get_non_empty_nodes(&store);
705    let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
706    actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
707
708    // make sure leaves are returned correctly
709    let mut leaves = smt.bottom_leaves();
710    assert_eq!(leaves.next(), Some((leaf_node_b, vec![(key_b, val_b)])));
711    assert_eq!(leaves.next(), Some((leaf_node_a, vec![(key_a, val_a)])));
712    assert_eq!(leaves.next(), None);
713}
714
715// GET PROOF TESTS
716// ================================================================================================
717
718#[test]
719fn tsmt_get_proof() {
720    let mut smt = TieredSmt::default();
721
722    // --- insert a value into the tree ---------------------------------------
723    let raw_a = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000000_u64;
724    let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
725    let value_a = [ONE, ONE, ONE, ONE];
726    smt.insert(key_a, value_a);
727
728    // --- insert a value with the same 48-bit prefix into the tree -----------
729    let raw_b = 0b_01010101_01010101_11111111_11111111_10110101_10101010_10111100_00000000_u64;
730    let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
731    let value_b = [ONE, ONE, ONE, ZERO];
732    smt.insert(key_b, value_b);
733
734    let smt_alt = smt.clone();
735
736    // --- insert a value with the same 32-bit prefix into the tree -----------
737    let raw_c = 0b_01010101_01010101_11111111_11111111_11111101_10101010_10111100_00000000_u64;
738    let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
739    let value_c = [ONE, ONE, ZERO, ZERO];
740    smt.insert(key_c, value_c);
741
742    // --- insert a value with the same 64-bit prefix as A into the tree ------
743    let raw_d = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000000_u64;
744    let key_d = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_d)]);
745    let value_d = [ONE, ZERO, ZERO, ZERO];
746    smt.insert(key_d, value_d);
747
748    // at this point the tree looks as follows:
749    // - A and D are located in the same node at depth 64.
750    // - B is located at depth 64 and shares the same 48-bit prefix with A and D.
751    // - C is located at depth 48 and shares the same 32-bit prefix with A, B, and D.
752
753    // --- generate proof for key A and test that it verifies correctly -------
754    let proof = smt.prove(key_a);
755    assert!(proof.verify_membership(&key_a, &value_a, &smt.root()));
756
757    assert!(!proof.verify_membership(&key_a, &value_b, &smt.root()));
758    assert!(!proof.verify_membership(&key_a, &EMPTY_WORD, &smt.root()));
759    assert!(!proof.verify_membership(&key_b, &value_a, &smt.root()));
760    assert!(!proof.verify_membership(&key_a, &value_a, &smt_alt.root()));
761
762    assert_eq!(proof.get(&key_a), Some(value_a));
763    assert_eq!(proof.get(&key_b), None);
764
765    // since A and D are stored in the same node, we should be able to use the proof to verify
766    // membership of D
767    assert!(proof.verify_membership(&key_d, &value_d, &smt.root()));
768    assert_eq!(proof.get(&key_d), Some(value_d));
769
770    // --- generate proof for key B and test that it verifies correctly -------
771    let proof = smt.prove(key_b);
772    assert!(proof.verify_membership(&key_b, &value_b, &smt.root()));
773
774    assert!(!proof.verify_membership(&key_b, &value_a, &smt.root()));
775    assert!(!proof.verify_membership(&key_b, &EMPTY_WORD, &smt.root()));
776    assert!(!proof.verify_membership(&key_a, &value_b, &smt.root()));
777    assert!(!proof.verify_membership(&key_b, &value_b, &smt_alt.root()));
778
779    assert_eq!(proof.get(&key_b), Some(value_b));
780    assert_eq!(proof.get(&key_a), None);
781
782    // --- generate proof for key C and test that it verifies correctly -------
783    let proof = smt.prove(key_c);
784    assert!(proof.verify_membership(&key_c, &value_c, &smt.root()));
785
786    assert!(!proof.verify_membership(&key_c, &value_a, &smt.root()));
787    assert!(!proof.verify_membership(&key_c, &EMPTY_WORD, &smt.root()));
788    assert!(!proof.verify_membership(&key_a, &value_c, &smt.root()));
789    assert!(!proof.verify_membership(&key_c, &value_c, &smt_alt.root()));
790
791    assert_eq!(proof.get(&key_c), Some(value_c));
792    assert_eq!(proof.get(&key_b), None);
793
794    // --- generate proof for key D and test that it verifies correctly -------
795    let proof = smt.prove(key_d);
796    assert!(proof.verify_membership(&key_d, &value_d, &smt.root()));
797
798    assert!(!proof.verify_membership(&key_d, &value_b, &smt.root()));
799    assert!(!proof.verify_membership(&key_d, &EMPTY_WORD, &smt.root()));
800    assert!(!proof.verify_membership(&key_b, &value_d, &smt.root()));
801    assert!(!proof.verify_membership(&key_d, &value_d, &smt_alt.root()));
802
803    assert_eq!(proof.get(&key_d), Some(value_d));
804    assert_eq!(proof.get(&key_b), None);
805
806    // since A and D are stored in the same node, we should be able to use the proof to verify
807    // membership of A
808    assert!(proof.verify_membership(&key_a, &value_a, &smt.root()));
809    assert_eq!(proof.get(&key_a), Some(value_a));
810
811    // --- generate proof for an empty key at depth 64 ------------------------
812    // this key has the same 48-bit prefix as A but is different from B
813    let raw = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000011_u64;
814    let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
815
816    let proof = smt.prove(key);
817    assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root()));
818
819    assert!(!proof.verify_membership(&key, &value_a, &smt.root()));
820    assert!(!proof.verify_membership(&key, &EMPTY_WORD, &smt_alt.root()));
821
822    assert_eq!(proof.get(&key), Some(EMPTY_WORD));
823    assert_eq!(proof.get(&key_b), None);
824
825    // the same proof should verify against any key with the same 64-bit prefix
826    let key2 = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw)]);
827    assert!(proof.verify_membership(&key2, &EMPTY_WORD, &smt.root()));
828    assert_eq!(proof.get(&key2), Some(EMPTY_WORD));
829
830    // but verifying if against a key with the same 63-bit prefix (or smaller) should fail
831    let raw3 = 0b_01010101_01010101_11111111_11111111_10110101_10101010_11111100_00000010_u64;
832    let key3 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw3)]);
833    assert!(!proof.verify_membership(&key3, &EMPTY_WORD, &smt.root()));
834    assert_eq!(proof.get(&key3), None);
835
836    // --- generate proof for an empty key at depth 48 ------------------------
837    // this key has the same 32-prefix as A, B, C, and D, but is different from C
838    let raw = 0b_01010101_01010101_11111111_11111111_00110101_10101010_11111100_00000000_u64;
839    let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
840
841    let proof = smt.prove(key);
842    assert!(proof.verify_membership(&key, &EMPTY_WORD, &smt.root()));
843
844    assert!(!proof.verify_membership(&key, &value_a, &smt.root()));
845    assert!(!proof.verify_membership(&key, &EMPTY_WORD, &smt_alt.root()));
846
847    assert_eq!(proof.get(&key), Some(EMPTY_WORD));
848    assert_eq!(proof.get(&key_b), None);
849
850    // the same proof should verify against any key with the same 48-bit prefix
851    let raw2 = 0b_01010101_01010101_11111111_11111111_00110101_10101010_01111100_00000000_u64;
852    let key2 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw2)]);
853    assert!(proof.verify_membership(&key2, &EMPTY_WORD, &smt.root()));
854    assert_eq!(proof.get(&key2), Some(EMPTY_WORD));
855
856    // but verifying against a key with the same 47-bit prefix (or smaller) should fail
857    let raw3 = 0b_01010101_01010101_11111111_11111111_00110101_10101011_11111100_00000000_u64;
858    let key3 = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw3)]);
859    assert!(!proof.verify_membership(&key3, &EMPTY_WORD, &smt.root()));
860    assert_eq!(proof.get(&key3), None);
861}
862
863// ERROR TESTS
864// ================================================================================================
865
866#[test]
867fn tsmt_node_not_available() {
868    let mut smt = TieredSmt::default();
869
870    let raw = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
871    let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
872    let value = [ONE; WORD_SIZE];
873
874    // build an index which is just below the inserted leaf node
875    let index = NodeIndex::make(17, raw >> 47);
876
877    // since we haven't inserted the node yet, we should be able to get node and path to this index
878    assert!(smt.get_node(index).is_ok());
879    assert!(smt.get_path(index).is_ok());
880
881    smt.insert(key, value);
882
883    // but once the node is inserted, everything under it should be unavailable
884    assert!(smt.get_node(index).is_err());
885    assert!(smt.get_path(index).is_err());
886
887    let index = NodeIndex::make(32, raw >> 32);
888    assert!(smt.get_node(index).is_err());
889    assert!(smt.get_path(index).is_err());
890
891    let index = NodeIndex::make(34, raw >> 30);
892    assert!(smt.get_node(index).is_err());
893    assert!(smt.get_path(index).is_err());
894
895    let index = NodeIndex::make(50, raw >> 14);
896    assert!(smt.get_node(index).is_err());
897    assert!(smt.get_path(index).is_err());
898
899    let index = NodeIndex::make(64, raw);
900    assert!(smt.get_node(index).is_err());
901    assert!(smt.get_path(index).is_err());
902}
903
904// HELPER FUNCTIONS
905// ================================================================================================
906
907fn get_init_root() -> RpoDigest {
908    EmptySubtreeRoots::empty_hashes(64)[0]
909}
910
911fn build_leaf_node(key: RpoDigest, value: Word, depth: u8) -> RpoDigest {
912    Rpo256::merge_in_domain(&[key, value.into()], depth.into())
913}
914
915fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest {
916    assert_eq!(keys.len(), values.len());
917
918    let mut elements = Vec::with_capacity(keys.len());
919    for (key, val) in keys.iter().zip(values.iter()) {
920        elements.extend_from_slice(key.as_elements());
921        elements.extend_from_slice(val.as_slice());
922    }
923
924    Rpo256::hash_elements(&elements)
925}
926
927fn get_non_empty_nodes(store: &MerkleStore) -> Vec<InnerNodeInfo> {
928    store
929        .inner_nodes()
930        .filter(|node| !is_empty_subtree(&node.value))
931        .collect::<Vec<_>>()
932}
933
934fn is_empty_subtree(node: &RpoDigest) -> bool {
935    EmptySubtreeRoots::empty_hashes(255).contains(node)
936}