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