arnaucube
arnaucube/miden-cryptoPublic Self-hosted repository
miden-crypto/src/merkle/smt/full/tests.rs
src/merkle/smt/full/tests.rsRaw
  1use super::{Felt, LeafIndex, NodeIndex, Rpo256, RpoDigest, Smt, SmtLeaf, EMPTY_WORD, SMT_DEPTH};
  2use crate::{
  3    merkle::{EmptySubtreeRoots, MerkleStore},
  4    utils::{collections::*, Deserializable, Serializable},
  5    Word, ONE, WORD_SIZE,
  6};
  7
  8// SMT
  9// --------------------------------------------------------------------------------------------
 10
 11/// This test checks that inserting twice at the same key functions as expected. The test covers
 12/// only the case where the key is alone in its leaf
 13#[test]
 14fn test_smt_insert_at_same_key() {
 15    let mut smt = Smt::default();
 16    let mut store: MerkleStore = MerkleStore::default();
 17
 18    assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
 19
 20    let key_1: RpoDigest = {
 21        let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
 22
 23        RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
 24    };
 25    let key_1_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key_1).into();
 26
 27    let value_1 = [ONE; WORD_SIZE];
 28    let value_2 = [ONE + ONE; WORD_SIZE];
 29
 30    // Insert value 1 and ensure root is as expected
 31    {
 32        let leaf_node = build_empty_or_single_leaf_node(key_1, value_1);
 33        let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
 34
 35        let old_value_1 = smt.insert(key_1, value_1);
 36        assert_eq!(old_value_1, EMPTY_WORD);
 37
 38        assert_eq!(smt.root(), tree_root);
 39    }
 40
 41    // Insert value 2 and ensure root is as expected
 42    {
 43        let leaf_node = build_empty_or_single_leaf_node(key_1, value_2);
 44        let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
 45
 46        let old_value_2 = smt.insert(key_1, value_2);
 47        assert_eq!(old_value_2, value_1);
 48
 49        assert_eq!(smt.root(), tree_root);
 50    }
 51}
 52
 53/// This test checks that inserting twice at the same key functions as expected. The test covers
 54/// only the case where the leaf type is `SmtLeaf::Multiple`
 55#[test]
 56fn test_smt_insert_at_same_key_2() {
 57    // The most significant u64 used for both keys (to ensure they map to the same leaf)
 58    let key_msb: u64 = 42;
 59
 60    let key_already_present: RpoDigest =
 61        RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(key_msb)]);
 62    let key_already_present_index: NodeIndex =
 63        LeafIndex::<SMT_DEPTH>::from(key_already_present).into();
 64    let value_already_present = [ONE + ONE + ONE; WORD_SIZE];
 65
 66    let mut smt =
 67        Smt::with_entries(core::iter::once((key_already_present, value_already_present))).unwrap();
 68    let mut store: MerkleStore = {
 69        let mut store = MerkleStore::default();
 70
 71        let leaf_node = build_empty_or_single_leaf_node(key_already_present, value_already_present);
 72        store
 73            .set_node(*EmptySubtreeRoots::entry(SMT_DEPTH, 0), key_already_present_index, leaf_node)
 74            .unwrap();
 75        store
 76    };
 77
 78    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(key_msb)]);
 79    let key_1_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key_1).into();
 80
 81    assert_eq!(key_1_index, key_already_present_index);
 82
 83    let value_1 = [ONE; WORD_SIZE];
 84    let value_2 = [ONE + ONE; WORD_SIZE];
 85
 86    // Insert value 1 and ensure root is as expected
 87    {
 88        // Note: key_1 comes first because it is smaller
 89        let leaf_node = build_multiple_leaf_node(&[
 90            (key_1, value_1),
 91            (key_already_present, value_already_present),
 92        ]);
 93        let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
 94
 95        let old_value_1 = smt.insert(key_1, value_1);
 96        assert_eq!(old_value_1, EMPTY_WORD);
 97
 98        assert_eq!(smt.root(), tree_root);
 99    }
