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