miden-crypto/src/merkle/tiered_smt/values.rs
src/merkle/tiered_smt/values.rsRaw
1use super::{get_key_prefix, BTreeMap, LeafNodeIndex, 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 (i.e. [RpoDigest]) as well as by the 64-bit key
27/// prefix.
28#[derive(Debug, Default, Clone, PartialEq, Eq)]
29pub struct ValueStore {
30 values: BTreeMap<u64, StoreEntry>,
31}
32
33impl ValueStore {
34 // PUBLIC ACCESSORS
35 // --------------------------------------------------------------------------------------------
36
37 /// Returns a reference to the value stored under the specified key, or None if there is no
38 /// value associated with the specified key.
39 pub fn get(&self, key: &RpoDigest) -> Option<&Word> {
40 let prefix = get_key_prefix(key);
41 self.values.get(&prefix).and_then(|entry| entry.get(key))
42 }
43
44 /// Returns the first key-value pair such that the key prefix is greater than or equal to the
45 /// specified prefix.
46 pub fn get_first(&self, prefix: u64) -> Option<&(RpoDigest, Word)> {
47 self.range(prefix..).next()
48 }
49
50 /// Returns the first key-value pair such that the key prefix is greater than or equal to the
51 /// specified prefix and the key value is not equal to the exclude_key value.
52 pub fn get_first_filtered(
53 &self,
54 prefix: u64,
55 exclude_key: &RpoDigest,
56 ) -> Option<&(RpoDigest, Word)> {
57 self.range(prefix..).find(|(key, _)| key != exclude_key)
58 }
59
60 /// Returns a vector with key-value pairs for all keys with the specified 64-bit prefix, or
61 /// None if no keys with the specified prefix are present in this store.
62 pub fn get_all(&self, prefix: u64) -> Option<Vec<(RpoDigest, Word)>> {
63 self.values.get(&prefix).map(|entry| match entry {
64 StoreEntry::Single(kv_pair) => vec![*kv_pair],
65 StoreEntry::List(kv_pairs) => kv_pairs.clone(),
66 })
67 }
68
69 /// Returns information about a sibling of a leaf node with the specified index, but only if
70 /// this is the only sibling the leaf has in some subtree starting at the first tier.
71 ///
72 /// For example, if `index` is an index at depth 32, and there is a leaf node at depth 32 with
73 /// the same root at depth 16 as `index`, we say that this leaf is a lone sibling.
74 ///
75 /// The returned tuple contains: they key-value pair of the sibling as well as the index of
76 /// the node for the root of the common subtree in which both nodes are leaves.
77 ///
78 /// This method assumes that the key-value pair for the specified index has already been
79 /// removed from the store.
80 pub fn get_lone_sibling(
81 &self,
82 index: LeafNodeIndex,
83 ) -> Option<(&RpoDigest, &Word, LeafNodeIndex)> {
84 // iterate over tiers from top to bottom, looking at the tiers which are strictly above
85 // the depth of the index. This implies that only tiers at depth 32 and 48 will be
86 // considered. For each tier, check if the parent of the index at the higher tier
87 // contains a single node. The fist tier (depth 16) is excluded because we cannot move
88 // nodes at depth 16 to a higher tier. This implies that nodes at the first tier will
89 // never have "lone siblings".
90 for &tier_depth in TIER_DEPTHS.iter().filter(|&t| index.depth() > *t) {
91 // compute the index of the root at a higher tier
92 let mut parent_index = index;
93 parent_index.move_up_to(tier_depth);
94
95 // find the lone sibling, if any; we need to handle the "last node" at a given tier
96 // separately specify the bounds for the search correctly.
97 let start_prefix = parent_index.value() << (MAX_DEPTH - tier_depth);
98 let sibling = if start_prefix.leading_ones() as u8 == tier_depth {
99 let mut iter = self.range(start_prefix..);
100 iter.next().filter(|_| iter.next().is_none())
101 } else {
102 let end_prefix = (parent_index.value() + 1) << (MAX_DEPTH - tier_depth);
103 let mut iter = self.range(start_prefix..end_prefix);
104 iter.next().filter(|_| iter.next().is_none())
105 };
106
107 if let Some((key, value)) = sibling {
108 return Some((key, value, parent_index));
109 }
110 }
111
112 None
113 }
114
115 /// Returns an iterator over all key-value pairs in this store.
116 pub fn iter(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
117 self.values.iter().flat_map(|(_, entry)| entry.iter())
118 }
119
120 // STATE MUTATORS
121 // --------------------------------------------------------------------------------------------
122
123 /// Inserts the specified key-value pair into this store and returns the value previously
124 /// associated with the specified key.
125 ///
126 /// If no value was previously associated with the specified key, None is returned.
127 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
128 let prefix = get_key_prefix(&key);
129 match self.values.entry(prefix) {
130 Entry::Occupied(mut entry) => entry.get_mut().insert(key, value),
131 Entry::Vacant(entry) => {
132 entry.insert(StoreEntry::new(key, value));
133 None
134 }
135 }
136 }
137
138 /// Removes the key-value pair for the specified key from this store and returns the value
139 /// associated with this key.
140 ///
141 /// If no value was associated with the specified key, None is returned.
142 pub fn remove(&mut self, key: &RpoDigest) -> Option<Word> {
143 let prefix = get_key_prefix(key);
144 match self.values.entry(prefix) {
145 Entry::Occupied(mut entry) => {
146 let (value, remove_entry) = entry.get_mut().remove(key);
147 if remove_entry {
148 entry.remove_entry();
149 }
150 value
151 }
152 Entry::Vacant(_) => None,
153 }
154 }
155
156 // HELPER METHODS
157 // --------------------------------------------------------------------------------------------
158
159 /// Returns an iterator over all key-value pairs contained in this store such that the most
160 /// significant 64 bits of the key lay within the specified bounds.
161 ///
162 /// The order of iteration is from the smallest to the largest key.
163 fn range<R: RangeBounds<u64>>(&self, bounds: R) -> impl Iterator<Item = &(RpoDigest, Word)> {
164 self.values.range(bounds).flat_map(|(_, entry)| entry.iter())
165 }
166}
167
168// VALUE NODE
169// ================================================================================================
170
171/// An entry in the [ValueStore].
172///
173/// An entry can contain either a single key-value pair or a vector of key-value pairs sorted by
174/// key.
175#[derive(Debug, Clone, PartialEq, Eq)]
176pub enum StoreEntry {
177 Single((RpoDigest, Word)),
178 List(Vec<(RpoDigest, Word)>),
179}
180
181impl StoreEntry {
182 // CONSTRUCTOR
183 // --------------------------------------------------------------------------------------------
184 /// Returns a new [StoreEntry] instantiated with a single key-value pair.
185 pub fn new(key: RpoDigest, value: Word) -> Self {
186 Self::Single((key, value))
187 }
188
189 // PUBLIC ACCESSORS
190 // --------------------------------------------------------------------------------------------
191
192 /// Returns the value associated with the specified key, or None if this entry does not contain
193 /// a value associated with the specified key.
194 pub fn get(&self, key: &RpoDigest) -> Option<&Word> {
195 match self {
196 StoreEntry::Single(kv_pair) => {
197 if kv_pair.0 == *key {
198 Some(&kv_pair.1)
199 } else {
200 None
201 }
202 }
203 StoreEntry::List(kv_pairs) => {
204 match kv_pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, key)) {
205 Ok(pos) => Some(&kv_pairs[pos].1),
206 Err(_) => None,
207 }
208 }
209 }
210 }
211
212 /// Returns an iterator over all key-value pairs in this entry.
213 pub fn iter(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
214 EntryIterator {
215 entry: self,
216 pos: 0,
217 }
218 }
219
220 // STATE MUTATORS
221 // --------------------------------------------------------------------------------------------
222
223 /// Inserts the specified key-value pair into this entry and returns the value previously
224 /// associated with the specified key, or None if no value was associated with the specified
225 /// key.
226 ///
227 /// If a new key is inserted, this will also transform a `SingleEntry` into a `ListEntry`.
228 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
229 match self {
230 StoreEntry::Single(kv_pair) => {
231 // if the key is already in this entry, update the value and return
232 if kv_pair.0 == key {
233 let old_value = kv_pair.1;
234 kv_pair.1 = value;
235 return Some(old_value);
236 }
237
238 // transform the entry into a list entry, and make sure the key-value pairs
239 // are sorted by key
240 let mut pairs = vec![*kv_pair, (key, value)];
241 pairs.sort_by(|a, b| cmp_digests(&a.0, &b.0));
242
243 *self = StoreEntry::List(pairs);
244 None
245 }
246 StoreEntry::List(pairs) => {
247 match pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, &key)) {
248 Ok(pos) => {
249 let old_value = pairs[pos].1;
250 pairs[pos].1 = value;
251 Some(old_value)
252 }
253 Err(pos) => {
254 pairs.insert(pos, (key, value));
255 None
256 }
257 }
258 }
259 }
260 }
261
262 /// Removes the key-value pair with the specified key from this entry, and returns the value
263 /// of the removed pair. If the entry did not contain a key-value pair for the specified key,
264 /// None is returned.
265 ///
266 /// If the last last key-value pair was removed from the entry, the second tuple value will
267 /// be set to true.
268 pub fn remove(&mut self, key: &RpoDigest) -> (Option<Word>, bool) {
269 match self {
270 StoreEntry::Single(kv_pair) => {
271 if kv_pair.0 == *key {
272 (Some(kv_pair.1), true)
273 } else {
274 (None, false)
275 }
276 }
277 StoreEntry::List(kv_pairs) => {
278 match kv_pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, key)) {
279 Ok(pos) => {
280 let kv_pair = kv_pairs.remove(pos);
281 if kv_pairs.len() == 1 {
282 *self = StoreEntry::Single(kv_pairs[0]);
283 }
284 (Some(kv_pair.1), false)
285 }
286 Err(_) => (None, false),
287 }
288 }
289 }
290 }
291}
292
293/// A custom iterator over key-value pairs of a [StoreEntry].
294///
295/// For a `SingleEntry` this returns only one value, but for `ListEntry`, this iterates over the
296/// entire list of key-value pairs.
297pub struct EntryIterator<'a> {
298 entry: &'a StoreEntry,
299 pos: usize,
300}
301
302impl<'a> Iterator for EntryIterator<'a> {
303 type Item = &'a (RpoDigest, Word);
304
305 fn next(&mut self) -> Option<Self::Item> {
306 match self.entry {
307 StoreEntry::Single(kv_pair) => {
308 if self.pos == 0 {
309 self.pos = 1;
310 Some(kv_pair)
311 } else {
312 None
313 }
314 }
315 StoreEntry::List(kv_pairs) => {
316 if self.pos >= kv_pairs.len() {
317 None
318 } else {
319 let kv_pair = &kv_pairs[self.pos];
320 self.pos += 1;
321 Some(kv_pair)
322 }
323 }
324 }
325 }
326}
327
328// HELPER FUNCTIONS
329// ================================================================================================
330
331/// Compares two digests element-by-element using their integer representations starting with the
332/// most significant element.
333fn cmp_digests(d1: &RpoDigest, d2: &RpoDigest) -> Ordering {
334 let d1 = Word::from(d1);
335 let d2 = Word::from(d2);
336
337 for (v1, v2) in d1.iter().zip(d2.iter()).rev() {
338 let v1 = v1.as_int();
339 let v2 = v2.as_int();
340 if v1 != v2 {
341 return v1.cmp(&v2);
342 }
343 }
344
345 Ordering::Equal
346}
347
348// TESTS
349// ================================================================================================
350
351#[cfg(test)]
352mod tests {
353 use super::{LeafNodeIndex, RpoDigest, StoreEntry, ValueStore};
354 use crate::{Felt, ONE, WORD_SIZE, ZERO};
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 = LeafNodeIndex::make(32, 0b_10101010_10101010_00011111_11111110);
573 let parent_index = LeafNodeIndex::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 = LeafNodeIndex::make(32, 0b_11111111_11111111_00011111_11111111);
578 let parent_index = LeafNodeIndex::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 = LeafNodeIndex::make(32, 0b_11101010_10101010);
583 assert_eq!(store.get_lone_sibling(index), None);
584 }
585}