miden-crypto/src/merkle/simple_smt/tests.rs
src/merkle/simple_smt/tests.rsRaw
1use super::{
2 super::{int_to_node, InnerNodeInfo, MerkleError, MerkleTree, RpoDigest, SimpleSmt},
3 NodeIndex, Rpo256, Vec, Word,
4};
5use proptest::prelude::*;
6use rand_utils::prng_array;
7
8const KEYS4: [u64; 4] = [0, 1, 2, 3];
9const KEYS8: [u64; 8] = [0, 1, 2, 3, 4, 5, 6, 7];
10
11const VALUES4: [Word; 4] = [
12 int_to_node(1),
13 int_to_node(2),
14 int_to_node(3),
15 int_to_node(4),
16];
17
18const VALUES8: [Word; 8] = [
19 int_to_node(1),
20 int_to_node(2),
21 int_to_node(3),
22 int_to_node(4),
23 int_to_node(5),
24 int_to_node(6),
25 int_to_node(7),
26 int_to_node(8),
27];
28
29const ZERO_VALUES8: [Word; 8] = [int_to_node(0); 8];
30
31#[test]
32fn build_empty_tree() {
33 let smt = SimpleSmt::new(3).unwrap();
34 let mt = MerkleTree::new(ZERO_VALUES8.to_vec()).unwrap();
35 assert_eq!(mt.root(), smt.root());
36}
37
38#[test]
39fn empty_digests_are_consistent() {
40 let depth = 5;
41 let root = SimpleSmt::new(depth).unwrap().root();
42 let computed: [RpoDigest; 2] = (0..depth).fold([Default::default(); 2], |state, _| {
43 let digest = Rpo256::merge(&state);
44 [digest; 2]
45 });
46
47 assert_eq!(Word::from(computed[0]), root);
48}
49
50#[test]
51fn build_sparse_tree() {
52 let mut smt = SimpleSmt::new(3).unwrap();
53 let mut values = ZERO_VALUES8.to_vec();
54
55 // insert single value
56 let key = 6;
57 let new_node = int_to_node(7);
58 values[key as usize] = new_node;
59 smt.insert_leaf(key, new_node)
60 .expect("Failed to insert leaf");
61 let mt2 = MerkleTree::new(values.clone()).unwrap();
62 assert_eq!(mt2.root(), smt.root());
63 assert_eq!(
64 mt2.get_path(NodeIndex::make(3, 6)).unwrap(),
65 smt.get_path(NodeIndex::make(3, 6)).unwrap()
66 );
67
68 // insert second value at distinct leaf branch
69 let key = 2;
70 let new_node = int_to_node(3);
71 values[key as usize] = new_node;
72 smt.insert_leaf(key, new_node)
73 .expect("Failed to insert leaf");
74 let mt3 = MerkleTree::new(values).unwrap();
75 assert_eq!(mt3.root(), smt.root());
76 assert_eq!(
77 mt3.get_path(NodeIndex::make(3, 2)).unwrap(),
78 smt.get_path(NodeIndex::make(3, 2)).unwrap()
79 );
80}
81
82#[test]
83fn build_full_tree() {
84 let tree = SimpleSmt::new(2)
85 .unwrap()
86 .with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter()))
87 .unwrap();
88
89 let (root, node2, node3) = compute_internal_nodes();
90 assert_eq!(root, tree.root());
91 assert_eq!(node2, tree.get_node(NodeIndex::make(1, 0)).unwrap());
92 assert_eq!(node3, tree.get_node(NodeIndex::make(1, 1)).unwrap());
93}
94
95#[test]
96fn get_values() {
97 let tree = SimpleSmt::new(2)
98 .unwrap()
99 .with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter()))
100 .unwrap();
101
102 // check depth 2
103 assert_eq!(VALUES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap());
104 assert_eq!(VALUES4[1], tree.get_node(NodeIndex::make(2, 1)).unwrap());
105 assert_eq!(VALUES4[2], tree.get_node(NodeIndex::make(2, 2)).unwrap());
106 assert_eq!(VALUES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap());
107}
108
109#[test]
110fn get_path() {
111 let tree = SimpleSmt::new(2)
112 .unwrap()
113 .with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter()))
114 .unwrap();
115
116 let (_, node2, node3) = compute_internal_nodes();
117
118 // check depth 2
119 assert_eq!(
120 vec![VALUES4[1], node3],
121 *tree.get_path(NodeIndex::make(2, 0)).unwrap()
122 );
123 assert_eq!(
124 vec![VALUES4[0], node3],
125 *tree.get_path(NodeIndex::make(2, 1)).unwrap()
126 );
127 assert_eq!(
128 vec![VALUES4[3], node2],
129 *tree.get_path(NodeIndex::make(2, 2)).unwrap()
130 );
131 assert_eq!(
132 vec![VALUES4[2], node2],
133 *tree.get_path(NodeIndex::make(2, 3)).unwrap()
134 );
135
136 // check depth 1
137 assert_eq!(vec![node3], *tree.get_path(NodeIndex::make(1, 0)).unwrap());
138 assert_eq!(vec![node2], *tree.get_path(NodeIndex::make(1, 1)).unwrap());
139}
140
141#[test]
142fn test_parent_node_iterator() -> Result<(), MerkleError> {
143 let tree = SimpleSmt::new(2)
144 .unwrap()
145 .with_leaves(KEYS4.into_iter().zip(VALUES4.into_iter()))
146 .unwrap();
147
148 // check depth 2
149 assert_eq!(VALUES4[0], tree.get_node(NodeIndex::make(2, 0)).unwrap());
150 assert_eq!(VALUES4[1], tree.get_node(NodeIndex::make(2, 1)).unwrap());
151 assert_eq!(VALUES4[2], tree.get_node(NodeIndex::make(2, 2)).unwrap());
152 assert_eq!(VALUES4[3], tree.get_node(NodeIndex::make(2, 3)).unwrap());
153
154 // get parent nodes
155 let root = tree.root();
156 let l1n0 = tree.get_node(NodeIndex::make(1, 0))?;
157 let l1n1 = tree.get_node(NodeIndex::make(1, 1))?;
158 let l2n0 = tree.get_node(NodeIndex::make(2, 0))?;
159 let l2n1 = tree.get_node(NodeIndex::make(2, 1))?;
160 let l2n2 = tree.get_node(NodeIndex::make(2, 2))?;
161 let l2n3 = tree.get_node(NodeIndex::make(2, 3))?;
162
163 let nodes: Vec<InnerNodeInfo> = tree.inner_nodes().collect();
164 let expected = vec![
165 InnerNodeInfo {
166 value: root.into(),
167 left: l1n0.into(),
168 right: l1n1.into(),
169 },
170 InnerNodeInfo {
171 value: l1n0.into(),
172 left: l2n0.into(),
173 right: l2n1.into(),
174 },
175 InnerNodeInfo {
176 value: l1n1.into(),
177 left: l2n2.into(),
178 right: l2n3.into(),
179 },
180 ];
181 assert_eq!(nodes, expected);
182
183 Ok(())
184}
185
186#[test]
187fn update_leaf() {
188 let mut tree = SimpleSmt::new(3)
189 .unwrap()
190 .with_leaves(KEYS8.into_iter().zip(VALUES8.into_iter()))
191 .unwrap();
192
193 // update one value
194 let key = 3;
195 let new_node = int_to_node(9);
196 let mut expected_values = VALUES8.to_vec();
197 expected_values[key] = new_node;
198 let expected_tree = SimpleSmt::new(3)
199 .unwrap()
200 .with_leaves(KEYS8.into_iter().zip(expected_values.clone().into_iter()))
201 .unwrap();
202
203 tree.update_leaf(key as u64, new_node).unwrap();
204 assert_eq!(expected_tree.root, tree.root);
205
206 // update another value
207 let key = 6;
208 let new_node = int_to_node(10);
209 expected_values[key] = new_node;
210 let expected_tree = SimpleSmt::new(3)
211 .unwrap()
212 .with_leaves(KEYS8.into_iter().zip(expected_values.into_iter()))
213 .unwrap();
214
215 tree.update_leaf(key as u64, new_node).unwrap();
216 assert_eq!(expected_tree.root, tree.root);
217}
218
219#[test]
220fn small_tree_opening_is_consistent() {
221 // ____k____
222 // / \
223 // _i_ _j_
224 // / \ / \
225 // e f g h
226 // / \ / \ / \ / \
227 // a b 0 0 c 0 0 d
228
229 let z = Word::from(RpoDigest::default());
230
231 let a = Word::from(Rpo256::merge(&[z.into(); 2]));
232 let b = Word::from(Rpo256::merge(&[a.into(); 2]));
233 let c = Word::from(Rpo256::merge(&[b.into(); 2]));
234 let d = Word::from(Rpo256::merge(&[c.into(); 2]));
235
236 let e = Word::from(Rpo256::merge(&[a.into(), b.into()]));
237 let f = Word::from(Rpo256::merge(&[z.into(), z.into()]));
238 let g = Word::from(Rpo256::merge(&[c.into(), z.into()]));
239 let h = Word::from(Rpo256::merge(&[z.into(), d.into()]));
240
241 let i = Word::from(Rpo256::merge(&[e.into(), f.into()]));
242 let j = Word::from(Rpo256::merge(&[g.into(), h.into()]));
243
244 let k = Word::from(Rpo256::merge(&[i.into(), j.into()]));
245
246 let depth = 3;
247 let entries = vec![(0, a), (1, b), (4, c), (7, d)];
248 let tree = SimpleSmt::new(depth).unwrap().with_leaves(entries).unwrap();
249
250 assert_eq!(tree.root(), Word::from(k));
251
252 let cases: Vec<(u8, u64, Vec<Word>)> = vec![
253 (3, 0, vec![b, f, j]),
254 (3, 1, vec![a, f, j]),
255 (3, 4, vec![z, h, i]),
256 (3, 7, vec![z, g, i]),
257 (2, 0, vec![f, j]),
258 (2, 1, vec![e, j]),
259 (2, 2, vec![h, i]),
260 (2, 3, vec![g, i]),
261 (1, 0, vec![j]),
262 (1, 1, vec![i]),
263 ];
264
265 for (depth, key, path) in cases {
266 let opening = tree.get_path(NodeIndex::make(depth, key)).unwrap();
267
268 assert_eq!(path, *opening);
269 }
270}
271
272proptest! {
273 #[test]
274 fn arbitrary_openings_single_leaf(
275 depth in SimpleSmt::MIN_DEPTH..SimpleSmt::MAX_DEPTH,
276 key in prop::num::u64::ANY,
277 leaf in prop::num::u64::ANY,
278 ) {
279 let mut tree = SimpleSmt::new(depth).unwrap();
280
281 let key = key % (1 << depth as u64);
282 let leaf = int_to_node(leaf);
283
284 tree.insert_leaf(key, leaf.into()).unwrap();
285 tree.get_leaf_path(key).unwrap();
286
287 // traverse to root, fetching all paths
288 for d in 1..depth {
289 let k = key >> (depth - d);
290 tree.get_path(NodeIndex::make(d, k)).unwrap();
291 }
292 }
293
294 #[test]
295 fn arbitrary_openings_multiple_leaves(
296 depth in SimpleSmt::MIN_DEPTH..SimpleSmt::MAX_DEPTH,
297 count in 2u8..10u8,
298 ref seed in any::<[u8; 32]>()
299 ) {
300 let mut tree = SimpleSmt::new(depth).unwrap();
301 let mut seed = *seed;
302 let leaves = (1 << depth) - 1;
303
304 for _ in 0..count {
305 seed = prng_array(seed);
306
307 let mut key = [0u8; 8];
308 let mut leaf = [0u8; 8];
309
310 key.copy_from_slice(&seed[..8]);
311 leaf.copy_from_slice(&seed[8..16]);
312
313 let key = u64::from_le_bytes(key);
314 let key = key % leaves;
315 let leaf = u64::from_le_bytes(leaf);
316 let leaf = int_to_node(leaf);
317
318 tree.insert_leaf(key, leaf).unwrap();
319 tree.get_leaf_path(key).unwrap();
320 }
321 }
322}
323
324// HELPER FUNCTIONS
325// --------------------------------------------------------------------------------------------
326
327fn compute_internal_nodes() -> (Word, Word, Word) {
328 let node2 = Rpo256::hash_elements(&[VALUES4[0], VALUES4[1]].concat());
329 let node3 = Rpo256::hash_elements(&[VALUES4[2], VALUES4[3]].concat());
330 let root = Rpo256::merge(&[node2, node3]);
331
332 (root.into(), node2.into(), node3.into())
333}