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