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 // STATE MUTATORS
112 // --------------------------------------------------------------------------------------------
113
114 /// Inserts the specified key-value pair into this store and returns the value previously
115 /// associated with the specified key.
116 ///
117 /// If no value was previously associated with the specified key, None is returned.
118 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
119 let prefix = get_key_prefix(&key);
120 match self.values.entry(prefix) {
121 Entry::Occupied(mut entry) => entry.get_mut().insert(key, value),
122 Entry::Vacant(entry) => {
123 entry.insert(StoreEntry::new(key, value));
124 None
125 }
126 }
127 }
128
129 /// Removes the key-value pair for the specified key from this store and returns the value
130 /// associated with this key.
131 ///
132 /// If no value was associated with the specified key, None is returned.
133 pub fn remove(&mut self, key: &RpoDigest) -> Option<Word> {
134 let prefix = get_key_prefix(key);
135 match self.values.entry(prefix) {
136 Entry::Occupied(mut entry) => {
137 let (value, remove_entry) = entry.get_mut().remove(key);
138 if remove_entry {
139 entry.remove_entry();
140 }
141 value
142 }
143 Entry::Vacant(_) => None,
144 }
145 }
146
147 // HELPER METHODS
148 // --------------------------------------------------------------------------------------------
149
150 /// Returns an iterator over all key-value pairs contained in this store such that the most
151 /// significant 64 bits of the key lay within the specified bounds.
152 ///
153 /// The order of iteration is from the smallest to the largest key.
154 fn range<R: RangeBounds<u64>>(&self, bounds: R) -> impl Iterator<Item = &(RpoDigest, Word)> {
155 self.values.range(bounds).flat_map(|(_, entry)| entry.iter())
156 }
157}
158
159// VALUE NODE
160// ================================================================================================
161
162/// An entry in the [ValueStore].
163///
164/// An entry can contain either a single key-value pair or a vector of key-value pairs sorted by
165/// key.
166#[derive(Debug, Clone, PartialEq, Eq)]
167pub enum StoreEntry {
168 Single((RpoDigest, Word)),
169 List(Vec<(RpoDigest, Word)>),
170}
171
172impl StoreEntry {
173 // CONSTRUCTOR
174 // --------------------------------------------------------------------------------------------
175 /// Returns a new [StoreEntry] instantiated with a single key-value pair.
176 pub fn new(key: RpoDigest, value: Word) -> Self {
177 Self::Single((key, value))
178 }
179
180 // PUBLIC ACCESSORS
181 // --------------------------------------------------------------------------------------------
182
183 /// Returns the value associated with the specified key, or None if this entry does not contain
184 /// a value associated with the specified key.
185 pub fn get(&self, key: &RpoDigest) -> Option<&Word> {
186 match self {
187 StoreEntry::Single(kv_pair) => {
188 if kv_pair.0 == *key {
189 Some(&kv_pair.1)
190 } else {
191 None
192 }
193 }
194 StoreEntry::List(kv_pairs) => {
195 match kv_pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, key)) {
196 Ok(pos) => Some(&kv_pairs[pos].1),
197 Err(_) => None,
198 }
199 }
200 }
201 }
202
203 /// Returns an iterator over all key-value pairs in this entry.
204 pub fn iter(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
205 EntryIterator {
206 entry: self,
207 pos: 0,
208 }
209 }
210
211 // STATE MUTATORS
212 // --------------------------------------------------------------------------------------------
213
214 /// Inserts the specified key-value pair into this entry and returns the value previously
215 /// associated with the specified key, or None if no value was associated with the specified
216 /// key.
217 ///
218 /// If a new key is inserted, this will also transform a `SingleEntry` into a `ListEntry`.
219 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
220 match self {
221 StoreEntry::Single(kv_pair) => {
222 // if the key is already in this entry, update the value and return
223 if kv_pair.0 == key {
224 let old_value = kv_pair.1;
225 kv_pair.1 = value;
226 return Some(old_value);
227 }
228
229 // transform the entry into a list entry, and make sure the key-value pairs
230 // are sorted by key
231 let mut pairs = vec![*kv_pair, (key, value)];
232 pairs.sort_by(|a, b| cmp_digests(&a.0, &b.0));
233
234 *self = StoreEntry::List(pairs);
235 None
236 }
237 StoreEntry::List(pairs) => {
238 match pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, &key)) {
239 Ok(pos) => {
240 let old_value = pairs[pos].1;
241 pairs[pos].1 = value;
242 Some(old_value)
243 }
244 Err(pos) => {
245 pairs.insert(pos, (key, value));
246 None
247 }
248 }
249 }
250 }
251 }
252
253 /// Removes the key-value pair with the specified key from this entry, and returns the value
254 /// of the removed pair. If the entry did not contain a key-value pair for the specified key,
255 /// None is returned.
256 ///
257 /// If the last last key-value pair was removed from the entry, the second tuple value will
258 /// be set to true.
259 pub fn remove(&mut self, key: &RpoDigest) -> (Option<Word>, bool) {
260 match self {
261 StoreEntry::Single(kv_pair) => {
262 if kv_pair.0 == *key {
263 (Some(kv_pair.1), true)
264 } else {
265 (None, false)
266 }
267 }
268 StoreEntry::List(kv_pairs) => {
269 match kv_pairs.binary_search_by(|kv_pair| cmp_digests(&kv_pair.0, key)) {
270 Ok(pos) => {
271 let kv_pair = kv_pairs.remove(pos);
272 if kv_pairs.len() == 1 {
273 *self = StoreEntry::Single(kv_pairs[0]);
274 }
275 (Some(kv_pair.1), false)
276 }
277 Err(_) => (None, false),
278 }
279 }
280 }
281 }
282}
283
284/// A custom iterator over key-value pairs of a [StoreEntry].
285///
286/// For a `SingleEntry` this returns only one value, but for `ListEntry`, this iterates over the
287/// entire list of key-value pairs.
288pub struct EntryIterator<'a> {
289 entry: &'a StoreEntry,
290 pos: usize,
291}
292
293impl<'a> Iterator for EntryIterator<'a> {
294 type Item = &'a (RpoDigest, Word);
295
296 fn next(&mut self) -> Option<Self::Item> {
297 match self.entry {
298 StoreEntry::Single(kv_pair) => {
299 if self.pos == 0 {
300 self.pos = 1;
301 Some(kv_pair)
302 } else {
303 None
304 }
305 }
306 StoreEntry::List(kv_pairs) => {
307 if self.pos >= kv_pairs.len() {
308 None
309 } else {
310 let kv_pair = &kv_pairs[self.pos];
311 self.pos += 1;
312 Some(kv_pair)
313 }
314 }
315 }
316 }
317}
318
319// HELPER FUNCTIONS
320// ================================================================================================
321
322/// Compares two digests element-by-element using their integer representations starting with the
323/// most significant element.
324fn cmp_digests(d1: &RpoDigest, d2: &RpoDigest) -> Ordering {
325 let d1 = Word::from(d1);
326 let d2 = Word::from(d2);
327
328 for (v1, v2) in d1.iter().zip(d2.iter()).rev() {
329 let v1 = v1.as_int();
330 let v2 = v2.as_int();
331 if v1 != v2 {
332 return v1.cmp(&v2);
333 }
334 }
335
336 Ordering::Equal
337}
338
339// TESTS
340// ================================================================================================
341
342#[cfg(test)]
343mod tests {
344
345 use super::{RpoDigest, ValueStore};
346 use crate::{
347 merkle::{tiered_smt::values::StoreEntry, NodeIndex},
348 Felt, ONE, WORD_SIZE, ZERO,
349 };
350
351 #[test]
352 fn test_insert() {
353 let mut store = ValueStore::default();
354
355 // insert the first key-value pair into the store
356 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
357 let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]);
358 let value_a = [ONE; WORD_SIZE];
359
360 assert!(store.insert(key_a, value_a).is_none());
361 assert_eq!(store.values.len(), 1);
362
363 let entry = store.values.get(&raw_a).unwrap();
364 let expected_entry = StoreEntry::Single((key_a, value_a));
365 assert_eq!(entry, &expected_entry);
366
367 // insert a key-value pair with a different key into the store; since the keys are
368 // different, another entry is added to the values map
369 let raw_b = 0b_11111110_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
370 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
371 let value_b = [ONE, ZERO, ONE, ZERO];
372
373 assert!(store.insert(key_b, value_b).is_none());
374 assert_eq!(store.values.len(), 2);
375
376 let entry1 = store.values.get(&raw_a).unwrap();
377 let expected_entry1 = StoreEntry::Single((key_a, value_a));
378 assert_eq!(entry1, &expected_entry1);
379
380 let entry2 = store.values.get(&raw_b).unwrap();
381 let expected_entry2 = StoreEntry::Single((key_b, value_b));
382 assert_eq!(entry2, &expected_entry2);
383
384 // insert a key-value pair with the same 64-bit key prefix as the first key; this should
385 // transform the first entry into a List entry
386 let key_c = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]);
387 let value_c = [ONE, ONE, ZERO, ZERO];
388
389 assert!(store.insert(key_c, value_c).is_none());
390 assert_eq!(store.values.len(), 2);
391
392 let entry1 = store.values.get(&raw_a).unwrap();
393 let expected_entry1 = StoreEntry::List(vec![(key_c, value_c), (key_a, value_a)]);
394 assert_eq!(entry1, &expected_entry1);
395
396 let entry2 = store.values.get(&raw_b).unwrap();
397 let expected_entry2 = StoreEntry::Single((key_b, value_b));
398 assert_eq!(entry2, &expected_entry2);
399
400 // replace values for keys a and b
401 let value_a2 = [ONE, ONE, ONE, ZERO];
402 let value_b2 = [ZERO, ZERO, ZERO, ONE];
403
404 assert_eq!(store.insert(key_a, value_a2), Some(value_a));
405 assert_eq!(store.values.len(), 2);
406
407 assert_eq!(store.insert(key_b, value_b2), Some(value_b));
408 assert_eq!(store.values.len(), 2);
409
410 let entry1 = store.values.get(&raw_a).unwrap();
411 let expected_entry1 = StoreEntry::List(vec![(key_c, value_c), (key_a, value_a2)]);
412 assert_eq!(entry1, &expected_entry1);
413
414 let entry2 = store.values.get(&raw_b).unwrap();
415 let expected_entry2 = StoreEntry::Single((key_b, value_b2));
416 assert_eq!(entry2, &expected_entry2);
417
418 // insert one more key-value pair with the same 64-bit key-prefix as the first key
419 let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
420 let value_d = [ZERO, ONE, ZERO, ZERO];
421
422 assert!(store.insert(key_d, value_d).is_none());
423 assert_eq!(store.values.len(), 2);
424
425 let entry1 = store.values.get(&raw_a).unwrap();
426 let expected_entry1 =
427 StoreEntry::List(vec![(key_c, value_c), (key_a, value_a2), (key_d, value_d)]);
428 assert_eq!(entry1, &expected_entry1);
429
430 let entry2 = store.values.get(&raw_b).unwrap();
431 let expected_entry2 = StoreEntry::Single((key_b, value_b2));
432 assert_eq!(entry2, &expected_entry2);
433 }
434
435 #[test]
436 fn test_remove() {
437 // populate the value store
438 let mut store = ValueStore::default();
439
440 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
441 let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]);
442 let value_a = [ONE; WORD_SIZE];
443 store.insert(key_a, value_a);
444
445 let raw_b = 0b_11111110_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
446 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
447 let value_b = [ONE, ZERO, ONE, ZERO];
448 store.insert(key_b, value_b);
449
450 let key_c = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]);
451 let value_c = [ONE, ONE, ZERO, ZERO];
452 store.insert(key_c, value_c);
453
454 let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
455 let value_d = [ZERO, ONE, ZERO, ZERO];
456 store.insert(key_d, value_d);
457
458 assert_eq!(store.values.len(), 2);
459
460 let entry1 = store.values.get(&raw_a).unwrap();
461 let expected_entry1 =
462 StoreEntry::List(vec![(key_c, value_c), (key_a, value_a), (key_d, value_d)]);
463 assert_eq!(entry1, &expected_entry1);
464
465 let entry2 = store.values.get(&raw_b).unwrap();
466 let expected_entry2 = StoreEntry::Single((key_b, value_b));
467 assert_eq!(entry2, &expected_entry2);
468
469 // remove non-existent keys
470 let key_e = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_a)]);
471 assert!(store.remove(&key_e).is_none());
472
473 let raw_f = 0b_11111110_11111111_00011111_11111111_10010110_10010011_11100000_00000000_u64;
474 let key_f = RpoDigest::from([ZERO, ZERO, ONE, Felt::new(raw_f)]);
475 assert!(store.remove(&key_f).is_none());
476
477 // remove keys from the list entry
478 assert_eq!(store.remove(&key_c).unwrap(), value_c);
479 let entry1 = store.values.get(&raw_a).unwrap();
480 let expected_entry1 = StoreEntry::List(vec![(key_a, value_a), (key_d, value_d)]);
481 assert_eq!(entry1, &expected_entry1);
482
483 assert_eq!(store.remove(&key_a).unwrap(), value_a);
484 let entry1 = store.values.get(&raw_a).unwrap();
485 let expected_entry1 = StoreEntry::Single((key_d, value_d));
486 assert_eq!(entry1, &expected_entry1);
487
488 assert_eq!(store.remove(&key_d).unwrap(), value_d);
489 assert!(store.values.get(&raw_a).is_none());
490 assert_eq!(store.values.len(), 1);
491
492 // remove a key from a single entry
493 assert_eq!(store.remove(&key_b).unwrap(), value_b);
494 assert!(store.values.get(&raw_b).is_none());
495 assert_eq!(store.values.len(), 0);
496 }
497
498 #[test]
499 fn test_range() {
500 // populate the value store
501 let mut store = ValueStore::default();
502
503 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
504 let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]);
505 let value_a = [ONE; WORD_SIZE];
506 store.insert(key_a, value_a);
507
508 let raw_b = 0b_11111110_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
509 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
510 let value_b = [ONE, ZERO, ONE, ZERO];
511 store.insert(key_b, value_b);
512
513 let key_c = RpoDigest::from([ONE, ONE, ZERO, Felt::new(raw_a)]);
514 let value_c = [ONE, ONE, ZERO, ZERO];
515 store.insert(key_c, value_c);
516
517 let key_d = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_a)]);
518 let value_d = [ZERO, ONE, ZERO, ZERO];
519 store.insert(key_d, value_d);
520
521 let raw_e = 0b_10101000_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
522 let key_e = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_e)]);
523 let value_e = [ZERO, ZERO, ZERO, ONE];
524 store.insert(key_e, value_e);
525
526 // check the entire range
527 let mut iter = store.range(..u64::MAX);
528 assert_eq!(iter.next(), Some(&(key_e, value_e)));
529 assert_eq!(iter.next(), Some(&(key_c, value_c)));
530 assert_eq!(iter.next(), Some(&(key_a, value_a)));
531 assert_eq!(iter.next(), Some(&(key_d, value_d)));
532 assert_eq!(iter.next(), Some(&(key_b, value_b)));
533 assert_eq!(iter.next(), None);
534
535 // check all but e
536 let mut iter = store.range(raw_a..u64::MAX);
537 assert_eq!(iter.next(), Some(&(key_c, value_c)));
538 assert_eq!(iter.next(), Some(&(key_a, value_a)));
539 assert_eq!(iter.next(), Some(&(key_d, value_d)));
540 assert_eq!(iter.next(), Some(&(key_b, value_b)));
541 assert_eq!(iter.next(), None);
542
543 // check all but e and b
544 let mut iter = store.range(raw_a..raw_b);
545 assert_eq!(iter.next(), Some(&(key_c, value_c)));
546 assert_eq!(iter.next(), Some(&(key_a, value_a)));
547 assert_eq!(iter.next(), Some(&(key_d, value_d)));
548 assert_eq!(iter.next(), None);
549 }
550
551 #[test]
552 fn test_get_lone_sibling() {
553 // populate the value store
554 let mut store = ValueStore::default();
555
556 let raw_a = 0b_10101010_10101010_00011111_11111111_10010110_10010011_11100000_00000000_u64;
557 let key_a = RpoDigest::from([ZERO, ONE, ONE, Felt::new(raw_a)]);
558 let value_a = [ONE; WORD_SIZE];
559 store.insert(key_a, value_a);
560
561 let raw_b = 0b_11111111_11111111_00011111_11111111_10010110_10010011_11100000_00000000_u64;
562 let key_b = RpoDigest::from([ONE, ONE, ONE, Felt::new(raw_b)]);
563 let value_b = [ONE, ZERO, ONE, ZERO];
564 store.insert(key_b, value_b);
565
566 // check sibling node for `a`
567 let index = NodeIndex::make(32, 0b_10101010_10101010_00011111_11111110);
568 let parent_index = NodeIndex::make(16, 0b_10101010_10101010);
569 assert_eq!(store.get_lone_sibling(index), Some((&key_a, &value_a, parent_index)));
570
571 // check sibling node for `b`
572 let index = NodeIndex::make(32, 0b_11111111_11111111_00011111_11111111);
573 let parent_index = NodeIndex::make(16, 0b_11111111_11111111);
574 assert_eq!(store.get_lone_sibling(index), Some((&key_b, &value_b, parent_index)));
575
576 // check some other sibling for some other index
577 let index = NodeIndex::make(32, 0b_11101010_10101010);
578 assert_eq!(store.get_lone_sibling(index), None);
579 }
580}