miden-crypto/src/merkle/tiered_smt/tests.rs
src/merkle/tiered_smt/tests.rsRaw
1use super::{
2 super::{super::ONE, Felt, MerkleStore, WORD_SIZE, ZERO},
3 get_remaining_path, EmptySubtreeRoots, InnerNodeInfo, NodeIndex, Rpo256, RpoDigest, TieredSmt,
4 Vec, Word,
5};
6
7#[test]
8fn tsmt_insert_one() {
9 let mut smt = TieredSmt::default();
10 let mut store = MerkleStore::default();
11
12 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
13 let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
14 let value = [ONE; WORD_SIZE];
15
16 // since the tree is empty, the first node will be inserted at depth 16 and the index will be
17 // 16 most significant bits of the key
18 let index = NodeIndex::make(16, raw >> 48);
19 let leaf_node = build_leaf_node(key, value, 16);
20 let tree_root = store.set_node(smt.root().into(), index, leaf_node.into()).unwrap().root;
21
22 smt.insert(key, value);
23
24 assert_eq!(smt.root(), tree_root.into());
25
26 // make sure the value was inserted, and the node is at the expected index
27 assert_eq!(smt.get_value(key), value);
28 assert_eq!(smt.get_node(index).unwrap(), leaf_node);
29
30 // make sure the paths we get from the store and the tree match
31 let expected_path = store.get_path(tree_root, index).unwrap();
32 assert_eq!(smt.get_path(index).unwrap(), expected_path.path);
33
34 // make sure inner nodes match
35 let expected_nodes = get_non_empty_nodes(&store);
36 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
37 assert_eq!(actual_nodes.len(), expected_nodes.len());
38 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
39
40 // make sure leaves are returned correctly
41 let mut leaves = smt.upper_leaves();
42 assert_eq!(leaves.next(), Some((leaf_node, key, value)));
43 assert_eq!(leaves.next(), None);
44}
45
46#[test]
47fn tsmt_insert_two_16() {
48 let mut smt = TieredSmt::default();
49 let mut store = MerkleStore::default();
50
51 // --- insert the first value ---------------------------------------------
52 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
53 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
54 let val_a = [ONE; WORD_SIZE];
55 smt.insert(key_a, val_a);
56
57 // --- insert the second value --------------------------------------------
58 // the key for this value has the same 16-bit prefix as the key for the first value,
59 // thus, on insertions, both values should be pushed to depth 32 tier
60 let raw_b = 0b_10101010_10101010_10011111_11111111_10010110_10010011_11100000_00000000_u64;
61 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
62 let val_b = [Felt::new(2); WORD_SIZE];
63 smt.insert(key_b, val_b);
64
65 // --- build Merkle store with equivalent data ----------------------------
66 let mut tree_root = get_init_root();
67 let index_a = NodeIndex::make(32, raw_a >> 32);
68 let leaf_node_a = build_leaf_node(key_a, val_a, 32);
69 tree_root = store
70 .set_node(tree_root.into(), index_a, leaf_node_a.into())
71 .unwrap()
72 .root
73 .into();
74
75 let index_b = NodeIndex::make(32, raw_b >> 32);
76 let leaf_node_b = build_leaf_node(key_b, val_b, 32);
77 tree_root = store
78 .set_node(tree_root.into(), index_b, leaf_node_b.into())
79 .unwrap()
80 .root
81 .into();
82
83 // --- verify that data is consistent between store and tree --------------
84
85 assert_eq!(smt.root(), tree_root.into());
86
87 assert_eq!(smt.get_value(key_a), val_a);
88 assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
89 let expected_path = store.get_path(tree_root.into(), index_a).unwrap().path;
90 assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
91
92 assert_eq!(smt.get_value(key_b), val_b);
93 assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
94 let expected_path = store.get_path(tree_root.into(), index_b).unwrap().path;
95 assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
96
97 // make sure inner nodes match - the store contains more entries because it keeps track of
98 // all prior state - so, we don't check that the number of inner nodes is the same in both
99 let expected_nodes = get_non_empty_nodes(&store);
100 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
101 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
102
103 // make sure leaves are returned correctly
104 let mut leaves = smt.upper_leaves();
105 assert_eq!(leaves.next(), Some((leaf_node_a, key_a, val_a)));
106 assert_eq!(leaves.next(), Some((leaf_node_b, key_b, val_b)));
107 assert_eq!(leaves.next(), None);
108}
109
110#[test]
111fn tsmt_insert_two_32() {
112 let mut smt = TieredSmt::default();
113 let mut store = MerkleStore::default();
114
115 // --- insert the first value ---------------------------------------------
116 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
117 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
118 let val_a = [ONE; WORD_SIZE];
119 smt.insert(key_a, val_a);
120
121 // --- insert the second value --------------------------------------------
122 // the key for this value has the same 32-bit prefix as the key for the first value,
123 // thus, on insertions, both values should be pushed to depth 48 tier
124 let raw_b = 0b_10101010_10101010_00011111_11111111_00010110_10010011_11100000_00000000_u64;
125 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
126 let val_b = [Felt::new(2); WORD_SIZE];
127 smt.insert(key_b, val_b);
128
129 // --- build Merkle store with equivalent data ----------------------------
130 let mut tree_root = get_init_root();
131 let index_a = NodeIndex::make(48, raw_a >> 16);
132 let leaf_node_a = build_leaf_node(key_a, val_a, 48);
133 tree_root = store
134 .set_node(tree_root.into(), index_a, leaf_node_a.into())
135 .unwrap()
136 .root
137 .into();
138
139 let index_b = NodeIndex::make(48, raw_b >> 16);
140 let leaf_node_b = build_leaf_node(key_b, val_b, 48);
141 tree_root = store
142 .set_node(tree_root.into(), index_b, leaf_node_b.into())
143 .unwrap()
144 .root
145 .into();
146
147 // --- verify that data is consistent between store and tree --------------
148
149 assert_eq!(smt.root(), tree_root.into());
150
151 assert_eq!(smt.get_value(key_a), val_a);
152 assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
153 let expected_path = store.get_path(tree_root.into(), index_a).unwrap().path;
154 assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
155
156 assert_eq!(smt.get_value(key_b), val_b);
157 assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
158 let expected_path = store.get_path(tree_root.into(), index_b).unwrap().path;
159 assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
160
161 // make sure inner nodes match - the store contains more entries because it keeps track of
162 // all prior state - so, we don't check that the number of inner nodes is the same in both
163 let expected_nodes = get_non_empty_nodes(&store);
164 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
165 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
166}
167
168#[test]
169fn tsmt_insert_three() {
170 let mut smt = TieredSmt::default();
171 let mut store = MerkleStore::default();
172
173 // --- insert the first value ---------------------------------------------
174 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
175 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
176 let val_a = [ONE; WORD_SIZE];
177 smt.insert(key_a, val_a);
178
179 // --- insert the second value --------------------------------------------
180 // the key for this value has the same 16-bit prefix as the key for the first value,
181 // thus, on insertions, both values should be pushed to depth 32 tier
182 let raw_b = 0b_10101010_10101010_10011111_11111111_10010110_10010011_11100000_00000000_u64;
183 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
184 let val_b = [Felt::new(2); WORD_SIZE];
185 smt.insert(key_b, val_b);
186
187 // --- insert the third value ---------------------------------------------
188 // the key for this value has the same 16-bit prefix as the keys for the first two,
189 // values; thus, on insertions, it will be inserted into depth 32 tier, but will not
190 // affect locations of the other two values
191 let raw_c = 0b_10101010_10101010_11011111_11111111_10010110_10010011_11100000_00000000_u64;
192 let key_c = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_c)]);
193 let val_c = [Felt::new(3); WORD_SIZE];
194 smt.insert(key_c, val_c);
195
196 // --- build Merkle store with equivalent data ----------------------------
197 let mut tree_root = get_init_root();
198 let index_a = NodeIndex::make(32, raw_a >> 32);
199 let leaf_node_a = build_leaf_node(key_a, val_a, 32);
200 tree_root = store
201 .set_node(tree_root.into(), index_a, leaf_node_a.into())
202 .unwrap()
203 .root
204 .into();
205
206 let index_b = NodeIndex::make(32, raw_b >> 32);
207 let leaf_node_b = build_leaf_node(key_b, val_b, 32);
208 tree_root = store
209 .set_node(tree_root.into(), index_b, leaf_node_b.into())
210 .unwrap()
211 .root
212 .into();
213
214 let index_c = NodeIndex::make(32, raw_c >> 32);
215 let leaf_node_c = build_leaf_node(key_c, val_c, 32);
216 tree_root = store
217 .set_node(tree_root.into(), index_c, leaf_node_c.into())
218 .unwrap()
219 .root
220 .into();
221
222 // --- verify that data is consistent between store and tree --------------
223
224 assert_eq!(smt.root(), tree_root.into());
225
226 assert_eq!(smt.get_value(key_a), val_a);
227 assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
228 let expected_path = store.get_path(tree_root.into(), index_a).unwrap().path;
229 assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
230
231 assert_eq!(smt.get_value(key_b), val_b);
232 assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
233 let expected_path = store.get_path(tree_root.into(), index_b).unwrap().path;
234 assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
235
236 assert_eq!(smt.get_value(key_c), val_c);
237 assert_eq!(smt.get_node(index_c).unwrap(), leaf_node_c);
238 let expected_path = store.get_path(tree_root.into(), index_c).unwrap().path;
239 assert_eq!(smt.get_path(index_c).unwrap(), expected_path);
240
241 // make sure inner nodes match - the store contains more entries because it keeps track of
242 // all prior state - so, we don't check that the number of inner nodes is the same in both
243 let expected_nodes = get_non_empty_nodes(&store);
244 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
245 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
246}
247
248#[test]
249fn tsmt_update() {
250 let mut smt = TieredSmt::default();
251 let mut store = MerkleStore::default();
252
253 // --- insert a value into the tree ---------------------------------------
254 let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
255 let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
256 let value_a = [ONE; WORD_SIZE];
257 smt.insert(key, value_a);
258
259 // --- update the value ---------------------------------------------------
260 let value_b = [Felt::new(2); WORD_SIZE];
261 smt.insert(key, value_b);
262
263 // --- verify consistency -------------------------------------------------
264 let mut tree_root = get_init_root();
265 let index = NodeIndex::make(16, raw >> 48);
266 let leaf_node = build_leaf_node(key, value_b, 16);
267 tree_root = store.set_node(tree_root.into(), index, leaf_node.into()).unwrap().root.into();
268
269 assert_eq!(smt.root(), tree_root.into());
270
271 assert_eq!(smt.get_value(key), value_b);
272 assert_eq!(smt.get_node(index).unwrap(), leaf_node);
273 let expected_path = store.get_path(tree_root.into(), index).unwrap().path;
274 assert_eq!(smt.get_path(index).unwrap(), expected_path);
275
276 // make sure inner nodes match - the store contains more entries because it keeps track of
277 // all prior state - so, we don't check that the number of inner nodes is the same in both
278 let expected_nodes = get_non_empty_nodes(&store);
279 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
280 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
281}
282
283// BOTTOM TIER TESTS
284// ================================================================================================
285
286#[test]
287fn tsmt_bottom_tier() {
288 let mut smt = TieredSmt::default();
289 let mut store = MerkleStore::default();
290
291 // common prefix for the keys
292 let prefix = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
293
294 // --- insert the first value ---------------------------------------------
295 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(prefix)]);
296 let val_a = [ONE; WORD_SIZE];
297 smt.insert(key_a, val_a);
298
299 // --- insert the second value --------------------------------------------
300 // this key has the same 64-bit prefix and thus both values should end up in the same
301 // node at depth 64
302 let key_b = RpoDigest::from([ZERO, ONE, ONE, Felt::new(prefix)]);
303 let val_b = [Felt::new(2); WORD_SIZE];
304 smt.insert(key_b, val_b);
305
306 // --- build Merkle store with equivalent data ----------------------------
307 let index = NodeIndex::make(64, prefix);
308 // to build bottom leaf we sort by key starting with the least significant element, thus
309 // key_b is smaller than key_a.
310 let leaf_node = build_bottom_leaf_node(&[key_b, key_a], &[val_b, val_a]);
311 let mut tree_root = get_init_root();
312 tree_root = store.set_node(tree_root.into(), index, leaf_node.into()).unwrap().root.into();
313
314 // --- verify that data is consistent between store and tree --------------
315
316 assert_eq!(smt.root(), tree_root.into());
317
318 assert_eq!(smt.get_value(key_a), val_a);
319 assert_eq!(smt.get_value(key_b), val_b);
320
321 assert_eq!(smt.get_node(index).unwrap(), leaf_node);
322 let expected_path = store.get_path(tree_root.into(), index).unwrap().path;
323 assert_eq!(smt.get_path(index).unwrap(), expected_path);
324
325 // make sure inner nodes match - the store contains more entries because it keeps track of
326 // all prior state - so, we don't check that the number of inner nodes is the same in both
327 let expected_nodes = get_non_empty_nodes(&store);
328 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
329 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
330
331 // make sure leaves are returned correctly
332 let mut leaves = smt.bottom_leaves();
333 assert_eq!(leaves.next(), Some((leaf_node, vec![(key_b, val_b), (key_a, val_a)])));
334 assert_eq!(leaves.next(), None);
335}
336
337#[test]
338fn tsmt_bottom_tier_two() {
339 let mut smt = TieredSmt::default();
340 let mut store = MerkleStore::default();
341
342 // --- insert the first value ---------------------------------------------
343 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
344 let key_a = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
345 let val_a = [ONE; WORD_SIZE];
346 smt.insert(key_a, val_a);
347
348 // --- insert the second value --------------------------------------------
349 // the key for this value has the same 48-bit prefix as the key for the first value,
350 // thus, on insertions, both should end up in different nodes at depth 64
351 let raw_b = 0b_10101010_10101010_00011111_11111111_10010110_10010011_01100000_00000000_u64;
352 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
353 let val_b = [Felt::new(2); WORD_SIZE];
354 smt.insert(key_b, val_b);
355
356 // --- build Merkle store with equivalent data ----------------------------
357 let mut tree_root = get_init_root();
358 let index_a = NodeIndex::make(64, raw_a);
359 let leaf_node_a = build_bottom_leaf_node(&[key_a], &[val_a]);
360 tree_root = store
361 .set_node(tree_root.into(), index_a, leaf_node_a.into())
362 .unwrap()
363 .root
364 .into();
365
366 let index_b = NodeIndex::make(64, raw_b);
367 let leaf_node_b = build_bottom_leaf_node(&[key_b], &[val_b]);
368 tree_root = store
369 .set_node(tree_root.into(), index_b, leaf_node_b.into())
370 .unwrap()
371 .root
372 .into();
373
374 // --- verify that data is consistent between store and tree --------------
375
376 assert_eq!(smt.root(), tree_root.into());
377
378 assert_eq!(smt.get_value(key_a), val_a);
379 assert_eq!(smt.get_node(index_a).unwrap(), leaf_node_a);
380 let expected_path = store.get_path(tree_root.into(), index_a).unwrap().path;
381 assert_eq!(smt.get_path(index_a).unwrap(), expected_path);
382
383 assert_eq!(smt.get_value(key_b), val_b);
384 assert_eq!(smt.get_node(index_b).unwrap(), leaf_node_b);
385 let expected_path = store.get_path(tree_root.into(), index_b).unwrap().path;
386 assert_eq!(smt.get_path(index_b).unwrap(), expected_path);
387
388 // make sure inner nodes match - the store contains more entries because it keeps track of
389 // all prior state - so, we don't check that the number of inner nodes is the same in both
390 let expected_nodes = get_non_empty_nodes(&store);
391 let actual_nodes = smt.inner_nodes().collect::<Vec<_>>();
392 actual_nodes.iter().for_each(|node| assert!(expected_nodes.contains(node)));
393
394 // make sure leaves are returned correctly
395 let mut leaves = smt.bottom_leaves();
396 assert_eq!(leaves.next(), Some((leaf_node_b, vec![(key_b, val_b)])));
397 assert_eq!(leaves.next(), Some((leaf_node_a, vec![(key_a, val_a)])));
398 assert_eq!(leaves.next(), None);
399}
400
401// ERROR TESTS
402// ================================================================================================
403
404#[test]
405fn tsmt_node_not_available() {
406 let mut smt = TieredSmt::default();
407
408 let raw = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
409 let key = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
410 let value = [ONE; WORD_SIZE];
411
412 // build an index which is just below the inserted leaf node
413 let index = NodeIndex::make(17, raw >> 47);
414
415 // since we haven't inserted the node yet, we should be able to get node and path to this index
416 assert!(smt.get_node(index).is_ok());
417 assert!(smt.get_path(index).is_ok());
418
419 smt.insert(key, value);
420
421 // but once the node is inserted, everything under it should be unavailable
422 assert!(smt.get_node(index).is_err());
423 assert!(smt.get_path(index).is_err());
424
425 let index = NodeIndex::make(32, raw >> 32);
426 assert!(smt.get_node(index).is_err());
427 assert!(smt.get_path(index).is_err());
428
429 let index = NodeIndex::make(34, raw >> 30);
430 assert!(smt.get_node(index).is_err());
431 assert!(smt.get_path(index).is_err());
432
433 let index = NodeIndex::make(50, raw >> 14);
434 assert!(smt.get_node(index).is_err());
435 assert!(smt.get_path(index).is_err());
436
437 let index = NodeIndex::make(64, raw);
438 assert!(smt.get_node(index).is_err());
439 assert!(smt.get_path(index).is_err());
440}
441
442// HELPER FUNCTIONS
443// ================================================================================================
444
445fn get_init_root() -> Word {
446 EmptySubtreeRoots::empty_hashes(64)[0].into()
447}
448
449fn build_leaf_node(key: RpoDigest, value: Word, depth: u8) -> RpoDigest {
450 let remaining_path = get_remaining_path(key, depth as u32);
451 Rpo256::merge_in_domain(&[remaining_path, value.into()], depth.into())
452}
453
454fn build_bottom_leaf_node(keys: &[RpoDigest], values: &[Word]) -> RpoDigest {
455 assert_eq!(keys.len(), values.len());
456
457 let mut elements = Vec::with_capacity(keys.len());
458 for (key, val) in keys.iter().zip(values.iter()) {
459 let mut key = Word::from(key);
460 key[3] = ZERO;
461 elements.extend_from_slice(&key);
462 elements.extend_from_slice(val.as_slice());
463 }
464
465 Rpo256::hash_elements(&elements)
466}
467
468fn get_non_empty_nodes(store: &MerkleStore) -> Vec<InnerNodeInfo> {
469 store
470 .inner_nodes()
471 .filter(|node| !is_empty_subtree(&RpoDigest::from(node.value)))
472 .collect::<Vec<_>>()
473}
474
475fn is_empty_subtree(node: &RpoDigest) -> bool {
476 EmptySubtreeRoots::empty_hashes(255).contains(&node)
477}