miden-crypto/src/utils/kv_map.rs
src/utils/kv_map.rsRaw
1use super::{collections::ApplyDiff, diff::Diff};
2use core::cell::RefCell;
3use winter_utils::{
4 collections::{btree_map::IntoIter, BTreeMap, BTreeSet},
5 Box,
6};
7
8// KEY-VALUE MAP TRAIT
9// ================================================================================================
10
11/// A trait that defines the interface for a key-value map.
12pub trait KvMap<K: Ord + Clone, V: Clone>:
13 Extend<(K, V)> + FromIterator<(K, V)> + IntoIterator<Item = (K, V)>
14{
15 fn get(&self, key: &K) -> Option<&V>;
16 fn contains_key(&self, key: &K) -> bool;
17 fn len(&self) -> usize;
18 fn is_empty(&self) -> bool {
19 self.len() == 0
20 }
21 fn insert(&mut self, key: K, value: V) -> Option<V>;
22 fn remove(&mut self, key: &K) -> Option<V>;
23
24 fn iter(&self) -> Box<dyn Iterator<Item = (&K, &V)> + '_>;
25}
26
27// BTREE MAP `KvMap` IMPLEMENTATION
28// ================================================================================================
29
30impl<K: Ord + Clone, V: Clone> KvMap<K, V> for BTreeMap<K, V> {
31 fn get(&self, key: &K) -> Option<&V> {
32 self.get(key)
33 }
34
35 fn contains_key(&self, key: &K) -> bool {
36 self.contains_key(key)
37 }
38
39 fn len(&self) -> usize {
40 self.len()
41 }
42
43 fn insert(&mut self, key: K, value: V) -> Option<V> {
44 self.insert(key, value)
45 }
46
47 fn remove(&mut self, key: &K) -> Option<V> {
48 self.remove(key)
49 }
50
51 fn iter(&self) -> Box<dyn Iterator<Item = (&K, &V)> + '_> {
52 Box::new(self.iter())
53 }
54}
55
56// RECORDING MAP
57// ================================================================================================
58
59/// A [RecordingMap] that records read requests to the underlying key-value map.
60///
61/// The data recorder is used to generate a proof for read requests.
62///
63/// The [RecordingMap] is composed of three parts:
64/// - `data`: which contains the current set of key-value pairs in the map.
65/// - `updates`: which tracks keys for which values have been changed since the map was
66/// instantiated. updates include both insertions, removals and updates of values under existing
67/// keys.
68/// - `trace`: which contains the key-value pairs from the original data which have been accesses
69/// since the map was instantiated.
70#[derive(Debug, Default, Clone, Eq, PartialEq)]
71pub struct RecordingMap<K, V> {
72 data: BTreeMap<K, V>,
73 updates: BTreeSet<K>,
74 trace: RefCell<BTreeMap<K, V>>,
75}
76
77impl<K: Ord + Clone, V: Clone> RecordingMap<K, V> {
78 // CONSTRUCTOR
79 // --------------------------------------------------------------------------------------------
80 /// Returns a new [RecordingMap] instance initialized with the provided key-value pairs.
81 /// ([BTreeMap]).
82 pub fn new(init: impl IntoIterator<Item = (K, V)>) -> Self {
83 RecordingMap {
84 data: init.into_iter().collect(),
85 updates: BTreeSet::new(),
86 trace: RefCell::new(BTreeMap::new()),
87 }
88 }
89
90 // PUBLIC ACCESSORS
91 // --------------------------------------------------------------------------------------------
92
93 pub fn inner(&self) -> &BTreeMap<K, V> {
94 &self.data
95 }
96
97 // FINALIZER
98 // --------------------------------------------------------------------------------------------
99
100 /// Consumes the [RecordingMap] and returns a ([BTreeMap], [BTreeMap]) tuple. The first
101 /// element of the tuple is a map that represents the state of the map at the time `.finalize()`
102 /// is called. The second element contains the key-value pairs from the initial data set that
103 /// were read during recording.
104 pub fn finalize(self) -> (BTreeMap<K, V>, BTreeMap<K, V>) {
105 (self.data, self.trace.take())
106 }
107
108 // TEST HELPERS
109 // --------------------------------------------------------------------------------------------
110
111 #[cfg(test)]
112 pub fn trace_len(&self) -> usize {
113 self.trace.borrow().len()
114 }
115
116 #[cfg(test)]
117 pub fn updates_len(&self) -> usize {
118 self.updates.len()
119 }
120}
121
122impl<K: Ord + Clone, V: Clone> KvMap<K, V> for RecordingMap<K, V> {
123 // PUBLIC ACCESSORS
124 // --------------------------------------------------------------------------------------------
125
126 /// Returns a reference to the value associated with the given key if the value exists.
127 ///
128 /// If the key is part of the initial data set, the key access is recorded.
129 fn get(&self, key: &K) -> Option<&V> {
130 self.data.get(key).map(|value| {
131 if !self.updates.contains(key) {
132 self.trace.borrow_mut().insert(key.clone(), value.clone());
133 }
134 value
135 })
136 }
137
138 /// Returns a boolean to indicate whether the given key exists in the data set.
139 ///
140 /// If the key is part of the initial data set, the key access is recorded.
141 fn contains_key(&self, key: &K) -> bool {
142 self.get(key).is_some()
143 }
144
145 /// Returns the number of key-value pairs in the data set.
146 fn len(&self) -> usize {
147 self.data.len()
148 }
149
150 // MUTATORS
151 // --------------------------------------------------------------------------------------------
152
153 /// Inserts a key-value pair into the data set.
154 ///
155 /// If the key already exists in the data set, the value is updated and the old value is
156 /// returned.
157 fn insert(&mut self, key: K, value: V) -> Option<V> {
158 let new_update = self.updates.insert(key.clone());
159 self.data.insert(key.clone(), value).map(|old_value| {
160 if new_update {
161 self.trace.borrow_mut().insert(key, old_value.clone());
162 }
163 old_value
164 })
165 }
166
167 /// Removes a key-value pair from the data set.
168 ///
169 /// If the key exists in the data set, the old value is returned.
170 fn remove(&mut self, key: &K) -> Option<V> {
171 self.data.remove(key).map(|old_value| {
172 let new_update = self.updates.insert(key.clone());
173 if new_update {
174 self.trace.borrow_mut().insert(key.clone(), old_value.clone());
175 }
176 old_value
177 })
178 }
179
180 // ITERATION
181 // --------------------------------------------------------------------------------------------
182
183 /// Returns an iterator over the key-value pairs in the data set.
184 fn iter(&self) -> Box<dyn Iterator<Item = (&K, &V)> + '_> {
185 Box::new(self.data.iter())
186 }
187}
188
189impl<K: Clone + Ord, V: Clone> Extend<(K, V)> for RecordingMap<K, V> {
190 fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
191 iter.into_iter().for_each(move |(k, v)| {
192 self.insert(k, v);
193 });
194 }
195}
196
197impl<K: Clone + Ord, V: Clone> FromIterator<(K, V)> for RecordingMap<K, V> {
198 fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
199 Self::new(iter)
200 }
201}
202
203impl<K: Clone + Ord, V: Clone> IntoIterator for RecordingMap<K, V> {
204 type Item = (K, V);
205 type IntoIter = IntoIter<K, V>;
206
207 fn into_iter(self) -> Self::IntoIter {
208 self.data.into_iter()
209 }
210}
211
212// KV MAP DIFF
213// ================================================================================================
214/// [KvMapDiff] stores the difference between two key-value maps.
215///
216/// The [KvMapDiff] is composed of two parts:
217/// - `updates` - a map of key-value pairs that were updated in the second map compared to the
218/// first map. This includes new key-value pairs.
219/// - `removed` - a set of keys that were removed from the second map compared to the first map.
220#[derive(Debug, Clone)]
221pub struct KvMapDiff<K, V> {
222 pub updated: BTreeMap<K, V>,
223 pub removed: BTreeSet<K>,
224}
225
226impl<K, V> KvMapDiff<K, V> {
227 // CONSTRUCTOR
228 // --------------------------------------------------------------------------------------------
229 /// Creates a new [KvMapDiff] instance.
230 pub fn new() -> Self {
231 KvMapDiff {
232 updated: BTreeMap::new(),
233 removed: BTreeSet::new(),
234 }
235 }
236}
237
238impl<K, V> Default for KvMapDiff<K, V> {
239 fn default() -> Self {
240 Self::new()
241 }
242}
243
244impl<K: Ord + Clone, V: Clone + PartialEq, T: KvMap<K, V>> Diff<K, V> for T {
245 type DiffType = KvMapDiff<K, V>;
246
247 fn diff(&self, other: &T) -> Self::DiffType {
248 let mut diff = KvMapDiff::default();
249 for (k, v) in self.iter() {
250 if let Some(other_value) = other.get(k) {
251 if v != other_value {
252 diff.updated.insert(k.clone(), other_value.clone());
253 }
254 } else {
255 diff.removed.insert(k.clone());
256 }
257 }
258 for (k, v) in other.iter() {
259 if self.get(k).is_none() {
260 diff.updated.insert(k.clone(), v.clone());
261 }
262 }
263 diff
264 }
265}
266
267impl<K: Ord + Clone, V: Clone, T: KvMap<K, V>> ApplyDiff<K, V> for T {
268 type DiffType = KvMapDiff<K, V>;
269
270 fn apply(&mut self, diff: Self::DiffType) {
271 for (k, v) in diff.updated {
272 self.insert(k, v);
273 }
274 for k in diff.removed {
275 self.remove(&k);
276 }
277 }
278}
279
280// TESTS
281// ================================================================================================
282
283#[cfg(test)]
284mod tests {
285 use super::*;
286
287 const ITEMS: [(u64, u64); 5] = [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4)];
288
289 #[test]
290 fn test_get_item() {
291 // instantiate a recording map
292 let map = RecordingMap::new(ITEMS.to_vec());
293
294 // get a few items
295 let get_items = [0, 1, 2];
296 for key in get_items.iter() {
297 map.get(key);
298 }
299
300 // convert the map into a proof
301 let (_, proof) = map.finalize();
302
303 // check that the proof contains the expected values
304 for (key, value) in ITEMS.iter() {
305 match get_items.contains(key) {
306 true => assert_eq!(proof.get(key), Some(value)),
307 false => assert_eq!(proof.get(key), None),
308 }
309 }
310 }
311
312 #[test]
313 fn test_contains_key() {
314 // instantiate a recording map
315 let map = RecordingMap::new(ITEMS.to_vec());
316
317 // check if the map contains a few items
318 let get_items = [0, 1, 2];
319 for key in get_items.iter() {
320 map.contains_key(key);
321 }
322
323 // convert the map into a proof
324 let (_, proof) = map.finalize();
325
326 // check that the proof contains the expected values
327 for (key, _) in ITEMS.iter() {
328 match get_items.contains(key) {
329 true => assert!(proof.contains_key(key)),
330 false => assert!(!proof.contains_key(key)),
331 }
332 }
333 }
334
335 #[test]
336 fn test_len() {
337 // instantiate a recording map
338 let mut map = RecordingMap::new(ITEMS.to_vec());
339 // length of the map should be equal to the number of items
340 assert_eq!(map.len(), ITEMS.len());
341
342 // inserting entry with key that already exists should not change the length, but it does
343 // add entries to the trace and update sets
344 map.insert(4, 5);
345 assert_eq!(map.len(), ITEMS.len());
346 assert_eq!(map.trace_len(), 1);
347 assert_eq!(map.updates_len(), 1);
348
349 // inserting entry with new key should increase the length; it should also record the key
350 // as an updated key, but the trace length does not change since old values were not touched
351 map.insert(5, 5);
352 assert_eq!(map.len(), ITEMS.len() + 1);
353 assert_eq!(map.trace_len(), 1);
354 assert_eq!(map.updates_len(), 2);
355
356 // get some items so that they are saved in the trace; this should record original items
357 // in the trace, but should not affect the set of updates
358 let get_items = [0, 1, 2];
359 for key in get_items.iter() {
360 map.contains_key(key);
361 }
362 assert_eq!(map.trace_len(), 4);
363 assert_eq!(map.updates_len(), 2);
364
365 // read the same items again, this should not have any effect on either length, trace, or
366 // the set of updates
367 let get_items = [0, 1, 2];
368 for key in get_items.iter() {
369 map.contains_key(key);
370 }
371 assert_eq!(map.trace_len(), 4);
372 assert_eq!(map.updates_len(), 2);
373
374 // read a newly inserted item; this should not affect either length, trace, or the set of
375 // updates
376 let _val = map.get(&5).unwrap();
377 assert_eq!(map.trace_len(), 4);
378 assert_eq!(map.updates_len(), 2);
379
380 // update a newly inserted item; this should not affect either length, trace, or the set
381 // of updates
382 map.insert(5, 11);
383 assert_eq!(map.trace_len(), 4);
384 assert_eq!(map.updates_len(), 2);
385
386 // Note: The length reported by the proof will be different to the length originally
387 // reported by the map.
388 let (_, proof) = map.finalize();
389
390 // length of the proof should be equal to get_items + 1. The extra item is the original
391 // value at key = 4u64
392 assert_eq!(proof.len(), get_items.len() + 1);
393 }
394
395 #[test]
396 fn test_iter() {
397 let mut map = RecordingMap::new(ITEMS.to_vec());
398 assert!(map.iter().all(|(x, y)| ITEMS.contains(&(*x, *y))));
399
400 // when inserting entry with key that already exists the iterator should return the new value
401 let new_value = 5;
402 map.insert(4, new_value);
403 assert_eq!(map.iter().count(), ITEMS.len());
404 assert!(map.iter().all(|(x, y)| if x == &4 {
405 y == &new_value
406 } else {
407 ITEMS.contains(&(*x, *y))
408 }));
409 }
410
411 #[test]
412 fn test_is_empty() {
413 // instantiate an empty recording map
414 let empty_map: RecordingMap<u64, u64> = RecordingMap::default();
415 assert!(empty_map.is_empty());
416
417 // instantiate a non-empty recording map
418 let map = RecordingMap::new(ITEMS.to_vec());
419 assert!(!map.is_empty());
420 }
421
422 #[test]
423 fn test_remove() {
424 let mut map = RecordingMap::new(ITEMS.to_vec());
425
426 // remove an item that exists
427 let key = 0;
428 let value = map.remove(&key).unwrap();
429 assert_eq!(value, ITEMS[0].1);
430 assert_eq!(map.len(), ITEMS.len() - 1);
431 assert_eq!(map.trace_len(), 1);
432 assert_eq!(map.updates_len(), 1);
433
434 // add the item back and then remove it again
435 let key = 0;
436 let value = 0;
437 map.insert(key, value);
438 let value = map.remove(&key).unwrap();
439 assert_eq!(value, 0);
440 assert_eq!(map.len(), ITEMS.len() - 1);
441 assert_eq!(map.trace_len(), 1);
442 assert_eq!(map.updates_len(), 1);
443
444 // remove an item that does not exist
445 let key = 100;
446 let value = map.remove(&key);
447 assert_eq!(value, None);
448 assert_eq!(map.len(), ITEMS.len() - 1);
449 assert_eq!(map.trace_len(), 1);
450 assert_eq!(map.updates_len(), 1);
451
452 // insert a new item and then remove it
453 let key = 100;
454 let value = 100;
455 map.insert(key, value);
456 let value = map.remove(&key).unwrap();
457 assert_eq!(value, 100);
458 assert_eq!(map.len(), ITEMS.len() - 1);
459 assert_eq!(map.trace_len(), 1);
460 assert_eq!(map.updates_len(), 2);
461
462 // convert the map into a proof
463 let (_, proof) = map.finalize();
464
465 // check that the proof contains the expected values
466 for (key, value) in ITEMS.iter() {
467 match key {
468 0 => assert_eq!(proof.get(key), Some(value)),
469 _ => assert_eq!(proof.get(key), None),
470 }
471 }
472 }
473
474 #[test]
475 fn test_kv_map_diff() {
476 let mut initial_state = ITEMS.into_iter().collect::<BTreeMap<_, _>>();
477 let mut map = RecordingMap::new(initial_state.clone());
478
479 // remove an item that exists
480 let key = 0;
481 let _value = map.remove(&key).unwrap();
482
483 // add a new item
484 let key = 100;
485 let value = 100;
486 map.insert(key, value);
487
488 // update an existing item
489 let key = 1;
490 let value = 100;
491 map.insert(key, value);
492
493 // compute a diff
494 let diff = initial_state.diff(map.inner());
495 assert!(diff.updated.len() == 2);
496 assert!(diff.updated.iter().all(|(k, v)| [(100, 100), (1, 100)].contains(&(*k, *v))));
497 assert!(diff.removed.len() == 1);
498 assert!(diff.removed.first() == Some(&0));
499
500 // apply the diff to the initial state and assert the contents are the same as the map
501 initial_state.apply(diff);
502 assert!(initial_state.iter().eq(map.iter()));
503 }
504}