miden-crypto/src/merkle/smt/tests.rs
src/merkle/smt/tests.rsRaw
1use alloc::{collections::BTreeMap, vec::Vec};
2
3use super::{
4 build_subtree, InnerNode, LeafIndex, NodeIndex, PairComputations, SmtLeaf, SparseMerkleTree,
5 SubtreeLeaf, SubtreeLeavesIter, COLS_PER_SUBTREE, SUBTREE_DEPTH,
6};
7use crate::{
8 hash::rpo::RpoDigest,
9 merkle::{Smt, SMT_DEPTH},
10 Felt, Word, ONE,
11};
12
13fn smtleaf_to_subtree_leaf(leaf: &SmtLeaf) -> SubtreeLeaf {
14 SubtreeLeaf {
15 col: leaf.index().index.value(),
16 hash: leaf.hash(),
17 }
18}
19
20#[test]
21fn test_sorted_pairs_to_leaves() {
22 let entries: Vec<(RpoDigest, Word)> = vec![
23 // Subtree 0.
24 (RpoDigest::new([ONE, ONE, ONE, Felt::new(16)]), [ONE; 4]),
25 (RpoDigest::new([ONE, ONE, ONE, Felt::new(17)]), [ONE; 4]),
26 // Leaf index collision.
27 (RpoDigest::new([ONE, ONE, Felt::new(10), Felt::new(20)]), [ONE; 4]),
28 (RpoDigest::new([ONE, ONE, Felt::new(20), Felt::new(20)]), [ONE; 4]),
29 // Subtree 1. Normal single leaf again.
30 (RpoDigest::new([ONE, ONE, ONE, Felt::new(400)]), [ONE; 4]), // Subtree boundary.
31 (RpoDigest::new([ONE, ONE, ONE, Felt::new(401)]), [ONE; 4]),
32 // Subtree 2. Another normal leaf.
33 (RpoDigest::new([ONE, ONE, ONE, Felt::new(1024)]), [ONE; 4]),
34 ];
35
36 let control = Smt::with_entries_sequential(entries.clone()).unwrap();
37
38 let control_leaves: Vec<SmtLeaf> = {
39 let mut entries_iter = entries.iter().cloned();
40 let mut next_entry = || entries_iter.next().unwrap();
41 let control_leaves = vec![
42 // Subtree 0.
43 SmtLeaf::Single(next_entry()),
44 SmtLeaf::Single(next_entry()),
45 SmtLeaf::new_multiple(vec![next_entry(), next_entry()]).unwrap(),
46 // Subtree 1.
47 SmtLeaf::Single(next_entry()),
48 SmtLeaf::Single(next_entry()),
49 // Subtree 2.
50 SmtLeaf::Single(next_entry()),
51 ];
52 assert_eq!(entries_iter.next(), None);
53 control_leaves
54 };
55
56 let control_subtree_leaves: Vec<Vec<SubtreeLeaf>> = {
57 let mut control_leaves_iter = control_leaves.iter();
58 let mut next_leaf = || control_leaves_iter.next().unwrap();
59
60 let control_subtree_leaves: Vec<Vec<SubtreeLeaf>> = [
61 // Subtree 0.
62 vec![next_leaf(), next_leaf(), next_leaf()],
63 // Subtree 1.
64 vec![next_leaf(), next_leaf()],
65 // Subtree 2.
66 vec![next_leaf()],
67 ]
68 .map(|subtree| subtree.into_iter().map(smtleaf_to_subtree_leaf).collect())
69 .to_vec();
70 assert_eq!(control_leaves_iter.next(), None);
71 control_subtree_leaves
72 };
73
74 let subtrees: PairComputations<u64, SmtLeaf> = Smt::sorted_pairs_to_leaves(entries);
75 // This will check that the hashes, columns, and subtree assignments all match.
76 assert_eq!(subtrees.leaves, control_subtree_leaves);
77
78 // Flattening and re-separating out the leaves into subtrees should have the same result.
79 let mut all_leaves: Vec<SubtreeLeaf> = subtrees.leaves.clone().into_iter().flatten().collect();
80 let re_grouped: Vec<Vec<_>> = SubtreeLeavesIter::from_leaves(&mut all_leaves).collect();
81 assert_eq!(subtrees.leaves, re_grouped);
82
83 // Then finally we might as well check the computed leaf nodes too.
84 let control_leaves: BTreeMap<u64, SmtLeaf> = control
85 .leaves()
86 .map(|(index, value)| (index.index.value(), value.clone()))
87 .collect();
88
89 for (column, test_leaf) in subtrees.nodes {
90 if test_leaf.is_empty() {
91 continue;
92 }
93 let control_leaf = control_leaves
94 .get(&column)
95 .unwrap_or_else(|| panic!("no leaf node found for column {column}"));
96 assert_eq!(control_leaf, &test_leaf);
97 }
98}
99
100// Helper for the below tests.
101fn generate_entries(pair_count: u64) -> Vec<(RpoDigest, Word)> {
102 (0..pair_count)
103 .map(|i| {
104 let leaf_index = ((i as f64 / pair_count as f64) * (pair_count as f64)) as u64;
105 let key = RpoDigest::new([ONE, ONE, Felt::new(i), Felt::new(leaf_index)]);
106 let value = [ONE, ONE, ONE, Felt::new(i)];
107 (key, value)
108 })
109 .collect()
110}
111
112#[test]
113fn test_single_subtree() {
114 // A single subtree's worth of leaves.
115 const PAIR_COUNT: u64 = COLS_PER_SUBTREE;
116
117 let entries = generate_entries(PAIR_COUNT);
118
119 let control = Smt::with_entries_sequential(entries.clone()).unwrap();
120
121 // `entries` should already be sorted by nature of how we constructed it.
122 let leaves = Smt::sorted_pairs_to_leaves(entries).leaves;
123 let leaves = leaves.into_iter().next().unwrap();
124
125 let (first_subtree, subtree_root) = build_subtree(leaves, SMT_DEPTH, SMT_DEPTH);
126 assert!(!first_subtree.is_empty());
127
128 // The inner nodes computed from that subtree should match the nodes in our control tree.
129 for (index, node) in first_subtree.into_iter() {
130 let control = control.get_inner_node(index);
131 assert_eq!(
132 control, node,
133 "subtree-computed node at index {index:?} does not match control",
134 );
135 }
136
137 // The root returned should also match the equivalent node in the control tree.
138 let control_root_index =
139 NodeIndex::new(SMT_DEPTH - SUBTREE_DEPTH, subtree_root.col).expect("Valid root index");
140 let control_root_node = control.get_inner_node(control_root_index);
141 let control_hash = control_root_node.hash();
142 assert_eq!(
143 control_hash, subtree_root.hash,
144 "Subtree-computed root at index {control_root_index:?} does not match control"
145 );
146}
147
148// Test that not just can we compute a subtree correctly, but we can feed the results of one
149// subtree into computing another. In other words, test that `build_subtree()` is correctly
150// composable.
151#[test]
152fn test_two_subtrees() {
153 // Two subtrees' worth of leaves.
154 const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 2;
155
156 let entries = generate_entries(PAIR_COUNT);
157
158 let control = Smt::with_entries_sequential(entries.clone()).unwrap();
159
160 let PairComputations { leaves, .. } = Smt::sorted_pairs_to_leaves(entries);
161 // With two subtrees' worth of leaves, we should have exactly two subtrees.
162 let [first, second]: [Vec<_>; 2] = leaves.try_into().unwrap();
163 assert_eq!(first.len() as u64, PAIR_COUNT / 2);
164 assert_eq!(first.len(), second.len());
165
166 let mut current_depth = SMT_DEPTH;
167 let mut next_leaves: Vec<SubtreeLeaf> = Default::default();
168
169 let (first_nodes, first_root) = build_subtree(first, SMT_DEPTH, current_depth);
170 next_leaves.push(first_root);
171
172 let (second_nodes, second_root) = build_subtree(second, SMT_DEPTH, current_depth);
173 next_leaves.push(second_root);
174
175 // All new inner nodes + the new subtree-leaves should be 512, for one depth-cycle.
176 let total_computed = first_nodes.len() + second_nodes.len() + next_leaves.len();
177 assert_eq!(total_computed as u64, PAIR_COUNT);
178
179 // Verify the computed nodes of both subtrees.
180 let computed_nodes = first_nodes.clone().into_iter().chain(second_nodes);
181 for (index, test_node) in computed_nodes {
182 let control_node = control.get_inner_node(index);
183 assert_eq!(
184 control_node, test_node,
185 "subtree-computed node at index {index:?} does not match control",
186 );
187 }
188
189 current_depth -= SUBTREE_DEPTH;
190
191 let (nodes, root_leaf) = build_subtree(next_leaves, SMT_DEPTH, current_depth);
192 assert_eq!(nodes.len(), SUBTREE_DEPTH as usize);
193 assert_eq!(root_leaf.col, 0);
194
195 for (index, test_node) in nodes {
196 let control_node = control.get_inner_node(index);
197 assert_eq!(
198 control_node, test_node,
199 "subtree-computed node at index {index:?} does not match control",
200 );
201 }
202
203 let index = NodeIndex::new(current_depth - SUBTREE_DEPTH, root_leaf.col).unwrap();
204 let control_root = control.get_inner_node(index).hash();
205 assert_eq!(control_root, root_leaf.hash, "Root mismatch");
206}
207
208#[test]
209fn test_singlethreaded_subtrees() {
210 const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 64;
211
212 let entries = generate_entries(PAIR_COUNT);
213
214 let control = Smt::with_entries_sequential(entries.clone()).unwrap();
215
216 let mut accumulated_nodes: BTreeMap<NodeIndex, InnerNode> = Default::default();
217
218 let PairComputations {
219 leaves: mut leaf_subtrees,
220 nodes: test_leaves,
221 } = Smt::sorted_pairs_to_leaves(entries);
222
223 for current_depth in (SUBTREE_DEPTH..=SMT_DEPTH).step_by(SUBTREE_DEPTH as usize).rev() {
224 // There's no flat_map_unzip(), so this is the best we can do.
225 let (nodes, mut subtree_roots): (Vec<BTreeMap<_, _>>, Vec<SubtreeLeaf>) = leaf_subtrees
226 .into_iter()
227 .enumerate()
228 .map(|(i, subtree)| {
229 // Pre-assertions.
230 assert!(
231 subtree.is_sorted(),
232 "subtree {i} at bottom-depth {current_depth} is not sorted",
233 );
234 assert!(
235 !subtree.is_empty(),
236 "subtree {i} at bottom-depth {current_depth} is empty!",
237 );
238
239 // Do actual things.
240 let (nodes, subtree_root) = build_subtree(subtree, SMT_DEPTH, current_depth);
241
242 // Post-assertions.
243 for (&index, test_node) in nodes.iter() {
244 let control_node = control.get_inner_node(index);
245 assert_eq!(
246 test_node, &control_node,
247 "depth {} subtree {}: test node does not match control at index {:?}",
248 current_depth, i, index,
249 );
250 }
251
252 (nodes, subtree_root)
253 })
254 .unzip();
255
256 // Update state between each depth iteration.
257
258 leaf_subtrees = SubtreeLeavesIter::from_leaves(&mut subtree_roots).collect();
259 accumulated_nodes.extend(nodes.into_iter().flatten());
260
261 assert!(!leaf_subtrees.is_empty(), "on depth {current_depth}");
262 }
263
264 // Make sure the true leaves match, first checking length and then checking each individual
265 // leaf.
266 let control_leaves: BTreeMap<_, _> = control.leaves().collect();
267 let control_leaves_len = control_leaves.len();
268 let test_leaves_len = test_leaves.len();
269 assert_eq!(test_leaves_len, control_leaves_len);
270 for (col, ref test_leaf) in test_leaves {
271 let index = LeafIndex::new_max_depth(col);
272 let &control_leaf = control_leaves.get(&index).unwrap();
273 assert_eq!(test_leaf, control_leaf, "test leaf at column {col} does not match control");
274 }
275
276 // Make sure the inner nodes match, checking length first and then each individual leaf.
277 let control_nodes_len = control.inner_nodes().count();
278 let test_nodes_len = accumulated_nodes.len();
279 assert_eq!(test_nodes_len, control_nodes_len);
280 for (index, test_node) in accumulated_nodes.clone() {
281 let control_node = control.get_inner_node(index);
282 assert_eq!(test_node, control_node, "test node does not match control at {index:?}");
283 }
284
285 // After the last iteration of the above for loop, we should have the new root node actually
286 // in two places: one in `accumulated_nodes`, and the other as the "next leaves" return from
287 // `build_subtree()`. So let's check both!
288
289 let control_root = control.get_inner_node(NodeIndex::root());
290
291 // That for loop should have left us with only one leaf subtree...
292 let [leaf_subtree]: [Vec<_>; 1] = leaf_subtrees.try_into().unwrap();
293 // which itself contains only one 'leaf'...
294 let [root_leaf]: [SubtreeLeaf; 1] = leaf_subtree.try_into().unwrap();
295 // which matches the expected root.
296 assert_eq!(control.root(), root_leaf.hash);
297
298 // Likewise `accumulated_nodes` should contain a node at the root index...
299 assert!(accumulated_nodes.contains_key(&NodeIndex::root()));
300 // and it should match our actual root.
301 let test_root = accumulated_nodes.get(&NodeIndex::root()).unwrap();
302 assert_eq!(control_root, *test_root);
303 // And of course the root we got from each place should match.
304 assert_eq!(control.root(), root_leaf.hash);
305}
306
307/// The parallel version of `test_singlethreaded_subtree()`.
308#[test]
309#[cfg(feature = "concurrent")]
310fn test_multithreaded_subtrees() {
311 use rayon::prelude::*;
312
313 const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 64;
314
315 let entries = generate_entries(PAIR_COUNT);
316
317 let control = Smt::with_entries_sequential(entries.clone()).unwrap();
318
319 let mut accumulated_nodes: BTreeMap<NodeIndex, InnerNode> = Default::default();
320
321 let PairComputations {
322 leaves: mut leaf_subtrees,
323 nodes: test_leaves,
324 } = Smt::sorted_pairs_to_leaves(entries);
325
326 for current_depth in (SUBTREE_DEPTH..=SMT_DEPTH).step_by(SUBTREE_DEPTH as usize).rev() {
327 let (nodes, mut subtree_roots): (Vec<BTreeMap<_, _>>, Vec<SubtreeLeaf>) = leaf_subtrees
328 .into_par_iter()
329 .enumerate()
330 .map(|(i, subtree)| {
331 // Pre-assertions.
332 assert!(
333 subtree.is_sorted(),
334 "subtree {i} at bottom-depth {current_depth} is not sorted",
335 );
336 assert!(
337 !subtree.is_empty(),
338 "subtree {i} at bottom-depth {current_depth} is empty!",
339 );
340
341 let (nodes, subtree_root) = build_subtree(subtree, SMT_DEPTH, current_depth);
342
343 // Post-assertions.
344 for (&index, test_node) in nodes.iter() {
345 let control_node = control.get_inner_node(index);
346 assert_eq!(
347 test_node, &control_node,
348 "depth {} subtree {}: test node does not match control at index {:?}",
349 current_depth, i, index,
350 );
351 }
352
353 (nodes, subtree_root)
354 })
355 .unzip();
356
357 leaf_subtrees = SubtreeLeavesIter::from_leaves(&mut subtree_roots).collect();
358 accumulated_nodes.extend(nodes.into_iter().flatten());
359
360 assert!(!leaf_subtrees.is_empty(), "on depth {current_depth}");
361 }
362
363 // Make sure the true leaves match, checking length first and then each individual leaf.
364 let control_leaves: BTreeMap<_, _> = control.leaves().collect();
365 let control_leaves_len = control_leaves.len();
366 let test_leaves_len = test_leaves.len();
367 assert_eq!(test_leaves_len, control_leaves_len);
368 for (col, ref test_leaf) in test_leaves {
369 let index = LeafIndex::new_max_depth(col);
370 let &control_leaf = control_leaves.get(&index).unwrap();
371 assert_eq!(test_leaf, control_leaf);
372 }
373
374 // Make sure the inner nodes match, checking length first and then each individual leaf.
375 let control_nodes_len = control.inner_nodes().count();
376 let test_nodes_len = accumulated_nodes.len();
377 assert_eq!(test_nodes_len, control_nodes_len);
378 for (index, test_node) in accumulated_nodes.clone() {
379 let control_node = control.get_inner_node(index);
380 assert_eq!(test_node, control_node, "test node does not match control at {index:?}");
381 }
382
383 // After the last iteration of the above for loop, we should have the new root node actually
384 // in two places: one in `accumulated_nodes`, and the other as the "next leaves" return from
385 // `build_subtree()`. So let's check both!
386
387 let control_root = control.get_inner_node(NodeIndex::root());
388
389 // That for loop should have left us with only one leaf subtree...
390 let [leaf_subtree]: [_; 1] = leaf_subtrees.try_into().unwrap();
391 // which itself contains only one 'leaf'...
392 let [root_leaf]: [_; 1] = leaf_subtree.try_into().unwrap();
393 // which matches the expected root.
394 assert_eq!(control.root(), root_leaf.hash);
395
396 // Likewise `accumulated_nodes` should contain a node at the root index...
397 assert!(accumulated_nodes.contains_key(&NodeIndex::root()));
398 // and it should match our actual root.
399 let test_root = accumulated_nodes.get(&NodeIndex::root()).unwrap();
400 assert_eq!(control_root, *test_root);
401 // And of course the root we got from each place should match.
402 assert_eq!(control.root(), root_leaf.hash);
403}
404
405#[test]
406#[cfg(feature = "concurrent")]
407fn test_with_entries_parallel() {
408 const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 64;
409
410 let entries = generate_entries(PAIR_COUNT);
411
412 let control = Smt::with_entries_sequential(entries.clone()).unwrap();
413
414 let smt = Smt::with_entries(entries.clone()).unwrap();
415 assert_eq!(smt.root(), control.root());
416 assert_eq!(smt, control);
417}