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).map(|value| {
130 if !self.updates.contains(key) {
131 self.trace.borrow_mut().insert(key.clone(), value.clone());
132 }
133 value
134 })
135 }
136
137 /// Returns a boolean to indicate whether the given key exists in the data set.
138 ///
139 /// If the key is part of the initial data set, the key access is recorded.
140 fn contains_key(&self, key: &K) -> bool {
141 self.get(key).is_some()
142 }
143
144 /// Returns the number of key-value pairs in the data set.
145 fn len(&self) -> usize {
146 self.data.len()
147 }
148
149 // MUTATORS
150 // --------------------------------------------------------------------------------------------
151
152 /// Inserts a key-value pair into the data set.
153 ///
154 /// If the key already exists in the data set, the value is updated and the old value is
155 /// returned.
156 fn insert(&mut self, key: K, value: V) -> Option<V> {
157 let new_update = self.updates.insert(key.clone());
158 self.data.insert(key.clone(), value).map(|old_value| {
159 if new_update {
160 self.trace.borrow_mut().insert(key, old_value.clone());
161 }
162 old_value
163 })
164 }
165
166 /// Removes a key-value pair from the data set.
167 ///
168 /// If the key exists in the data set, the old value is returned.
169 fn remove(&mut self, key: &K) -> Option<V> {
170 self.data.remove(key).map(|old_value| {
171 let new_update = self.updates.insert(key.clone());
172 if new_update {
173 self.trace.borrow_mut().insert(key.clone(), old_value.clone());
174 }
175 old_value
176 })
177 }
178
179 // ITERATION
180 // --------------------------------------------------------------------------------------------
181
182 /// Returns an iterator over the key-value pairs in the data set.
183 fn iter(&self) -> Box<dyn Iterator<Item = (&K, &V)> + '_> {
184 Box::new(self.data.iter())
185 }
186}
187
188impl<K: Clone + Ord, V: Clone> Extend<(K, V)> for RecordingMap<K, V> {
189 fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
190 iter.into_iter().for_each(move |(k, v)| {
191 self.insert(k, v);
192 });
193 }
194}
195
196impl<K: Clone + Ord, V: Clone> FromIterator<(K, V)> for RecordingMap<K, V> {
197 fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
198 Self::new(iter)
199 }
200}
201
202impl<K: Clone + Ord, V: Clone> IntoIterator for RecordingMap<K, V> {
203 type Item = (K, V);
204 type IntoIter = alloc::collections::btree_map::IntoIter<K, V>;
205
206 fn into_iter(self) -> Self::IntoIter {
207 self.data.into_iter()
208 }
209}
210
211// TESTS
212// ================================================================================================
213
214#[cfg(test)]
215mod tests {
216 use super::*;
217
218 const ITEMS: [(u64, u64); 5] = [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4)];
219
220 #[test]
221 fn test_get_item() {
222 // instantiate a recording map
223 let map = RecordingMap::new(ITEMS.to_vec());
224
225 // get a few items
226 let get_items = [0, 1, 2];
227 for key in get_items.iter() {
228 map.get(key);
229 }
230
231 // convert the map into a proof
232 let (_, proof) = map.finalize();
233
234 // check that the proof contains the expected values
235 for (key, value) in ITEMS.iter() {
236 match get_items.contains(key) {
237 true => assert_eq!(proof.get(key), Some(value)),
238 false => assert_eq!(proof.get(key), None),
239 }
240 }
241 }
242
243 #[test]
244 fn test_contains_key() {
245 // instantiate a recording map
246 let map = RecordingMap::new(ITEMS.to_vec());
247
248 // check if the map contains a few items
249 let get_items = [0, 1, 2];
250 for key in get_items.iter() {
251 map.contains_key(key);
252 }
253
254 // convert the map into a proof
255 let (_, proof) = map.finalize();
256
257 // check that the proof contains the expected values
258 for (key, _) in ITEMS.iter() {
259 match get_items.contains(key) {
260 true => assert!(proof.contains_key(key)),
261 false => assert!(!proof.contains_key(key)),
262 }
263 }
264 }
265
266 #[test]
267 fn test_len() {
268 // instantiate a recording map
269 let mut map = RecordingMap::new(ITEMS.to_vec());
270 // length of the map should be equal to the number of items
271 assert_eq!(map.len(), ITEMS.len());
272
273 // inserting entry with key that already exists should not change the length, but it does
274 // add entries to the trace and update sets
275 map.insert(4, 5);
276 assert_eq!(map.len(), ITEMS.len());
277 assert_eq!(map.trace_len(), 1);
278 assert_eq!(map.updates_len(), 1);
279
280 // inserting entry with new key should increase the length; it should also record the key
281 // as an updated key, but the trace length does not change since old values were not touched
282 map.insert(5, 5);
283 assert_eq!(map.len(), ITEMS.len() + 1);
284 assert_eq!(map.trace_len(), 1);
285 assert_eq!(map.updates_len(), 2);
286
287 // get some items so that they are saved in the trace; this should record original items
288 // in the trace, but should not affect the set of updates
289 let get_items = [0, 1, 2];
290 for key in get_items.iter() {
291 map.contains_key(key);
292 }
293 assert_eq!(map.trace_len(), 4);
294 assert_eq!(map.updates_len(), 2);
295
296 // read the same items again, this should not have any effect on either length, trace, or
297 // the set of updates
298 let get_items = [0, 1, 2];
299 for key in get_items.iter() {
300 map.contains_key(key);
301 }
302 assert_eq!(map.trace_len(), 4);
303 assert_eq!(map.updates_len(), 2);
304
305 // read a newly inserted item; this should not affect either length, trace, or the set of
306 // updates
307 let _val = map.get(&5).unwrap();
308 assert_eq!(map.trace_len(), 4);
309 assert_eq!(map.updates_len(), 2);
310
311 // update a newly inserted item; this should not affect either length, trace, or the set
312 // of updates
313 map.insert(5, 11);
314 assert_eq!(map.trace_len(), 4);
315 assert_eq!(map.updates_len(), 2);
316
317 // Note: The length reported by the proof will be different to the length originally
318 // reported by the map.
319 let (_, proof) = map.finalize();
320
321 // length of the proof should be equal to get_items + 1. The extra item is the original
322 // value at key = 4u64
323 assert_eq!(proof.len(), get_items.len() + 1);
324 }
325
326 #[test]
327 fn test_iter() {
328 let mut map = RecordingMap::new(ITEMS.to_vec());
329 assert!(map.iter().all(|(x, y)| ITEMS.contains(&(*x, *y))));
330
331 // when inserting entry with key that already exists the iterator should return the new value
332 let new_value = 5;
333 map.insert(4, new_value);
334 assert_eq!(map.iter().count(), ITEMS.len());
335 assert!(map.iter().all(|(x, y)| if x == &4 {
336 y == &new_value
337 } else {
338 ITEMS.contains(&(*x, *y))
339 }));
340 }
341
342 #[test]
343 fn test_is_empty() {
344 // instantiate an empty recording map
345 let empty_map: RecordingMap<u64, u64> = RecordingMap::default();
346 assert!(empty_map.is_empty());
347
348 // instantiate a non-empty recording map
349 let map = RecordingMap::new(ITEMS.to_vec());
350 assert!(!map.is_empty());
351 }
352
353 #[test]
354 fn test_remove() {
355 let mut map = RecordingMap::new(ITEMS.to_vec());
356
357 // remove an item that exists
358 let key = 0;
359 let value = map.remove(&key).unwrap();
360 assert_eq!(value, ITEMS[0].1);
361 assert_eq!(map.len(), ITEMS.len() - 1);
362 assert_eq!(map.trace_len(), 1);
363 assert_eq!(map.updates_len(), 1);
364
365 // add the item back and then remove it again
366 let key = 0;
367 let value = 0;
368 map.insert(key, value);
369 let value = map.remove(&key).unwrap();
370 assert_eq!(value, 0);
371 assert_eq!(map.len(), ITEMS.len() - 1);
372 assert_eq!(map.trace_len(), 1);
373 assert_eq!(map.updates_len(), 1);
374
375 // remove an item that does not exist
376 let key = 100;
377 let value = map.remove(&key);
378 assert_eq!(value, None);
379 assert_eq!(map.len(), ITEMS.len() - 1);
380 assert_eq!(map.trace_len(), 1);
381 assert_eq!(map.updates_len(), 1);
382
383 // insert a new item and then remove it
384 let key = 100;
385 let value = 100;
386 map.insert(key, value);
387 let value = map.remove(&key).unwrap();
388 assert_eq!(value, 100);
389 assert_eq!(map.len(), ITEMS.len() - 1);
390 assert_eq!(map.trace_len(), 1);
391 assert_eq!(map.updates_len(), 2);
392
393 // convert the map into a proof
394 let (_, proof) = map.finalize();
395
396 // check that the proof contains the expected values
397 for (key, value) in ITEMS.iter() {
398 match key {
399 0 => assert_eq!(proof.get(key), Some(value)),
400 _ => assert_eq!(proof.get(key), None),
401 }
402 }
403 }
404}