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