unc_sdk/collections/
lookup_map.rs

1//! A persistent map without iterators. Unlike `unc_sdk::collections::UnorderedMap` this map
2//! doesn't store keys and values separately in vectors, so it can't iterate over keys. But it
3//! makes this map more efficient in the number of reads and writes.
4use std::marker::PhantomData;
5
6use borsh::{to_vec, BorshDeserialize, BorshSerialize};
7
8use crate::collections::append_slice;
9use crate::{env, IntoStorageKey};
10use unc_sdk_macros::unc;
11
12const ERR_KEY_SERIALIZATION: &str = "Cannot serialize key with Borsh";
13const ERR_VALUE_DESERIALIZATION: &str = "Cannot deserialize value with Borsh";
14const ERR_VALUE_SERIALIZATION: &str = "Cannot serialize value with Borsh";
15
16/// An non-iterable implementation of a map that stores its content directly on the trie.
17#[unc(inside_uncsdk)]
18pub struct LookupMap<K, V> {
19    key_prefix: Vec<u8>,
20    #[borsh(skip)]
21    el: PhantomData<(K, V)>,
22}
23
24impl<K, V> LookupMap<K, V> {
25    /// Create a new map. Use `key_prefix` as a unique prefix for keys.
26    ///
27    /// # Examples
28    ///
29    /// ```
30    /// use unc_sdk::collections::LookupMap;
31    /// let mut map: LookupMap<String, String> = LookupMap::new(b"m");
32    /// ```
33    pub fn new<S>(key_prefix: S) -> Self
34    where
35        S: IntoStorageKey,
36    {
37        Self { key_prefix: key_prefix.into_storage_key(), el: PhantomData }
38    }
39
40    fn raw_key_to_storage_key(&self, raw_key: &[u8]) -> Vec<u8> {
41        append_slice(&self.key_prefix, raw_key)
42    }
43
44    /// Returns `true` if the serialized key is present in the map.
45    fn contains_key_raw(&self, key_raw: &[u8]) -> bool {
46        let storage_key = self.raw_key_to_storage_key(key_raw);
47        env::storage_has_key(&storage_key)
48    }
49
50    /// Returns the serialized value corresponding to the serialized key.
51    fn get_raw(&self, key_raw: &[u8]) -> Option<Vec<u8>> {
52        let storage_key = self.raw_key_to_storage_key(key_raw);
53        env::storage_read(&storage_key)
54    }
55
56    /// Inserts a serialized key-value pair into the map.
57    /// If the map did not have this key present, `None` is returned. Otherwise returns
58    /// a serialized value. Note, the keys that have the same hash value are undistinguished by
59    /// the implementation.
60    pub fn insert_raw(&mut self, key_raw: &[u8], value_raw: &[u8]) -> Option<Vec<u8>> {
61        let storage_key = self.raw_key_to_storage_key(key_raw);
62        if env::storage_write(&storage_key, value_raw) {
63            Some(env::storage_get_evicted().unwrap())
64        } else {
65            None
66        }
67    }
68
69    /// Removes a serialized key from the map, returning the serialized value at the key if the key
70    /// was previously in the map.
71    pub fn remove_raw(&mut self, key_raw: &[u8]) -> Option<Vec<u8>> {
72        let storage_key = self.raw_key_to_storage_key(key_raw);
73        if env::storage_remove(&storage_key) {
74            Some(env::storage_get_evicted().unwrap())
75        } else {
76            None
77        }
78    }
79}
80
81impl<K, V> LookupMap<K, V>
82where
83    K: BorshSerialize,
84    V: BorshSerialize + BorshDeserialize,
85{
86    fn serialize_key(key: &K) -> Vec<u8> {
87        match to_vec(key) {
88            Ok(x) => x,
89            Err(_) => env::panic_str(ERR_KEY_SERIALIZATION),
90        }
91    }
92
93    fn deserialize_value(raw_value: &[u8]) -> V {
94        match V::try_from_slice(raw_value) {
95            Ok(x) => x,
96            Err(_) => env::panic_str(ERR_VALUE_DESERIALIZATION),
97        }
98    }
99
100    fn serialize_value(value: &V) -> Vec<u8> {
101        match to_vec(value) {
102            Ok(x) => x,
103            Err(_) => env::panic_str(ERR_VALUE_SERIALIZATION),
104        }
105    }
106
107    /// Returns true if the map contains a given key.
108    ///
109    /// # Examples
110    ///
111    /// ```
112    /// use unc_sdk::collections::LookupMap;
113    ///
114    /// let mut map: LookupMap<String, String> = LookupMap::new(b"m");
115    /// assert_eq!(map.contains_key(&"Toyota".into()), false);
116    ///
117    /// map.insert(&"Toyota".into(), &"Camry".into());
118    /// assert_eq!(map.contains_key(&"Toyota".into()), true);
119    /// ```
120    pub fn contains_key(&self, key: &K) -> bool {
121        self.contains_key_raw(&Self::serialize_key(key))
122    }
123
124    /// Returns the value corresponding to the key.
125    ///
126    /// # Examples
127    ///
128    /// ```
129    /// use unc_sdk::collections::LookupMap;
130    ///
131    /// let mut map: LookupMap<String, String> = LookupMap::new(b"m");
132    /// assert_eq!(map.get(&"Toyota".into()), None);
133    ///
134    /// map.insert(&"Toyota".into(), &"Camry".into());
135    /// assert_eq!(map.get(&"Toyota".into()), Some("Camry".into()));
136    /// ```
137    pub fn get(&self, key: &K) -> Option<V> {
138        self.get_raw(&Self::serialize_key(key)).map(|value_raw| Self::deserialize_value(&value_raw))
139    }
140
141    /// Removes a key from the map, returning the value at the key if the key was previously in the
142    /// map.
143    ///
144    /// # Examples
145    ///
146    /// ```
147    /// use unc_sdk::collections::LookupMap;
148    ///
149    /// let mut map: LookupMap<String, String> = LookupMap::new(b"m");
150    /// assert_eq!(map.remove(&"Toyota".into()), None);
151    ///
152    /// map.insert(&"Toyota".into(), &"Camry".into());
153    /// assert_eq!(map.remove(&"Toyota".into()), Some("Camry".into()));
154    /// ```
155    pub fn remove(&mut self, key: &K) -> Option<V> {
156        self.remove_raw(&Self::serialize_key(key))
157            .map(|value_raw| Self::deserialize_value(&value_raw))
158    }
159
160    /// Inserts a key-value pair into the map.
161    /// If the map did not have this key present, `None` is returned. Otherwise returns
162    /// a value. Note, the keys that have the same hash value are undistinguished by
163    /// the implementation.
164    ///
165    /// # Examples
166    ///
167    /// ```
168    /// use unc_sdk::collections::LookupMap;
169    ///
170    /// let mut map: LookupMap<String, String> = LookupMap::new(b"m");
171    /// assert_eq!(map.insert(&"Toyota".into(), &"Camry".into()), None);
172    /// assert_eq!(map.insert(&"Toyota".into(), &"Corolla".into()), Some("Camry".into()));
173    /// ```
174    pub fn insert(&mut self, key: &K, value: &V) -> Option<V> {
175        self.insert_raw(&Self::serialize_key(key), &Self::serialize_value(value))
176            .map(|value_raw| Self::deserialize_value(&value_raw))
177    }
178
179    /// Inserts all new key-values from the iterator and replaces values with existing keys
180    /// with new values returned from the iterator.
181    ///
182    /// # Examples
183    ///
184    /// ```
185    /// use unc_sdk::collections::LookupMap;
186    ///
187    /// let mut extendee: LookupMap<String, String> = LookupMap::new(b"m");
188    /// let mut source = vec![];
189    ///
190    /// source.push(("Toyota".into(), "Camry".into()));
191    /// source.push(("Nissan".into(), "Almera".into()));
192    /// source.push(("Ford".into(), "Mustang".into()));
193    /// source.push(("Chevrolet".into(), "Camaro".into()));
194    /// extendee.extend(source.into_iter());
195    /// ```
196    pub fn extend<IT: IntoIterator<Item = (K, V)>>(&mut self, iter: IT) {
197        for (el_key, el_value) in iter {
198            self.insert(&el_key, &el_value);
199        }
200    }
201}
202
203impl<K, V> std::fmt::Debug for LookupMap<K, V>
204where
205    K: std::fmt::Debug + BorshSerialize,
206    V: std::fmt::Debug + BorshSerialize + BorshDeserialize,
207{
208    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
209        f.debug_struct("LookupMap").field("key_prefix", &self.key_prefix).finish()
210    }
211}
212
213#[cfg(not(target_arch = "wasm32"))]
214#[cfg(test)]
215mod tests {
216    use crate::collections::LookupMap;
217    use rand::seq::SliceRandom;
218    use rand::{Rng, SeedableRng};
219    use std::collections::HashMap;
220
221    #[test]
222    pub fn test_insert_one() {
223        let mut map = LookupMap::new(b"m");
224        assert_eq!(None, map.insert(&1, &2));
225        assert_eq!(2, map.insert(&1, &3).unwrap());
226    }
227
228    #[test]
229    pub fn test_insert() {
230        let mut map = LookupMap::new(b"m");
231        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(0);
232        for _ in 0..500 {
233            let key = rng.gen::<u64>();
234            let value = rng.gen::<u64>();
235            map.insert(&key, &value);
236        }
237    }
238
239    #[test]
240    pub fn test_insert_has_key() {
241        let mut map = LookupMap::new(b"m");
242        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(0);
243        let mut key_to_value = HashMap::new();
244        for _ in 0..100 {
245            let key = rng.gen::<u64>();
246            let value = rng.gen::<u64>();
247            map.insert(&key, &value);
248            key_to_value.insert(key, value);
249        }
250        // Non existing
251        for _ in 0..100 {
252            let key = rng.gen::<u64>();
253            assert_eq!(map.contains_key(&key), key_to_value.contains_key(&key));
254        }
255        // Existing
256        for (key, _) in key_to_value.iter() {
257            assert!(map.contains_key(key));
258        }
259    }
260
261    #[test]
262    pub fn test_insert_remove() {
263        let mut map = LookupMap::new(b"m");
264        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(1);
265        let mut keys = vec![];
266        let mut key_to_value = HashMap::new();
267        for _ in 0..100 {
268            let key = rng.gen::<u64>();
269            let value = rng.gen::<u64>();
270            keys.push(key);
271            key_to_value.insert(key, value);
272            map.insert(&key, &value);
273        }
274        keys.shuffle(&mut rng);
275        for key in keys {
276            let actual = map.remove(&key).unwrap();
277            assert_eq!(actual, key_to_value[&key]);
278        }
279    }
280
281    #[test]
282    pub fn test_remove_last_reinsert() {
283        let mut map = LookupMap::new(b"m");
284        let key1 = 1u64;
285        let value1 = 2u64;
286        map.insert(&key1, &value1);
287        let key2 = 3u64;
288        let value2 = 4u64;
289        map.insert(&key2, &value2);
290
291        let actual_value2 = map.remove(&key2).unwrap();
292        assert_eq!(actual_value2, value2);
293
294        let actual_insert_value2 = map.insert(&key2, &value2);
295        assert_eq!(actual_insert_value2, None);
296    }
297
298    #[test]
299    pub fn test_insert_override_remove() {
300        let mut map = LookupMap::new(b"m");
301        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(2);
302        let mut keys = vec![];
303        let mut key_to_value = HashMap::new();
304        for _ in 0..100 {
305            let key = rng.gen::<u64>();
306            let value = rng.gen::<u64>();
307            keys.push(key);
308            key_to_value.insert(key, value);
309            map.insert(&key, &value);
310        }
311        keys.shuffle(&mut rng);
312        for key in &keys {
313            let value = rng.gen::<u64>();
314            let actual = map.insert(key, &value).unwrap();
315            assert_eq!(actual, key_to_value[key]);
316            key_to_value.insert(*key, value);
317        }
318        keys.shuffle(&mut rng);
319        for key in keys {
320            let actual = map.remove(&key).unwrap();
321            assert_eq!(actual, key_to_value[&key]);
322        }
323    }
324
325    #[test]
326    pub fn test_get_non_existent() {
327        let mut map = LookupMap::new(b"m");
328        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(3);
329        let mut key_to_value = HashMap::new();
330        for _ in 0..500 {
331            let key = rng.gen::<u64>() % 20_000;
332            let value = rng.gen::<u64>();
333            key_to_value.insert(key, value);
334            map.insert(&key, &value);
335        }
336        for _ in 0..500 {
337            let key = rng.gen::<u64>() % 20_000;
338            assert_eq!(map.get(&key), key_to_value.get(&key).cloned());
339        }
340    }
341
342    #[test]
343    pub fn test_extend() {
344        let mut map = LookupMap::new(b"m");
345        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(4);
346        let mut key_to_value = HashMap::new();
347        for _ in 0..100 {
348            let key = rng.gen::<u64>();
349            let value = rng.gen::<u64>();
350            key_to_value.insert(key, value);
351            map.insert(&key, &value);
352        }
353        for _ in 0..10 {
354            let mut tmp = vec![];
355            for _ in 0..=(rng.gen::<u64>() % 20 + 1) {
356                let key = rng.gen::<u64>();
357                let value = rng.gen::<u64>();
358                tmp.push((key, value));
359            }
360            key_to_value.extend(tmp.iter().cloned());
361            map.extend(tmp.iter().cloned());
362        }
363
364        for (key, value) in key_to_value {
365            assert_eq!(map.get(&key).unwrap(), value);
366        }
367    }
368
369    #[test]
370    fn test_debug() {
371        let map: LookupMap<u64, u64> = LookupMap::new(b"m");
372
373        assert_eq!(
374            format!("{:?}", map),
375            format!("LookupMap {{ key_prefix: {:?} }}", map.key_prefix)
376        );
377    }
378}
unc_sdk/collections/lookup_map.rs

unc_sdk/collections/
lookup_map.rs
