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