100
101    // Insert value 2 and ensure root is as expected
102    {
103        let leaf_node = build_multiple_leaf_node(&[
104            (key_1, value_2),
105            (key_already_present, value_already_present),
106        ]);
107        let tree_root = store.set_node(smt.root(), key_1_index, leaf_node).unwrap().root;
108
109        let old_value_2 = smt.insert(key_1, value_2);
110        assert_eq!(old_value_2, value_1);
111
112        assert_eq!(smt.root(), tree_root);
113    }
114}
115
116/// This test ensures that the root of the tree is as expected when we add/remove 3 items at 3
117/// different keys. This also tests that the merkle paths produced are as expected.
118#[test]
119fn test_smt_insert_and_remove_multiple_values() {
120    fn insert_values_and_assert_path(
121        smt: &mut Smt,
122        store: &mut MerkleStore,
123        key_values: &[(RpoDigest, Word)],
124    ) {
125        for &(key, value) in key_values {
126            let key_index: NodeIndex = LeafIndex::<SMT_DEPTH>::from(key).into();
127
128            let leaf_node = build_empty_or_single_leaf_node(key, value);
129            let tree_root = store.set_node(smt.root(), key_index, leaf_node).unwrap().root;
130
131            let _ = smt.insert(key, value);
132
133            assert_eq!(smt.root(), tree_root);
134
135            let expected_path = store.get_path(tree_root, key_index).unwrap();
136            assert_eq!(smt.open(&key).into_parts().0, expected_path.path);
137        }
138    }
139    let mut smt = Smt::default();
140    let mut store: MerkleStore = MerkleStore::default();
141
142    assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
143
144    let key_1: RpoDigest = {
145        let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
146
147        RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
148    };
149
150    let key_2: RpoDigest = {
151        let raw = 0b_11111111_11111111_11111111_11111111_11111111_11111111_11111111_11111111_u64;
152
153        RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
154    };
155
156    let key_3: RpoDigest = {
157        let raw = 0b_00000000_00000000_00000000_00000000_00000000_00000000_00000000_00000000_u64;
158
159        RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)])
160    };
161
162    let value_1 = [ONE; WORD_SIZE];
163    let value_2 = [ONE + ONE; WORD_SIZE];
164    let value_3 = [ONE + ONE + ONE; WORD_SIZE];
165
166    // Insert values in the tree
167    let key_values = [(key_1, value_1), (key_2, value_2), (key_3, value_3)];
168    insert_values_and_assert_path(&mut smt, &mut store, &key_values);
169
170    // Remove values from the tree
171    let key_empty_values = [(key_1, EMPTY_WORD), (key_2, EMPTY_WORD), (key_3, EMPTY_WORD)];
172    insert_values_and_assert_path(&mut smt, &mut store, &key_empty_values);
173
174    let empty_root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
175    assert_eq!(smt.root(), empty_root);
176
177    // an empty tree should have no leaves or inner nodes
178    assert!(smt.leaves.is_empty());
179    assert!(smt.inner_nodes.is_empty());
180}
181
182/// This tests that inserting the empty value does indeed remove the key-value contained at the
183/// leaf. We insert & remove 3 values at the same leaf to ensure that all cases are covered (empty,
184/// single, multiple).
185#[test]
186fn test_smt_removal() {
187    let mut smt = Smt::default();
188
189    let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
190
191    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
192    let key_2: RpoDigest =
193        RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
194    let key_3: RpoDigest =
195        RpoDigest::from([3_u32.into(), 3_u32.into(), 3_u32.into(), Felt::new(raw)]);
196
197    let value_1 = [ONE; WORD_SIZE];
198    let value_2 = [2_u32.into(); WORD_SIZE];
199    let value_3: [Felt; 4] = [3_u32.into(); WORD_SIZE];
200
201    // insert key-value 1
202    {
203        let old_value_1 = smt.insert(key_1, value_1);
204        assert_eq!(old_value_1, EMPTY_WORD);
205
206        assert_eq!(smt.get_leaf(&key_1), SmtLeaf::Single((key_1, value_1)));
207    }
208
209    // insert key-value 2
210    {
211        let old_value_2 = smt.insert(key_2, value_2);
212        assert_eq!(old_value_2, EMPTY_WORD);
213
214        assert_eq!(
215            smt.get_leaf(&key_2),
216            SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)])
217        );
218    }
219
220    // insert key-value 3
221    {
222        let old_value_3 = smt.insert(key_3, value_3);
223        assert_eq!(old_value_3, EMPTY_WORD);
224
225        assert_eq!(
226            smt.get_leaf(&key_3),
227            SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2), (key_3, value_3)])
228        );
229    }
230
231    // remove key 3
232    {
233        let old_value_3 = smt.insert(key_3, EMPTY_WORD);
234        assert_eq!(old_value_3, value_3);
235
236        assert_eq!(
237            smt.get_leaf(&key_3),
238            SmtLeaf::Multiple(vec![(key_1, value_1), (key_2, value_2)])
239        );
240    }
241
242    // remove key 2
243    {
244        let old_value_2 = smt.insert(key_2, EMPTY_WORD);
245        assert_eq!(old_value_2, value_2);
246
247        assert_eq!(smt.get_leaf(&key_2), SmtLeaf::Single((key_1, value_1)));
248    }
249
250    // remove key 1
251    {
252        let old_value_1 = smt.insert(key_1, EMPTY_WORD);
253        assert_eq!(old_value_1, value_1);
254
255        assert_eq!(smt.get_leaf(&key_1), SmtLeaf::new_empty(key_1.into()));
256    }
257}
258
259/// Tests that 2 key-value pairs stored in the same leaf have the same path
260#[test]
261fn test_smt_path_to_keys_in_same_leaf_are_equal() {
262    let raw = 0b_01101001_01101100_00011111_11111111_10010110_10010011_11100000_00000000_u64;
263
264    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw)]);
265    let key_2: RpoDigest =
266        RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), Felt::new(raw)]);
267
268    let value_1 = [ONE; WORD_SIZE];
269    let value_2 = [2_u32.into(); WORD_SIZE];
270
271    let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
272
273    assert_eq!(smt.open(&key_1), smt.open(&key_2));
274}
275
276/// Tests that an empty leaf hashes to the empty word
277#[test]
278fn test_empty_leaf_hash() {
279    let smt = Smt::default();
280
281    let leaf = smt.get_leaf(&RpoDigest::default());
282    assert_eq!(leaf.hash(), EMPTY_WORD.into());
283}
284
285/// Tests that `get_value()` works as expected
286#[test]
287fn test_smt_get_value() {
288    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, ONE]);
289    let key_2: RpoDigest =
290        RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), 2_u32.into()]);
291
292    let value_1 = [ONE; WORD_SIZE];
293    let value_2 = [2_u32.into(); WORD_SIZE];
294
295    let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
296
297    let returned_value_1 = smt.get_value(&key_1);
298    let returned_value_2 = smt.get_value(&key_2);
299
300    assert_eq!(value_1, returned_value_1);
301    assert_eq!(value_2, returned_value_2);
302
303    // Check that a key with no inserted value returns the empty word
304    let key_no_value =
305        RpoDigest::from([42_u32.into(), 42_u32.into(), 42_u32.into(), 42_u32.into()]);
306
307    assert_eq!(EMPTY_WORD, smt.get_value(&key_no_value));
308}
309
310/// Tests that `entries()` works as expected
311#[test]
312fn test_smt_entries() {
313    let key_1: RpoDigest = RpoDigest::from([ONE, ONE, ONE, ONE]);
314    let key_2: RpoDigest =
315        RpoDigest::from([2_u32.into(), 2_u32.into(), 2_u32.into(), 2_u32.into()]);
316
317    let value_1 = [ONE; WORD_SIZE];
318    let value_2 = [2_u32.into(); WORD_SIZE];
319
320    let smt = Smt::with_entries([(key_1, value_1), (key_2, value_2)]).unwrap();
321
322    let mut entries = smt.entries();
323
324    // Note: for simplicity, we assume the order `(k1,v1), (k2,v2)`. If a new implementation
325    // switches the order, it is OK to modify the order here as well.
326    assert_eq!(&(key_1, value_1), entries.next().unwrap());
327    assert_eq!(&(key_2, value_2), entries.next().unwrap());
328    assert!(entries.next().is_none());
329}
330
331// SMT LEAF
332// --------------------------------------------------------------------------------------------
333
334#[test]
335fn test_empty_smt_leaf_serialization() {
336    let empty_leaf = SmtLeaf::new_empty(LeafIndex::new_max_depth(42));
337
338    let mut serialized = empty_leaf.to_bytes();
339    // extend buffer with random bytes
340    serialized.extend([1, 2, 3, 4, 5]);
341    let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
342
343    assert_eq!(empty_leaf, deserialized);
344}
345
346#[test]
347fn test_single_smt_leaf_serialization() {
348    let single_leaf = SmtLeaf::new_single(
349        RpoDigest::from([10_u32.into(), 11_u32.into(), 12_u32.into(), 13_u32.into()]),
350        [1_u32.into(), 2_u32.into(), 3_u32.into(), 4_u32.into()],
351    );
352
353    let mut serialized = single_leaf.to_bytes();
354    // extend buffer with random bytes
355    serialized.extend([1, 2, 3, 4, 5]);
356    let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
357
358    assert_eq!(single_leaf, deserialized);
359}
360
361#[test]
362fn test_multiple_smt_leaf_serialization_success() {
363    let multiple_leaf = SmtLeaf::new_multiple(vec![
364        (
365            RpoDigest::from([10_u32.into(), 11_u32.into(), 12_u32.into(), 13_u32.into()]),
366            [1_u32.into(), 2_u32.into(), 3_u32.into(), 4_u32.into()],
367        ),
368        (
369            RpoDigest::from([100_u32.into(), 101_u32.into(), 102_u32.into(), 13_u32.into()]),
370            [11_u32.into(), 12_u32.into(), 13_u32.into(), 14_u32.into()],
371        ),
372    ])
373    .unwrap();
374
375    let mut serialized = multiple_leaf.to_bytes();
376    // extend buffer with random bytes
377    serialized.extend([1, 2, 3, 4, 5]);
378    let deserialized = SmtLeaf::read_from_bytes(&serialized).unwrap();
379
380    assert_eq!(multiple_leaf, deserialized);
381}
382
383// HELPERS
384// --------------------------------------------------------------------------------------------
385
386fn build_empty_or_single_leaf_node(key: RpoDigest, value: Word) -> RpoDigest {
387    if value == EMPTY_WORD {
388        SmtLeaf::new_empty(key.into()).hash()
389    } else {
390        SmtLeaf::Single((key, value)).hash()
391    }
392}
393
394fn build_multiple_leaf_node(kv_pairs: &[(RpoDigest, Word)]) -> RpoDigest {
395    let elements: Vec<Felt> = kv_pairs
396        .iter()
397        .flat_map(|(key, value)| {
398            let key_elements = key.into_iter();
399            let value_elements = (*value).into_iter();
400
401            key_elements.chain(value_elements)
402        })
403        .collect();
404
405    Rpo256::hash_elements(&elements)
406}