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