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