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