unc_sdk/store/lookup_map/
mod.rs

1mod entry;
2mod impls;
3
4use std::borrow::Borrow;
5use std::fmt;
6
7use borsh::{BorshDeserialize, BorshSerialize};
8use once_cell::unsync::OnceCell;
9use unc_sdk_macros::unc;
10
11use super::ERR_NOT_EXIST;
12use crate::store::key::{Identity, ToKey};
13use crate::utils::{EntryState, StableMap};
14use crate::{env, CacheEntry, IntoStorageKey};
15
16pub use entry::{Entry, OccupiedEntry, VacantEntry};
17
18const ERR_ELEMENT_DESERIALIZATION: &str = "Cannot deserialize element";
19const ERR_ELEMENT_SERIALIZATION: &str = "Cannot serialize element";
20
21/// A non-iterable, lazily loaded storage map that stores its content directly on the storage trie.
22///
23/// This map stores the values under a hash of the map's `prefix` and [`BorshSerialize`] of the key
24/// and transformed using the map's [`ToKey`] implementation.
25///
26/// The default hash function for [`LookupMap`] is [`Identity`] which just prefixes the serialized
27/// key object and uses these bytes as the key. This is to be backwards-compatible with
28/// [`collections::LookupMap`](crate::collections::LookupMap) and be fast for small keys.
29/// To use a custom function, use [`with_hasher`]. Alternative builtin hash functions can be found
30/// at [`unc_sdk::store::key`](crate::store::key).
31///
32/// # Examples
33/// ```
34/// use unc_sdk::store::LookupMap;
35///
36/// // Initializes a map, the generic types can be inferred to `LookupMap<String, u8, Identity>`
37/// // The `b"a"` parameter is a prefix for the storage keys of this data structure.
38/// let mut map = LookupMap::new(b"a");
39///
40/// map.set("test".to_string(), Some(7u8));
41/// assert!(map.contains_key("test"));
42/// assert_eq!(map.get("test"), Some(&7u8));
43///
44/// let prev = map.insert("test".to_string(), 5u8);
45/// assert_eq!(prev, Some(7u8));
46/// assert_eq!(map["test"], 5u8);
47/// ```
48///
49/// [`LookupMap`] also implements an [`Entry API`](Self::entry), which allows
50/// for more complex methods of getting, setting, updating and removing keys and
51/// their values:
52///
53/// ```
54/// use unc_sdk::store::LookupMap;
55///
56/// // type inference lets us omit an explicit type signature (which
57/// // would be `LookupMap<String, u8>` in this example).
58/// let mut player_stats = LookupMap::new(b"m");
59///
60/// fn random_stat_buff() -> u8 {
61///     // could actually return some random value here - let's just return
62///     // some fixed value for now
63///     42
64/// }
65///
66/// // insert a key only if it doesn't already exist
67/// player_stats.entry("health".to_string()).or_insert(100);
68///
69/// // insert a key using a function that provides a new value only if it
70/// // doesn't already exist
71/// player_stats.entry("defence".to_string()).or_insert_with(random_stat_buff);
72///
73/// // update a key, guarding against the key possibly not being set
74/// let stat = player_stats.entry("attack".to_string()).or_insert(100);
75/// *stat += random_stat_buff();
76/// ```
77///
78/// [`with_hasher`]: Self::with_hasher
79
80#[unc(inside_uncsdk)]
81pub struct LookupMap<K, V, H = Identity>
82where
83    K: BorshSerialize + Ord,
84    V: BorshSerialize,
85    H: ToKey,
86{
87    prefix: Box<[u8]>,
88    /// Cache for loads and intermediate changes to the underlying vector.
89    /// The cached entries are wrapped in a [`Box`] to avoid existing pointers from being
90    /// invalidated.
91    #[borsh(skip, bound(deserialize = ""))] // removes `core::default::Default` from `K`/`V`
92    cache: StableMap<K, EntryAndHash<V, H::KeyType>>,
93}
94
95struct EntryAndHash<V, T> {
96    value: OnceCell<CacheEntry<V>>,
97    hash: OnceCell<T>,
98}
99
100impl<V, T> Default for EntryAndHash<V, T> {
101    fn default() -> Self {
102        Self { value: Default::default(), hash: Default::default() }
103    }
104}
105
106impl<K, V, H> Drop for LookupMap<K, V, H>
107where
108    K: BorshSerialize + Ord,
109    V: BorshSerialize,
110    H: ToKey,
111{
112    fn drop(&mut self) {
113        self.flush()
114    }
115}
116
117impl<K, V, H> fmt::Debug for LookupMap<K, V, H>
118where
119    K: BorshSerialize + Ord,
120    V: BorshSerialize,
121    H: ToKey,
122{
123    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
124        f.debug_struct("LookupMap").field("prefix", &self.prefix).finish()
125    }
126}
127
128impl<K, V> LookupMap<K, V, Identity>
129where
130    K: BorshSerialize + Ord,
131    V: BorshSerialize,
132{
133    /// Create a new [`LookupMap`] with the prefix provided.
134    ///
135    /// This prefix can be anything that implements [`IntoStorageKey`]. The prefix is used when
136    /// storing and looking up values in storage to ensure no collisions with other collections.
137    ///
138    /// # Examples
139    ///
140    /// ```
141    /// use unc_sdk::store::LookupMap;
142    ///
143    /// let mut map: LookupMap<u32, String> = LookupMap::new(b"m");
144    /// ```
145    #[inline]
146    pub fn new<S>(prefix: S) -> Self
147    where
148        S: IntoStorageKey,
149    {
150        Self::with_hasher(prefix)
151    }
152}
153
154impl<K, V, H> LookupMap<K, V, H>
155where
156    K: BorshSerialize + Ord,
157    V: BorshSerialize,
158    H: ToKey,
159{
160    /// Initialize a [`LookupMap`] with a custom hash function.
161    ///
162    /// # Example
163    /// ```
164    /// use unc_sdk::store::{LookupMap, key::Keccak256};
165    ///
166    /// let map = LookupMap::<String, String, Keccak256>::with_hasher(b"m");
167    /// ```
168    pub fn with_hasher<S>(prefix: S) -> Self
169    where
170        S: IntoStorageKey,
171    {
172        Self { prefix: prefix.into_storage_key().into_boxed_slice(), cache: Default::default() }
173    }
174
175    /// Overwrites the current value for the given key.
176    ///
177    /// This function will not load the existing value from storage and return the value in storage.
178    /// Use [`LookupMap::insert`] if you need the previous value.
179    ///
180    /// Calling `set` with a `None` value will delete the entry from storage.
181    ///
182    /// # Example
183    /// ```
184    /// use unc_sdk::store::LookupMap;
185    ///
186    /// let mut map = LookupMap::new(b"m");
187    ///
188    /// map.set("test".to_string(), Some(7u8));
189    /// assert!(map.contains_key("test"));
190    ///
191    /// //Delete the entry from storage
192    /// map.set("test".to_string(), None);
193    /// assert!(!map.contains_key("test"));
194    /// ```
195    pub fn set(&mut self, key: K, value: Option<V>) {
196        let entry = self.cache.get_mut(key);
197        match entry.value.get_mut() {
198            Some(entry) => *entry.value_mut() = value,
199            None => {
200                let _ = entry.value.set(CacheEntry::new_modified(value));
201            }
202        }
203    }
204}
205
206impl<K, V, H> LookupMap<K, V, H>
207where
208    K: BorshSerialize + Ord,
209    V: BorshSerialize + BorshDeserialize,
210    H: ToKey,
211{
212    fn deserialize_element(bytes: &[u8]) -> V {
213        V::try_from_slice(bytes).unwrap_or_else(|_| env::panic_str(ERR_ELEMENT_DESERIALIZATION))
214    }
215
216    fn load_element<Q: ?Sized>(prefix: &[u8], key: &Q) -> (H::KeyType, Option<V>)
217    where
218        Q: BorshSerialize,
219        K: Borrow<Q>,
220    {
221        let key = H::to_key(prefix, key, &mut Vec::new());
222        let storage_bytes = env::storage_read(key.as_ref());
223        (key, storage_bytes.as_deref().map(Self::deserialize_element))
224    }
225
226    /// Returns a reference to the value corresponding to the key.
227    ///
228    /// The key may be any borrowed form of the map's key type, but
229    /// [`BorshSerialize`] and [`ToOwned<Owned = K>`](ToOwned) on the borrowed form *must* match those for
230    /// the key type.
231    ///
232    /// # Example
233    /// ```
234    /// use unc_sdk::store::LookupMap;
235    ///
236    /// let mut map: LookupMap<u32, String> = LookupMap::new(b"m");
237    ///
238    /// map.insert(1, "a".to_string());
239    /// assert_eq!(map.get(&1), Some(&"a".to_string()));
240    /// assert_eq!(map.get(&2), None);
241    /// ```
242    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
243    where
244        K: Borrow<Q>,
245        Q: BorshSerialize + ToOwned<Owned = K>,
246    {
247        //* ToOwned bound, which forces a clone, is required to be able to keep the key in the cache
248        let cached = self.cache.get(k.to_owned());
249        let entry = cached.value.get_or_init(|| {
250            let (key, element) = Self::load_element(&self.prefix, k);
251            let _ = cached.hash.set(key);
252            CacheEntry::new_cached(element)
253        });
254        entry.value().as_ref()
255    }
256
257    pub(crate) fn get_mut_inner<Q: ?Sized>(&mut self, k: &Q) -> &mut CacheEntry<V>
258    where
259        K: Borrow<Q>,
260        Q: BorshSerialize + ToOwned<Owned = K>,
261    {
262        let prefix = &self.prefix;
263        //* ToOwned bound, which forces a clone, is required to be able to keep the key in the cache
264        let entry = self.cache.get_mut(k.to_owned());
265        entry.value.get_or_init(|| {
266            let (key, value) = Self::load_element(prefix, k);
267            let _ = entry.hash.set(key);
268            CacheEntry::new_cached(value)
269        });
270        let entry = entry.value.get_mut().unwrap_or_else(|| env::abort());
271        entry
272    }
273
274    /// Returns a mutable reference to the value corresponding to the key.
275    ///
276    /// The key may be any borrowed form of the map's key type, but
277    /// [`BorshSerialize`] and [`ToOwned<Owned = K>`](ToOwned) on the borrowed form *must* match those for
278    /// the key type.
279    ///
280    /// # Example
281    /// ```
282    /// use unc_sdk::store::LookupMap;
283    ///
284    /// let mut map: LookupMap<u32, String> = LookupMap::new(b"m");
285    /// map.insert(1, "a".to_string());
286    /// if let Some(x) = map.get_mut(&1) {
287    ///     *x = "b".to_string();
288    ///     assert_eq!(map[&1], "b".to_string());
289    /// }
290    /// ```
291    pub fn get_mut<Q: ?Sized>(&mut self, k: &Q) -> Option<&mut V>
292    where
293        K: Borrow<Q>,
294        Q: BorshSerialize + ToOwned<Owned = K>,
295    {
296        self.get_mut_inner(k).value_mut().as_mut()
297    }
298
299    /// Inserts a key-value pair into the map.
300    ///
301    /// If the map did not have this key present, [`None`] is returned.
302    ///
303    /// If the map did have this key present, the value is updated, and the old
304    /// value is returned. The key is not updated, though; this matters for
305    /// types that can be `==` without being identical.
306    ///
307    /// # Example
308    /// ```
309    /// use unc_sdk::store::LookupMap;
310    ///
311    /// let mut map: LookupMap<u32, String> = LookupMap::new(b"m");
312    /// assert_eq!(map.insert(37, "a".to_string()), None);
313    /// assert_eq!(map.contains_key(&37), true);
314    ///
315    /// map.insert(37, "b".to_string());
316    /// assert_eq!(map.insert(37, "c".to_string()), Some("b".to_string()));
317    /// assert_eq!(map[&37], "c".to_string());
318    /// ```
319    pub fn insert(&mut self, k: K, v: V) -> Option<V>
320    where
321        K: Clone,
322    {
323        self.get_mut_inner(&k).replace(Some(v))
324    }
325
326    /// Returns `true` if the map contains a value for the specified key.
327    ///
328    /// The key may be any borrowed form of the map's key type, but
329    /// [`BorshSerialize`] and [`ToOwned<Owned = K>`](ToOwned) on the borrowed form *must* match those for
330    /// the key type.
331    ///
332    /// # Example
333    /// ```
334    /// use unc_sdk::store::LookupMap;
335    ///
336    /// let mut map: LookupMap<u32, String> = LookupMap::new(b"m");
337    /// map.insert(1, "a".to_string());
338    /// assert_eq!(map.contains_key(&1), true);
339    /// assert_eq!(map.contains_key(&2), false);
340    /// ```
341    pub fn contains_key<Q: ?Sized>(&self, k: &Q) -> bool
342    where
343        K: Borrow<Q>,
344        Q: BorshSerialize + ToOwned<Owned = K> + Ord,
345    {
346        // Check cache before checking storage
347        let contains = self
348            .cache
349            .map_value_ref(k, |v| v.value.get().and_then(|s| s.value().as_ref()).is_some());
350        if let Some(is_some) = contains {
351            return is_some;
352        }
353
354        // Value is not in cache, check if storage has value for given key.
355        let storage_key = H::to_key(&self.prefix, k, &mut Vec::new());
356        let contains = env::storage_has_key(storage_key.as_ref());
357
358        if !contains {
359            // If value not in cache and not in storage, can set a cached `None`
360            let cache = self.cache.get(k.to_owned());
361            let _ = cache.value.set(CacheEntry::new_cached(None));
362            let _ = cache.hash.set(storage_key);
363        }
364        contains
365    }
366
367    /// Removes a key from the map, returning the value at the key if the key
368    /// was previously in the map.
369    ///
370    /// The key may be any borrowed form of the map's key type, but
371    /// [`BorshSerialize`] and [`ToOwned<Owned = K>`](ToOwned) on the borrowed form *must* match those for
372    /// the key type.
373    ///
374    /// # Example
375    /// ```
376    /// use unc_sdk::store::LookupMap;
377    ///
378    /// let mut map: LookupMap<u32, String> = LookupMap::new(b"m");
379    /// map.insert(1, "a".to_string());
380    /// assert_eq!(map.remove(&1), Some("a".to_string()));
381    /// assert_eq!(map.remove(&1), None);
382    /// ```
383    pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
384    where
385        K: Borrow<Q>,
386        Q: BorshSerialize + ToOwned<Owned = K>,
387    {
388        self.get_mut_inner(k).replace(None)
389    }
390
391    /// Gets the given key's corresponding entry in the map for in-place manipulation.
392    /// ```
393    /// use unc_sdk::store::LookupMap;
394    ///
395    /// let mut count = LookupMap::new(b"m");
396    ///
397    /// for ch in [7, 2, 4, 7, 4, 1, 7] {
398    ///     let counter = count.entry(ch).or_insert(0);
399    ///     *counter += 1;
400    /// }
401    ///
402    /// assert_eq!(count[&4], 2);
403    /// assert_eq!(count[&7], 3);
404    /// assert_eq!(count[&1], 1);
405    /// assert_eq!(count.get(&8), None);
406    /// ```
407    pub fn entry(&mut self, key: K) -> Entry<K, V>
408    where
409        K: Clone,
410    {
411        let entry = self.get_mut_inner(&key);
412        if entry.value().is_some() {
413            // Value exists in cache and is `Some`
414            Entry::Occupied(OccupiedEntry { key, entry })
415        } else {
416            // Value exists in cache, but is `None`
417            Entry::Vacant(VacantEntry { key, entry })
418        }
419    }
420}
421
422impl<K, V, H> LookupMap<K, V, H>
423where
424    K: BorshSerialize + Ord,
425    V: BorshSerialize,
426    H: ToKey,
427{
428    /// Flushes the intermediate values of the map before this is called when the structure is
429    /// [`Drop`]ed. This will write all modified values to storage but keep all cached values
430    /// in memory.
431    pub fn flush(&mut self) {
432        let mut buf = Vec::new();
433        for (k, v) in self.cache.inner().iter_mut() {
434            if let Some(val) = v.value.get_mut() {
435                if val.is_modified() {
436                    let prefix = &self.prefix;
437                    let key = v.hash.get_or_init(|| {
438                        buf.clear();
439                        H::to_key(prefix, k, &mut buf)
440                    });
441                    match val.value().as_ref() {
442                        Some(modified) => {
443                            buf.clear();
444                            BorshSerialize::serialize(modified, &mut buf)
445                                .unwrap_or_else(|_| env::panic_str(ERR_ELEMENT_SERIALIZATION));
446                            env::storage_write(key.as_ref(), &buf);
447                        }
448                        None => {
449                            // Element was removed, clear the storage for the value
450                            env::storage_remove(key.as_ref());
451                        }
452                    }
453
454                    // Update state of flushed state as cached, to avoid duplicate writes/removes
455                    // while also keeping the cached values in memory.
456                    val.replace_state(EntryState::Cached);
457                }
458            }
459        }
460    }
461}
462
463#[cfg(not(target_arch = "wasm32"))]
464#[cfg(test)]
465mod tests {
466    use super::LookupMap;
467    use crate::env;
468    use crate::store::key::{Keccak256, ToKey};
469    use crate::test_utils::test_env::setup_free;
470    use arbitrary::{Arbitrary, Unstructured};
471    use rand::seq::SliceRandom;
472    use rand::RngCore;
473    use rand::{Rng, SeedableRng};
474    use std::collections::HashMap;
475
476    #[test]
477    fn test_insert() {
478        let mut map = LookupMap::new(b"m");
479        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(0);
480        for _ in 0..100 {
481            let key = rng.gen::<u64>();
482            let value = rng.gen::<u64>();
483            map.insert(key, value);
484            assert_eq!(*map.get(&key).unwrap(), value);
485        }
486    }
487
488    #[test]
489    fn test_insert_has_key() {
490        let mut map = LookupMap::new(b"m");
491        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(0);
492        let mut key_to_value = HashMap::new();
493        for _ in 0..100 {
494            let key = rng.gen::<u64>();
495            let value = rng.gen::<u64>();
496            map.insert(key, value);
497            key_to_value.insert(key, value);
498        }
499        // Non existing
500        for _ in 0..100 {
501            let key = rng.gen::<u64>();
502            assert_eq!(map.contains_key(&key), key_to_value.contains_key(&key));
503        }
504        // Existing
505        for (key, _) in key_to_value.iter() {
506            assert!(map.contains_key(key));
507        }
508    }
509
510    #[test]
511    fn test_insert_remove() {
512        let mut map = LookupMap::new(b"m");
513        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(1);
514        let mut keys = vec![];
515        let mut key_to_value = HashMap::new();
516        for _ in 0..100 {
517            let key = rng.gen::<u64>();
518            let value = rng.gen::<u64>();
519            keys.push(key);
520            key_to_value.insert(key, value);
521            map.insert(key, value);
522        }
523        keys.shuffle(&mut rng);
524        for key in keys {
525            let actual = map.remove(&key).unwrap();
526            assert_eq!(actual, key_to_value[&key]);
527        }
528    }
529
530    #[test]
531    fn test_remove_last_reinsert() {
532        let mut map = LookupMap::new(b"m");
533        let key1 = 1u64;
534        let value1 = 2u64;
535        map.insert(key1, value1);
536        let key2 = 3u64;
537        let value2 = 4u64;
538        map.insert(key2, value2);
539
540        let actual_value2 = map.remove(&key2).unwrap();
541        assert_eq!(actual_value2, value2);
542
543        let actual_insert_value2 = map.insert(key2, value2);
544        assert_eq!(actual_insert_value2, None);
545    }
546
547    #[test]
548    fn test_insert_override_remove() {
549        let mut map = LookupMap::new(b"m");
550        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(2);
551        let mut keys = vec![];
552        let mut key_to_value = HashMap::new();
553        for _ in 0..100 {
554            let key = rng.gen::<u64>();
555            let value = rng.gen::<u64>();
556            keys.push(key);
557            key_to_value.insert(key, value);
558            map.insert(key, value);
559        }
560        keys.shuffle(&mut rng);
561        for key in &keys {
562            let value = rng.gen::<u64>();
563            let actual = map.insert(*key, value).unwrap();
564            assert_eq!(actual, key_to_value[key]);
565            key_to_value.insert(*key, value);
566        }
567        keys.shuffle(&mut rng);
568        for key in keys {
569            let actual = map.remove(&key).unwrap();
570            assert_eq!(actual, key_to_value[&key]);
571        }
572    }
573
574    #[test]
575    fn test_get_non_existent() {
576        let mut map = LookupMap::new(b"m");
577        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(3);
578        let mut key_to_value = HashMap::new();
579        for _ in 0..500 {
580            let key = rng.gen::<u64>() % 20_000;
581            let value = rng.gen::<u64>();
582            key_to_value.insert(key, value);
583            map.insert(key, value);
584        }
585        for _ in 0..500 {
586            let key = rng.gen::<u64>() % 20_000;
587            assert_eq!(map.get(&key), key_to_value.get(&key));
588        }
589    }
590
591    #[test]
592    fn size_of_map() {
593        assert_eq!(core::mem::size_of::<LookupMap<u8, u8>>(), 48);
594    }
595
596    #[test]
597    fn identity_compat_v1() {
598        use crate::collections::LookupMap as LM1;
599
600        let mut lm1 = LM1::new(b"m");
601        lm1.insert(&8u8, &"Some value".to_string());
602        lm1.insert(&0, &"Other".to_string());
603        assert_eq!(lm1.get(&8), Some("Some value".to_string()));
604
605        let mut lm2 = LookupMap::new(b"m");
606        assert_eq!(lm2.get(&8u8), Some(&"Some value".to_string()));
607        assert_eq!(lm2.remove(&0), Some("Other".to_string()));
608        *lm2.get_mut(&8).unwrap() = "New".to_string();
609        lm2.flush();
610
611        assert!(!lm1.contains_key(&0));
612        assert_eq!(lm1.get(&8), Some("New".to_string()));
613    }
614
615    #[test]
616    fn test_extend() {
617        let mut map = LookupMap::new(b"m");
618        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(4);
619        let mut key_to_value = HashMap::new();
620        for _ in 0..100 {
621            let key = rng.gen::<u64>();
622            let value = rng.gen::<u64>();
623            key_to_value.insert(key, value);
624            map.insert(key, value);
625        }
626        for _ in 0..10 {
627            let mut tmp = vec![];
628            for _ in 0..=(rng.gen::<u64>() % 20 + 1) {
629                let key = rng.gen::<u64>();
630                let value = rng.gen::<u64>();
631                tmp.push((key, value));
632            }
633            key_to_value.extend(tmp.iter().cloned());
634            map.extend(tmp.iter().cloned());
635        }
636
637        for (key, value) in key_to_value {
638            assert_eq!(*map.get(&key).unwrap(), value);
639        }
640    }
641
642    #[test]
643    fn flush_on_drop() {
644        let mut map = LookupMap::<_, _, Keccak256>::with_hasher(b"m");
645
646        // Set a value, which does not write to storage yet
647        map.set(5u8, Some(8u8));
648
649        // Create duplicate which references same data
650        assert_eq!(map[&5], 8);
651
652        let storage_key = Keccak256::to_key(b"m", &5, &mut Vec::new());
653        assert!(!env::storage_has_key(&storage_key));
654
655        drop(map);
656
657        let dup_map = LookupMap::<u8, u8, Keccak256>::with_hasher(b"m");
658
659        // New map can now load the value
660        assert_eq!(dup_map[&5], 8);
661    }
662
663    #[derive(Arbitrary, Debug)]
664    enum Op {
665        Insert(u8, u8),
666        Set(u8, Option<u8>),
667        Remove(u8),
668        Flush,
669        Restore,
670        Get(u8),
671    }
672
673    #[test]
674    fn arbitrary() {
675        setup_free();
676
677        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(0);
678        let mut buf = vec![0; 4096];
679        for _ in 0..512 {
680            // Clear storage in-between runs
681            crate::mock::with_mocked_blockchain(|b| b.take_storage());
682            rng.fill_bytes(&mut buf);
683
684            let mut lm = LookupMap::new(b"l");
685            let mut hm = HashMap::new();
686            let u = Unstructured::new(&buf);
687            if let Ok(ops) = Vec::<Op>::arbitrary_take_rest(u) {
688                for op in ops {
689                    match op {
690                        Op::Insert(k, v) => {
691                            let r1 = lm.insert(k, v);
692                            let r2 = hm.insert(k, v);
693                            assert_eq!(r1, r2)
694                        }
695                        Op::Set(k, v) => {
696                            lm.set(k, v);
697
698                            if let Some(val) = v {
699                                hm.insert(k, val);
700                            } else {
701                                hm.remove(&k);
702                            }
703
704                            // Extra get just to make sure set happened correctly
705                            assert_eq!(lm.get(&k), hm.get(&k));
706                        }
707                        Op::Remove(k) => {
708                            let r1 = lm.remove(&k);
709                            let r2 = hm.remove(&k);
710                            assert_eq!(r1, r2)
711                        }
712                        Op::Flush => {
713                            lm.flush();
714                        }
715                        Op::Restore => {
716                            lm = LookupMap::new(b"l");
717                        }
718                        Op::Get(k) => {
719                            let r1 = lm.get(&k);
720                            let r2 = hm.get(&k);
721                            assert_eq!(r1, r2)
722                        }
723                    }
724                }
725            }
726        }
727    }
728}
729
730// Hashbrown-like tests.
731#[cfg(test)]
732mod test_map {
733    use super::Entry::{Occupied, Vacant};
734    use crate::store::LookupMap;
735    use borsh::{BorshDeserialize, BorshSerialize};
736    use std::cell::RefCell;
737    use std::usize;
738    use std::vec::Vec;
739
740    #[test]
741    fn test_insert() {
742        let mut m = LookupMap::new(b"b");
743        assert!(m.insert(1, 2).is_none());
744        assert!(m.insert(2, 4).is_none());
745        assert_eq!(*m.get(&1).unwrap(), 2);
746        assert_eq!(*m.get(&2).unwrap(), 4);
747    }
748
749    thread_local! { static DROP_VECTOR: RefCell<Vec<i32>> = const { RefCell::new(Vec::new()) }}
750
751    #[derive(Hash, PartialEq, Eq, BorshSerialize, BorshDeserialize, PartialOrd, Ord)]
752    struct Droppable {
753        k: usize,
754    }
755
756    impl Droppable {
757        fn new(k: usize) -> Droppable {
758            DROP_VECTOR.with(|slot| {
759                slot.borrow_mut()[k] += 1;
760            });
761
762            Droppable { k }
763        }
764    }
765
766    impl Drop for Droppable {
767        fn drop(&mut self) {
768            DROP_VECTOR.with(|slot| {
769                slot.borrow_mut()[self.k] -= 1;
770            });
771        }
772    }
773
774    impl Clone for Droppable {
775        fn clone(&self) -> Self {
776            Droppable::new(self.k)
777        }
778    }
779
780    #[test]
781    fn test_drops() {
782        DROP_VECTOR.with(|slot| {
783            *slot.borrow_mut() = vec![0; 200];
784        });
785
786        {
787            let mut m = LookupMap::new(b"b");
788
789            DROP_VECTOR.with(|v| {
790                for i in 0..200 {
791                    assert_eq!(v.borrow()[i], 0);
792                }
793            });
794
795            for i in 0..100 {
796                let d1 = Droppable::new(i);
797                let d2 = Droppable::new(i + 100);
798                m.insert(d1, d2);
799            }
800
801            DROP_VECTOR.with(|v| {
802                for i in 0..100 {
803                    assert_eq!(v.borrow()[i], 1);
804                }
805            });
806
807            for i in 0..50 {
808                let k = Droppable::new(i);
809                let v = m.remove(&k);
810
811                assert!(v.is_some());
812
813                DROP_VECTOR.with(|v| {
814                    assert_eq!(v.borrow()[i], 2);
815                    assert_eq!(v.borrow()[i + 100], 1);
816                });
817            }
818
819            DROP_VECTOR.with(|v| {
820                for i in 0..50 {
821                    assert_eq!(v.borrow()[i], 1);
822                    assert_eq!(v.borrow()[i + 100], 0);
823                }
824
825                for i in 50..100 {
826                    assert_eq!(v.borrow()[i], 1);
827                    assert_eq!(v.borrow()[i + 100], 1);
828                }
829            });
830        }
831
832        DROP_VECTOR.with(|v| {
833            for i in 0..200 {
834                assert_eq!(v.borrow()[i], 0);
835            }
836        });
837    }
838
839    #[test]
840    fn test_empty_remove() {
841        let mut m: LookupMap<i32, bool> = LookupMap::new(b"b");
842        assert_eq!(m.remove(&0), None);
843    }
844
845    #[test]
846    fn test_empty_entry() {
847        let mut m: LookupMap<i32, bool> = LookupMap::new(b"b");
848        match m.entry(0) {
849            Occupied(_) => panic!(),
850            Vacant(_) => {}
851        }
852        assert!(*m.entry(0).or_insert(true));
853    }
854
855    #[test]
856    #[cfg_attr(miri, ignore)] // FIXME: takes too long
857    fn test_lots_of_insertions() {
858        let mut m = LookupMap::new(b"b");
859
860        // Try this a few times to make sure we never screw up the LookupMap's
861        // internal state.
862        for _ in 0..10 {
863            for i in 1..1001 {
864                assert!(m.insert(i, i).is_none());
865
866                for j in 1..=i {
867                    let r = m.get(&j);
868                    assert_eq!(r, Some(&j));
869                }
870
871                for j in i + 1..1001 {
872                    let r = m.get(&j);
873                    assert_eq!(r, None);
874                }
875            }
876
877            for i in 1001..2001 {
878                assert!(!m.contains_key(&i));
879            }
880
881            // remove forwards
882            for i in 1..1001 {
883                assert!(m.remove(&i).is_some());
884
885                for j in 1..=i {
886                    assert!(!m.contains_key(&j));
887                }
888
889                for j in i + 1..1001 {
890                    assert!(m.contains_key(&j));
891                }
892            }
893
894            for i in 1..1001 {
895                assert!(!m.contains_key(&i));
896            }
897
898            for i in 1..1001 {
899                assert!(m.insert(i, i).is_none());
900            }
901
902            // remove backwards
903            for i in (1..1001).rev() {
904                assert!(m.remove(&i).is_some());
905
906                for j in i..1001 {
907                    assert!(!m.contains_key(&j));
908                }
909
910                for j in 1..i {
911                    assert!(m.contains_key(&j));
912                }
913            }
914        }
915    }
916
917    #[test]
918    fn test_find_mut() {
919        let mut m = LookupMap::new(b"b");
920        assert!(m.insert(1, 12).is_none());
921        assert!(m.insert(2, 8).is_none());
922        assert!(m.insert(5, 14).is_none());
923        let new = 100;
924        match m.get_mut(&5) {
925            None => panic!(),
926            Some(x) => *x = new,
927        }
928        assert_eq!(m.get(&5), Some(&new));
929    }
930
931    #[test]
932    fn test_insert_overwrite() {
933        let mut m = LookupMap::new(b"b");
934        assert!(m.insert(1, 2).is_none());
935        assert_eq!(*m.get(&1).unwrap(), 2);
936        assert!(m.insert(1, 3).is_some());
937        assert_eq!(*m.get(&1).unwrap(), 3);
938    }
939
940    #[test]
941    fn test_remove() {
942        let mut m = LookupMap::new(b"b");
943        m.insert(1, 2);
944        assert_eq!(m.remove(&1), Some(2));
945        assert_eq!(m.remove(&1), None);
946    }
947
948    #[test]
949    fn test_find() {
950        let mut m = LookupMap::new(b"b");
951        assert!(m.get(&1).is_none());
952        m.insert(1, 2);
953        match m.get(&1) {
954            None => panic!(),
955            Some(v) => assert_eq!(*v, 2),
956        }
957    }
958
959    #[test]
960    fn test_show() {
961        let mut map = LookupMap::new(b"b");
962        let empty: LookupMap<i32, i32> = LookupMap::new(b"c");
963
964        map.insert(1, 2);
965        map.insert(3, 4);
966
967        let map_str = format!("{:?}", map);
968
969        assert_eq!(map_str, "LookupMap { prefix: [98] }");
970        assert_eq!(format!("{:?}", empty), "LookupMap { prefix: [99] }");
971    }
972
973    #[test]
974    fn test_index() {
975        let mut map = LookupMap::new(b"b");
976
977        map.insert(1, 2);
978        map.insert(2, 1);
979        map.insert(3, 4);
980
981        assert_eq!(map[&2], 1);
982    }
983
984    #[test]
985    #[should_panic]
986    #[allow(clippy::unnecessary_operation)]
987    fn test_index_nonexistent() {
988        let mut map = LookupMap::new(b"b");
989
990        map.insert(1, 2);
991        map.insert(2, 1);
992        map.insert(3, 4);
993
994        #[allow(clippy::no_effect)] // false positive lint
995        map[&4];
996    }
997
998    #[test]
999    fn test_entry() {
1000        let mut map = LookupMap::new(b"b");
1001
1002        let xs = [(1, 10), (2, 20), (3, 30), (4, 40), (5, 50), (6, 60)];
1003
1004        for v in xs {
1005            map.insert(v.0, v.1);
1006        }
1007
1008        // Existing key (insert)
1009        match map.entry(1) {
1010            Vacant(_) => unreachable!(),
1011            Occupied(mut view) => {
1012                assert_eq!(view.get(), &10);
1013                assert_eq!(view.insert(100), 10);
1014            }
1015        }
1016        assert_eq!(map.get(&1).unwrap(), &100);
1017
1018        // Existing key (update)
1019        match map.entry(2) {
1020            Vacant(_) => unreachable!(),
1021            Occupied(mut view) => {
1022                let v = view.get_mut();
1023                let new_v = (*v) * 10;
1024                *v = new_v;
1025            }
1026        }
1027        assert_eq!(map.get(&2).unwrap(), &200);
1028
1029        // Existing key (take)
1030        match map.entry(3) {
1031            Vacant(_) => unreachable!(),
1032            Occupied(view) => {
1033                assert_eq!(view.remove(), 30);
1034            }
1035        }
1036        assert_eq!(map.get(&3), None);
1037
1038        // Inexistent key (insert)
1039        match map.entry(10) {
1040            Occupied(_) => unreachable!(),
1041            Vacant(view) => {
1042                assert_eq!(*view.insert(1000), 1000);
1043            }
1044        }
1045        assert_eq!(map.get(&10).unwrap(), &1000);
1046    }
1047
1048    #[test]
1049    fn test_extend_ref_kv_tuple() {
1050        let mut a = LookupMap::new(b"b");
1051        a.insert(0, 0);
1052
1053        let for_iter: Vec<(i32, i32)> = (0..100).map(|i| (i, i)).collect();
1054        a.extend(for_iter);
1055
1056        for item in 0..100 {
1057            assert_eq!(a[&item], item);
1058        }
1059    }
1060
1061    #[test]
1062    fn test_occupied_entry_key() {
1063        let mut a = LookupMap::new(b"b");
1064        let key = "hello there";
1065        let value = "value goes here";
1066        a.insert(key.to_string(), value.to_string());
1067        assert_eq!(a[key], value);
1068
1069        match a.entry(key.to_string()) {
1070            Vacant(_) => panic!(),
1071            Occupied(e) => assert_eq!(key, *e.key()),
1072        }
1073        assert_eq!(a[key], value);
1074    }
1075
1076    #[test]
1077    fn test_vacant_entry_key() {
1078        let mut a = LookupMap::new(b"b");
1079        let key = "hello there";
1080        let value = "value goes here".to_string();
1081
1082        match a.entry(key.to_string()) {
1083            Occupied(_) => panic!(),
1084            Vacant(e) => {
1085                assert_eq!(key, *e.key());
1086                e.insert(value.clone());
1087            }
1088        }
1089        assert_eq!(a[key], value);
1090    }
1091}
unc_sdk/store/lookup_map/mod.rs

unc_sdk/store/lookup_map/
mod.rs
