miden-crypto/src/merkle/tiered_smt/values.rs
src/merkle/tiered_smt/values.rsRaw
1use super::{get_key_prefix, is_leaf_node, BTreeMap, NodeIndex, RpoDigest, StarkField, Vec, Word};
2use crate::utils::vec;
3use core::{
4 cmp::{Ord, Ordering},
5 ops::RangeBounds,
6};
7use winter_utils::collections::btree_map::Entry;
8
9// CONSTANTS
10// ================================================================================================
11
12/// Depths at which leaves can exist in a tiered SMT.
13const TIER_DEPTHS: [u8; 4] = super::TieredSmt::TIER_DEPTHS;
14
15/// Maximum node depth. This is also the bottom tier of the tree.
16const MAX_DEPTH: u8 = super::TieredSmt::MAX_DEPTH;
17
18// VALUE STORE
19// ================================================================================================
20/// A store for key-value pairs for a Tiered Sparse Merkle tree.
21///
22/// The store is organized in a [BTreeMap] where keys are 64 most significant bits of a key, and
23/// the values are the corresponding key-value pairs (or a list of key-value pairs if more that
24/// a single key-value pair shares the same 64-bit prefix).
25///
26/// The store supports lookup by the full key as well as by the 64-bit key prefix.
27#[derive(Debug, Default, Clone, PartialEq, Eq)]
28pub struct ValueStore {
29 values: BTreeMap<u64, StoreEntry>,
30}
31
32impl ValueStore {
33 // PUBLIC ACCESSORS
34 // --------------------------------------------------------------------------------------------
35
36 /// Returns a reference to the value stored under the specified key, or None if there is no
37 /// value associated with the specified key.
38 pub fn get(&self, key: &RpoDigest) -> Option<&Word> {
39 let prefix = get_key_prefix(key);
40 self.values.get(&prefix).and_then(|entry| entry.get(key))
41 }
42
43 /// Returns the first key-value pair such that the key prefix is greater than or equal to the
44 /// specified prefix.
45 pub fn get_first(&self, prefix: u64) -> Option<&(RpoDigest, Word)> {
46 self.range(prefix..).next()
47 }
48
49 /// Returns the first key-value pair such that the key prefix is greater than or equal to the
50 /// specified prefix and the key value is not equal to the exclude_key value.
51 pub fn get_first_filtered(
52 &self,
53 prefix: u64,
54 exclude_key: &RpoDigest,
55 ) -> Option<&(RpoDigest, Word)> {
56 self.range(prefix..).find(|(key, _)| key != exclude_key)
57 }
58
59 /// Returns a vector with key-value pairs for all keys with the specified 64-bit prefix, or
60 /// None if no keys with the specified prefix are present in this store.
61 pub fn get_all(&self, prefix: u64) -> Option<Vec<(RpoDigest, Word)>> {
62 self.values.get(&prefix).map(|entry| match entry {
63 StoreEntry::Single(kv_pair) => vec![*kv_pair],
64 StoreEntry::List(kv_pairs) => kv_pairs.clone(),
65 })
66 }
67
68 /// Returns information about a sibling of a leaf node with the specified index, but only if
69 /// this is the only sibling the leaf has in some subtree starting at the first tier.
70 ///
71 /// For example, if `index` is an index at depth 32, and there is a leaf node at depth 32 with
72 /// the same root at depth 16 as `index`, we say that this leaf is a lone sibling.
73 ///
74 /// The returned tuple contains: they key-value pair of the sibling as well as the index of
75 /// the node for the root of the common subtree in which both nodes are leaves.
76 ///
77 /// This method assumes that the key-value pair for the specified index has already been
78 /// removed from the store.
79 pub fn get_lone_sibling(&self, index: NodeIndex) -> Option<(&RpoDigest, &Word, NodeIndex)> {
80 debug_assert!(is_leaf_node(&index));
81
82 // iterate over tiers from top to bottom, looking at the tiers which are strictly above
83 // the depth of the index. This implies that only tiers at depth 32 and 48 will be
84 // considered. For each tier, check if the parent of the index at the higher tier
85 // contains a single node.
86 for &tier in TIER_DEPTHS.iter().filter(|&t| index.depth() > *t) {
87 // compute the index of the root at a higher tier
88 let mut parent_index = index;
89 parent_index.move_up_to(tier);
90
91 // find the lone sibling, if any; we need to handle the "last node" at a given tier
92 // separately specify the bounds for the search correctly.
93 let start_prefix = parent_index.value() << (MAX_DEPTH - tier);
94 let sibling = if start_prefix.leading_ones() as u8 == tier {
95 let mut iter = self.range(start_prefix..);
96 iter.next().filter(|_| iter.next().is_none())
97 } else {
98 let end_prefix = (parent_index.value() + 1) << (MAX_DEPTH - tier);
99 let mut iter = self.range(start_prefix..end_prefix);
100 iter.next().filter(|_| iter.next().is_none())
101 };
102
103 if let Some((key, value)) = sibling {
104 return Some((key, value, parent_index));
105 }
106 }
107
108 None
109 }
110
111 /// Returns an iterator over all key-value pairs in this store.
112 pub fn iter(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
113 self.values.iter().flat_map(|(_, entry)| entry.iter())
114 }
115
116 // STATE MUTATORS
117 // --------------------------------------------------------------------------------------------
118
119 /// Inserts the specified key-value pair into this store and returns the value previously
120 /// associated with the specified key.
121 ///
122 /// If no value was previously associated with the specified key, None is returned.
123 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
124 let prefix = get_key_prefix(&key);
125 match self.values.entry(prefix) {
126 Entry::Occupied(mut entry) => entry.get_mut().insert(key, value),
127 Entry::Vacant(entry) => {
128 entry.insert(StoreEntry::new(key, value));
129 None
130 }
131 }
132 }
133
134 /// Removes the key-value pair for the specified key from this store and returns the value
135 /// associated with this key.
136 ///
137 /// If no value was associated with the specified key, None is returned.
138 pub fn remove(&mut self, key: &RpoDigest) -> Option<Word> {
139 let prefix = get_key_prefix(key);
140 match self.values.entry(prefix) {
141 Entry::Occupied(mut entry) => {
142 let (value, remove_entry) = entry.get_mut().remove(key);
143 if remove_entry {
144 entry.remove_entry();
145 }
146 value
147 }
148 Entry::Vacant(_) => None,
149 }
150 }
151
152 // HELPER METHODS
153 // --------------------------------------------------------------------------------------------
154
155 /// Returns an iterator over all key-value pairs contained in this store such that the most
156 /// significant 64 bits of the key lay within the specified bounds.
157 ///
158 /// The order of iteration is from the smallest to the largest key.
159 fn range<R: RangeBounds<u64>>(&self, bounds: R) -> impl Iterator<Item = &(RpoDigest, Word)> {
160 self.values.range(bounds).flat_map(|(_, entry)| entry.iter())
161 }
162}
163
164// VALUE NODE
165// ================================================================================================
166
167/// An entry in the [ValueStore].
168///
169/// An entry can contain either a single key-value pair or a vector of key-value pairs sorted by
170/// key.
171#[derive(Debug, Clone, PartialEq, Eq)]
172pub enum StoreEntry {
173 Single((RpoDigest, Word)),
174 List(Vec<(RpoDigest, Word)>),
175}
176
177impl StoreEntry {
178 // CONSTRUCTOR
179 // --------------------------------------------------------------------------------------------
180 /// Returns a new [StoreEntry] instantiated with a single key-value pair.
181 pub fn new(key: RpoDigest, value: Word) -> Self {
182 Self::Single((key, value))
183 }
184
185 // PUBLIC ACCESSORS
186 // --------------------------------------------------------------------------------------------
187
188 /// Returns the value associated with the specified key, or None if this entry does not contain
189 /// a value associated with the specified key.
190 pub fn get(&self, key: &RpoDigest) -> Option<&Word> {
191 match self {
192 StoreEntry::Single(kv_pair) => {
193 if kv_pair.0 == *key {
194 Some(&kv_pair.1)
195 } else {
196 None
197 }
198 }
199 StoreEntry::List(kv_pairs) => {
200 match kv_pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, key)) {
201 Ok(pos) => Some(&kv_pairs[pos].1),
202 Err(_) => None,
203 }
204 }
205 }
206 }
207
208 /// Returns an iterator over all key-value pairs in this entry.
209 pub fn iter(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
210 EntryIterator {
211 entry: self,
212 pos: 0,
213 }
214 }
215
216 // STATE MUTATORS
217 // --------------------------------------------------------------------------------------------
218
219 /// Inserts the specified key-value pair into this entry and returns the value previously
220 /// associated with the specified key, or None if no value was associated with the specified
221 /// key.
222 ///
223 /// If a new key is inserted, this will also transform a `SingleEntry` into a `ListEntry`.
224 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
225 match self {
226 StoreEntry::Single(kv_pair) => {
227 // if the key is already in this entry, update the value and return
228 if kv_pair.0 == key {
229 let old_value = kv_pair.1;
230 kv_pair.1 = value;
231 return Some(old_value);
232 }
233
234 // transform the entry into a list entry, and make sure the key-value pairs
235 // are sorted by key
236 let mut pairs = vec![*kv_pair, (key, value)];
237 pairs.sort_by(|a, b| cmp_digests(&a.0, &b.0));
238
239 *self = StoreEntry::List(pairs);
240 None
241 }
242 StoreEntry::List(pairs) => {
243 match pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, &key)) {
244 Ok(pos) => {
245 let old_value = pairs[pos].1;
246 pairs[pos].1 = value;
247 Some(old_value)
248 }
249 Err(pos) => {
250 pairs.insert(pos, (key, value));
251 None
252 }
253 }
254 }
255 }
256 }
257
258 /// Removes the key-value pair with the specified key from this entry, and returns the value
259 /// of the removed pair. If the entry did not contain a key-value pair for the specified key,
260 /// None is returned.
261 ///
262 /// If the last last key-value pair was removed from the entry, the second tuple value will
263 /// be set to true.
264 pub fn remove(&mut self, key: &RpoDigest) -> (Option<Word>, bool) {
265 match self {
266 StoreEntry::Single(kv_pair) => {
267 if kv_pair.0 == *key {
268 (Some(kv_pair.1), true)
269 } else {
270 (None, false)
271 }
272 }
273 StoreEntry::List(kv_pairs) => {
274 match kv_pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, key)) {
275 Ok(pos) => {
276 let kv_pair = kv_pairs.remove(pos);
277 if kv_pairs.len() == 1 {
278 *self = StoreEntry::Single(kv_pairs[0]);
279 }
280 (Some(kv_pair.1), false)
281 }
282 Err(_) => (None, false),
283 }
284 }
285 }
286 }
287}
288
289/// A custom iterator over key-value pairs of a [StoreEntry].
290///
291/// For a `SingleEntry` this returns only one value, but for `ListEntry`, this iterates over the
292/// entire list of key-value pairs.
293pub struct EntryIterator<'a> {
294 entry: &'a StoreEntry,
295 pos: usize,
296}
297
298impl<'a> Iterator for EntryIterator<'a> {
299 type Item = &'a (RpoDigest, Word);
300
301 fn next(&mut self) -> Option<Self::Item> {
302 match self.entry {
303 StoreEntry::Single(kv_pair) => {
304 if self.pos == 0 {
305 self.pos = 1;
306 Some(kv_pair)
307 } else {
308 None
309 }
310 }
311 StoreEntry::List(kv_pairs) => {
312 if self.pos >= kv_pairs.len() {
313 None
314 } else {
315 let kv_pair = &kv_pairs[self.pos];
316 self.pos += 1;
317 Some(kv_pair)
318 }
319 }
320 }
321 }
322}
323
324// HELPER FUNCTIONS
325// ================================================================================================
326
327/// Compares two digests element-by-element using their integer representations starting with the
328/// most significant element.
329fn cmp_digests(d1: &RpoDigest, d2: &RpoDigest) -> Ordering {
330 let d1 = Word::from(d1);
331 let d2 = Word::from(d2);
332
333 for (v1, v2) in d1.iter().zip(d2.iter()).rev() {
334 let v1 = v1.as_int();
335 let v2 = v2.as_int();
336 if v1 != v2 {
337 return v1.cmp(&v2);
338 }
339 }
340
341 Ordering::Equal
342}
343
344// TESTS
345// ================================================================================================
346
347#[cfg(test)]
348mod tests {
349
350 use super::{RpoDigest, ValueStore};
351 use crate::{
352 merkle::{tiered_smt::values::StoreEntry, NodeIndex},
353 Felt, ONE, WORD_SIZE, ZERO,
354 };
355
356 #[test]
357 fn test_insert() {
358 let mut store = ValueStore::default();
359
360 // insert the first key-value pair into the store
361 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
362 let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]);
363 let value_a = [ONE; WORD_SIZE];
364
365 assert!(store.insert(key_a, value_a).is_none());
366 assert_eq!(store.values.len(), 1);
367
368 let entry = store.values.get(&raw_a).unwrap();
369 let expected_entry = StoreEntry::Single((key_a, value_a));
370 assert_eq!(entry, &expected_entry);
371
372 // insert a key-value pair with a different key into the store; since the keys are
373 // different, another entry is added to the values map
374 let raw_b = 0b_11111110_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
375 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
376 let value_b = [ONE, ZERO, ONE, ZERO];
377
378 assert!(store.insert(key_b, value_b).is_none());
379 assert_eq!(store.values.len(), 2);
380
381 let entry1 = store.values.get(&raw_a).unwrap();
382 let expected_entry1 = StoreEntry::Single((key_a, value_a));
383 assert_eq!(entry1, &expected_entry1);
384
385 let entry2 = store.values.get(&raw_b).unwrap();
386 let expected_entry2 = StoreEntry::Single((key_b, value_b));
387 assert_eq!(entry2, &expected_entry2);
388
389 // insert a key-value pair with the same 64-bit key prefix as the first key; this should
390 // transform the first entry into a List entry
391 let key_c = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]);
392 let value_c = [ONE, ONE, ZERO, ZERO];
393
394 assert!(store.insert(key_c, value_c).is_none());
395 assert_eq!(store.values.len(), 2);
396
397 let entry1 = store.values.get(&raw_a).unwrap();
398 let expected_entry1 = StoreEntry::List(vec![(key_c, value_c), (key_a, value_a)]);
399 assert_eq!(entry1, &expected_entry1);
400
401 let entry2 = store.values.get(&raw_b).unwrap();
402 let expected_entry2 = StoreEntry::Single((key_b, value_b));
403 assert_eq!(entry2, &expected_entry2);
404
405 // replace values for keys a and b
406 let value_a2 = [ONE, ONE, ONE, ZERO];
407 let value_b2 = [ZERO, ZERO, ZERO, ONE];
408
409 assert_eq!(store.insert(key_a, value_a2), Some(value_a));
410 assert_eq!(store.values.len(), 2);
411
412 assert_eq!(store.insert(key_b, value_b2), Some(value_b));
413 assert_eq!(store.values.len(), 2);
414
415 let entry1 = store.values.get(&raw_a).unwrap();
416 let expected_entry1 = StoreEntry::List(vec![(key_c, value_c), (key_a, value_a2)]);
417 assert_eq!(entry1, &expected_entry1);
418
419 let entry2 = store.values.get(&raw_b).unwrap();
420 let expected_entry2 = StoreEntry::Single((key_b, value_b2));
421 assert_eq!(entry2, &expected_entry2);
422
423 // insert one more key-value pair with the same 64-bit key-prefix as the first key
424 let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
425 let value_d = [ZERO, ONE, ZERO, ZERO];
426
427 assert!(store.insert(key_d, value_d).is_none());
428 assert_eq!(store.values.len(), 2);
429
430 let entry1 = store.values.get(&raw_a).unwrap();
431 let expected_entry1 =
432 StoreEntry::List(vec![(key_c, value_c), (key_a, value_a2), (key_d, value_d)]);
433 assert_eq!(entry1, &expected_entry1);
434
435 let entry2 = store.values.get(&raw_b).unwrap();
436 let expected_entry2 = StoreEntry::Single((key_b, value_b2));
437 assert_eq!(entry2, &expected_entry2);
438 }
439
440 #[test]
441 fn test_remove() {
442 // populate the value store
443 let mut store = ValueStore::default();
444
445 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
446 let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]);
447 let value_a = [ONE; WORD_SIZE];
448 store.insert(key_a, value_a);
449
450 let raw_b = 0b_11111110_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
451 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
452 let value_b = [ONE, ZERO, ONE, ZERO];
453 store.insert(key_b, value_b);
454
455 let key_c = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]);
456 let value_c = [ONE, ONE, ZERO, ZERO];
457 store.insert(key_c, value_c);
458
459 let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
460 let value_d = [ZERO, ONE, ZERO, ZERO];
461 store.insert(key_d, value_d);
462
463 assert_eq!(store.values.len(), 2);
464
465 let entry1 = store.values.get(&raw_a).unwrap();
466 let expected_entry1 =
467 StoreEntry::List(vec![(key_c, value_c), (key_a, value_a), (key_d, value_d)]);
468 assert_eq!(entry1, &expected_entry1);
469
470 let entry2 = store.values.get(&raw_b).unwrap();
471 let expected_entry2 = StoreEntry::Single((key_b, value_b));
472 assert_eq!(entry2, &expected_entry2);
473
474 // remove non-existent keys
475 let key_e = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_a)]);
476 assert!(store.remove(&key_e).is_none());
477
478 let raw_f = 0b_11111110_11111111_00011111_11111111_10010110_10010011_11100000_00000000_u64;
479 let key_f = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_f)]);
480 assert!(store.remove(&key_f).is_none());
481
482 // remove keys from the list entry
483 assert_eq!(store.remove(&key_c).unwrap(), value_c);
484 let entry1 = store.values.get(&raw_a).unwrap();
485 let expected_entry1 = StoreEntry::List(vec![(key_a, value_a), (key_d, value_d)]);
486 assert_eq!(entry1, &expected_entry1);
487
488 assert_eq!(store.remove(&key_a).unwrap(), value_a);
489 let entry1 = store.values.get(&raw_a).unwrap();
490 let expected_entry1 = StoreEntry::Single((key_d, value_d));
491 assert_eq!(entry1, &expected_entry1);
492
493 assert_eq!(store.remove(&key_d).unwrap(), value_d);
494 assert!(store.values.get(&raw_a).is_none());
495 assert_eq!(store.values.len(), 1);
496
497 // remove a key from a single entry
498 assert_eq!(store.remove(&key_b).unwrap(), value_b);
499 assert!(store.values.get(&raw_b).is_none());
500 assert_eq!(store.values.len(), 0);
501 }
502
503 #[test]
504 fn test_range() {
505 // populate the value store
506 let mut store = ValueStore::default();
507
508 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
509 let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]);
510 let value_a = [ONE; WORD_SIZE];
511 store.insert(key_a, value_a);
512
513 let raw_b = 0b_11111110_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
514 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
515 let value_b = [ONE, ZERO, ONE, ZERO];
516 store.insert(key_b, value_b);
517
518 let key_c = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]);
519 let value_c = [ONE, ONE, ZERO, ZERO];
520 store.insert(key_c, value_c);
521
522 let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
523 let value_d = [ZERO, ONE, ZERO, ZERO];
524 store.insert(key_d, value_d);
525
526 let raw_e = 0b_10101000_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
527 let key_e = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_e)]);
528 let value_e = [ZERO, ZERO, ZERO, ONE];
529 store.insert(key_e, value_e);
530
531 // check the entire range
532 let mut iter = store.range(..u64::MAX);
533 assert_eq!(iter.next(), Some(&(key_e, value_e)));
534 assert_eq!(iter.next(), Some(&(key_c, value_c)));
535 assert_eq!(iter.next(), Some(&(key_a, value_a)));
536 assert_eq!(iter.next(), Some(&(key_d, value_d)));
537 assert_eq!(iter.next(), Some(&(key_b, value_b)));
538 assert_eq!(iter.next(), None);
539
540 // check all but e
541 let mut iter = store.range(raw_a..u64::MAX);
542 assert_eq!(iter.next(), Some(&(key_c, value_c)));
543 assert_eq!(iter.next(), Some(&(key_a, value_a)));
544 assert_eq!(iter.next(), Some(&(key_d, value_d)));
545 assert_eq!(iter.next(), Some(&(key_b, value_b)));
546 assert_eq!(iter.next(), None);
547
548 // check all but e and b
549 let mut iter = store.range(raw_a..raw_b);
550 assert_eq!(iter.next(), Some(&(key_c, value_c)));
551 assert_eq!(iter.next(), Some(&(key_a, value_a)));
552 assert_eq!(iter.next(), Some(&(key_d, value_d)));
553 assert_eq!(iter.next(), None);
554 }
555
556 #[test]
557 fn test_get_lone_sibling() {
558 // populate the value store
559 let mut store = ValueStore::default();
560
561 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
562 let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]);
563 let value_a = [ONE; WORD_SIZE];
564 store.insert(key_a, value_a);
565
566 let raw_b = 0b_11111111_11111111_00011111_11111111_10010110_10010011_11100000_00000000_u64;
567 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
568 let value_b = [ONE, ZERO, ONE, ZERO];
569 store.insert(key_b, value_b);
570
571 // check sibling node for `a`
572 let index = NodeIndex::make(32, 0b_10101010_10101010_00011111_11111110);
573 let parent_index = NodeIndex::make(16, 0b_10101010_10101010);
574 assert_eq!(store.get_lone_sibling(index), Some((&key_a, &value_a, parent_index)));
575
576 // check sibling node for `b`
577 let index = NodeIndex::make(32, 0b_11111111_11111111_00011111_11111111);
578 let parent_index = NodeIndex::make(16, 0b_11111111_11111111);
579 assert_eq!(store.get_lone_sibling(index), Some((&key_b, &value_b, parent_index)));
580
581 // check some other sibling for some other index
582 let index = NodeIndex::make(32, 0b_11101010_10101010);
583 assert_eq!(store.get_lone_sibling(index), None);
584 }
585}