near_sdk/collections/unordered_map/
mod.rs

1//! A map implemented on a trie. Unlike `std::collections::HashMap` the keys in this map are not
2//! hashed but are instead serialized.
3
4mod iter;
5pub use iter::Iter;
6
7use crate::collections::{append, append_slice, Vector};
8use crate::{env, IntoStorageKey};
9use borsh::{to_vec, BorshDeserialize, BorshSerialize};
10use near_sdk_macros::near;
11use std::mem::size_of;
12
13const ERR_INCONSISTENT_STATE: &str = "The collection is an inconsistent state. Did previous smart contract execution terminate unexpectedly?";
14const ERR_KEY_SERIALIZATION: &str = "Cannot serialize key with Borsh";
15const ERR_VALUE_DESERIALIZATION: &str = "Cannot deserialize value with Borsh";
16const ERR_VALUE_SERIALIZATION: &str = "Cannot serialize value with Borsh";
17
18/// An iterable implementation of a map that stores its content directly on the trie.
19#[near(inside_nearsdk)]
20pub struct UnorderedMap<K, V> {
21    key_index_prefix: Vec<u8>,
22    // ser/de is independent of `K` ser/de, `BorshSerialize`/`BorshDeserialize`/`BorshSchema` bounds removed
23    #[cfg_attr(not(feature = "abi"), borsh(bound(serialize = "", deserialize = "")))]
24    #[cfg_attr(
25        feature = "abi",
26        borsh(bound(serialize = "", deserialize = ""), schema(params = ""))
27    )]
28    keys: Vector<K>,
29    // ser/de is independent of `V` ser/de, `BorshSerialize`/`BorshDeserialize`/`BorshSchema` bounds removed
30    #[cfg_attr(not(feature = "abi"), borsh(bound(serialize = "", deserialize = "")))]
31    #[cfg_attr(
32        feature = "abi",
33        borsh(bound(serialize = "", deserialize = ""), schema(params = ""))
34    )]
35    values: Vector<V>,
36}
37
38impl<K, V> UnorderedMap<K, V> {
39    /// Returns the number of elements in the map, also referred to as its size.
40    ///
41    /// # Examples
42    ///
43    /// ```
44    /// use near_sdk::collections::UnorderedMap;
45    ///
46    /// let mut map: UnorderedMap<u8, u8> = UnorderedMap::new(b"m");
47    /// assert_eq!(map.len(), 0);
48    /// map.insert(&1, &1);
49    /// map.insert(&2, &2);
50    /// assert_eq!(map.len(), 2);
51    /// ```
52    pub fn len(&self) -> u64 {
53        let keys_len = self.keys.len();
54        let values_len = self.values.len();
55        if keys_len != values_len {
56            env::panic_str(ERR_INCONSISTENT_STATE)
57        } else {
58            keys_len
59        }
60    }
61
62    /// Returns `true` if the map contains no elements.
63    pub fn is_empty(&self) -> bool {
64        let keys_is_empty = self.keys.is_empty();
65        let values_is_empty = self.values.is_empty();
66        if keys_is_empty != values_is_empty {
67            env::panic_str(ERR_INCONSISTENT_STATE)
68        } else {
69            keys_is_empty
70        }
71    }
72
73    /// Create new map with zero elements. Use `prefix` as a unique identifier.
74    ///
75    /// # Examples
76    ///
77    /// ```
78    /// use near_sdk::collections::UnorderedMap;
79    /// let mut map: UnorderedMap<u8, u8> = UnorderedMap::new(b"m");
80    /// ```
81    pub fn new<S>(prefix: S) -> Self
82    where
83        S: IntoStorageKey,
84    {
85        let prefix = prefix.into_storage_key();
86        let key_index_prefix = append(&prefix, b'i');
87        let index_key_id = append(&prefix, b'k');
88        let index_value_id = append(&prefix, b'v');
89
90        Self {
91            key_index_prefix,
92            keys: Vector::new(index_key_id),
93            values: Vector::new(index_value_id),
94        }
95    }
96
97    fn serialize_index(index: u64) -> [u8; size_of::<u64>()] {
98        index.to_le_bytes()
99    }
100
101    fn deserialize_index(raw_index: &[u8]) -> u64 {
102        let mut result = [0u8; size_of::<u64>()];
103        result.copy_from_slice(raw_index);
104        u64::from_le_bytes(result)
105    }
106
107    fn raw_key_to_index_lookup(&self, raw_key: &[u8]) -> Vec<u8> {
108        append_slice(&self.key_index_prefix, raw_key)
109    }
110
111    /// Returns an index of the given raw key.
112    fn get_index_raw(&self, key_raw: &[u8]) -> Option<u64> {
113        let index_lookup = self.raw_key_to_index_lookup(key_raw);
114        env::storage_read(&index_lookup).map(|raw_index| Self::deserialize_index(&raw_index))
115    }
116
117    /// Returns the serialized value corresponding to the serialized key.
118    fn get_raw(&self, key_raw: &[u8]) -> Option<Vec<u8>> {
119        self.get_index_raw(key_raw).map(|index| match self.values.get_raw(index) {
120            Some(x) => x,
121            None => env::panic_str(ERR_INCONSISTENT_STATE),
122        })
123    }
124
125    /// Inserts a serialized key-value pair into the map.
126    /// If the map did not have this key present, `None` is returned. Otherwise returns
127    /// a serialized value. Note, the keys that have the same hash value are undistinguished by
128    /// the implementation.
129    pub fn insert_raw(&mut self, key_raw: &[u8], value_raw: &[u8]) -> Option<Vec<u8>> {
130        let index_lookup = self.raw_key_to_index_lookup(key_raw);
131        match env::storage_read(&index_lookup) {
132            Some(index_raw) => {
133                // The element already exists.
134                let index = Self::deserialize_index(&index_raw);
135                Some(self.values.replace_raw(index, value_raw))
136            }
137            None => {
138                // The element does not exist yet.
139                let next_index = self.len();
140                let next_index_raw = Self::serialize_index(next_index);
141                env::storage_write(&index_lookup, &next_index_raw);
142                self.keys.push_raw(key_raw);
143                self.values.push_raw(value_raw);
144                None
145            }
146        }
147    }
148
149    /// Removes a serialized key from the map, returning the serialized value at the key if the key
150    /// was previously in the map.
151    pub fn remove_raw(&mut self, key_raw: &[u8]) -> Option<Vec<u8>> {
152        let index_lookup = self.raw_key_to_index_lookup(key_raw);
153        match env::storage_read(&index_lookup) {
154            Some(index_raw) => {
155                #[allow(clippy::branches_sharing_code)]
156                if self.len() == 1 {
157                    // If there is only one element then swap remove simply removes it without
158                    // swapping with the last element.
159                    env::storage_remove(&index_lookup);
160                } else {
161                    // If there is more than one element then swap remove swaps it with the last
162                    // element.
163                    let last_key_raw = match self.keys.get_raw(self.len() - 1) {
164                        Some(x) => x,
165                        None => env::panic_str(ERR_INCONSISTENT_STATE),
166                    };
167                    env::storage_remove(&index_lookup);
168                    // If the removed element was the last element from keys, then we don't need to
169                    // reinsert the lookup back.
170                    if last_key_raw != key_raw {
171                        let last_lookup_key = self.raw_key_to_index_lookup(&last_key_raw);
172                        env::storage_write(&last_lookup_key, &index_raw);
173                    }
174                }
175                let index = Self::deserialize_index(&index_raw);
176                self.keys.swap_remove_raw(index);
177                Some(self.values.swap_remove_raw(index))
178            }
179            None => None,
180        }
181    }
182}
183
184impl<K, V> UnorderedMap<K, V>
185where
186    K: BorshSerialize + BorshDeserialize,
187    V: BorshSerialize + BorshDeserialize,
188{
189    fn serialize_key(key: &K) -> Vec<u8> {
190        match to_vec(key) {
191            Ok(x) => x,
192            Err(_) => env::panic_str(ERR_KEY_SERIALIZATION),
193        }
194    }
195
196    fn deserialize_value(raw_value: &[u8]) -> V {
197        match V::try_from_slice(raw_value) {
198            Ok(x) => x,
199            Err(_) => env::panic_str(ERR_VALUE_DESERIALIZATION),
200        }
201    }
202
203    fn serialize_value(value: &V) -> Vec<u8> {
204        match to_vec(value) {
205            Ok(x) => x,
206            Err(_) => env::panic_str(ERR_VALUE_SERIALIZATION),
207        }
208    }
209
210    /// Returns the value corresponding to the key.
211    ///
212    /// # Examples
213    ///
214    /// ```
215    /// use near_sdk::collections::UnorderedMap;
216    ///
217    /// let mut map: UnorderedMap<u8, u8> = UnorderedMap::new(b"m");
218    /// assert_eq!(map.get(&1), None);
219    /// map.insert(&1, &10);
220    /// assert_eq!(map.get(&1), Some(10));
221    /// ```
222    pub fn get(&self, key: &K) -> Option<V> {
223        self.get_raw(&Self::serialize_key(key)).map(|value_raw| Self::deserialize_value(&value_raw))
224    }
225
226    /// Removes a key from the map, returning the value at the key if the key was previously in the
227    /// map.
228    ///
229    /// # Examples
230    ///
231    /// ```
232    /// use near_sdk::collections::UnorderedMap;
233    ///
234    /// let mut map: UnorderedMap<u8, u8> = UnorderedMap::new(b"m");
235    /// assert_eq!(map.remove(&1), None);
236    /// map.insert(&1, &10);
237    /// assert_eq!(map.remove(&1), Some(10));
238    /// assert_eq!(map.len(), 0);
239    /// ```
240    pub fn remove(&mut self, key: &K) -> Option<V> {
241        self.remove_raw(&Self::serialize_key(key))
242            .map(|value_raw| Self::deserialize_value(&value_raw))
243    }
244
245    /// Inserts a key-value pair into the map.
246    /// If the map did not have this key present, `None` is returned. Otherwise returns
247    /// a value. Note, the keys that have the same hash value are undistinguished by
248    /// the implementation.
249    ///
250    /// # Examples
251    ///
252    /// ```
253    /// use near_sdk::collections::UnorderedMap;
254    ///
255    /// let mut map: UnorderedMap<u8, u8> = UnorderedMap::new(b"m");
256    /// map.insert(&1, &10);
257    /// assert_eq!(map.get(&1), Some(10));
258    /// assert_eq!(map.len(), 1);
259    /// ```
260    pub fn insert(&mut self, key: &K, value: &V) -> Option<V> {
261        self.insert_raw(&Self::serialize_key(key), &Self::serialize_value(value))
262            .map(|value_raw| Self::deserialize_value(&value_raw))
263    }
264
265    /// Clears the map, removing all elements.
266    ///
267    /// # Examples
268    ///
269    /// ```
270    /// use near_sdk::collections::UnorderedMap;
271    ///
272    /// let mut map: UnorderedMap<u8, u8> = UnorderedMap::new(b"m");
273    /// map.insert(&1, &10);
274    /// map.insert(&2, &20);
275    /// map.clear();
276    /// assert_eq!(map.len(), 0);
277    /// ```
278    pub fn clear(&mut self) {
279        for raw_key in self.keys.iter_raw() {
280            let index_lookup = self.raw_key_to_index_lookup(&raw_key);
281            env::storage_remove(&index_lookup);
282        }
283        self.keys.clear();
284        self.values.clear();
285    }
286
287    /// Copies elements into an `std::vec::Vec`.
288    pub fn to_vec(&self) -> std::vec::Vec<(K, V)> {
289        self.iter().collect()
290    }
291
292    /// An iterator visiting all keys. The iterator element type is `K`.
293    pub fn keys(&self) -> impl Iterator<Item = K> + '_ {
294        self.keys.iter()
295    }
296
297    /// An iterator visiting all values. The iterator element type is `V`.
298    pub fn values(&self) -> impl Iterator<Item = V> + '_ {
299        self.values.iter()
300    }
301
302    /// Iterate over deserialized keys and values.
303    pub fn iter(&self) -> Iter<K, V> {
304        Iter::new(self)
305    }
306
307    pub fn extend<IT: IntoIterator<Item = (K, V)>>(&mut self, iter: IT) {
308        for (el_key, el_value) in iter {
309            self.insert(&el_key, &el_value);
310        }
311    }
312
313    /// Returns a view of keys as a vector.
314    /// It's sometimes useful to have random access to the keys.
315    pub fn keys_as_vector(&self) -> &Vector<K> {
316        &self.keys
317    }
318
319    /// Returns a view of values as a vector.
320    /// It's sometimes useful to have random access to the values.
321    pub fn values_as_vector(&self) -> &Vector<V> {
322        &self.values
323    }
324}
325
326impl<K, V> std::fmt::Debug for UnorderedMap<K, V>
327where
328    K: std::fmt::Debug + BorshSerialize + BorshDeserialize,
329    V: std::fmt::Debug + BorshSerialize + BorshDeserialize,
330{
331    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
332        f.debug_struct("UnorderedMap")
333            .field("key_index_prefix", &self.key_index_prefix)
334            .field("keys", &self.keys)
335            .field("values", &self.values)
336            .finish()
337    }
338}
339
340#[cfg(not(target_arch = "wasm32"))]
341#[cfg(test)]
342mod tests {
343    use crate::collections::UnorderedMap;
344    use borsh::{BorshDeserialize, BorshSerialize};
345    use rand::seq::SliceRandom;
346    use rand::{Rng, SeedableRng};
347    use std::collections::{HashMap, HashSet};
348    use std::iter::FromIterator;
349    use std::sync::atomic::{AtomicUsize, Ordering};
350
351    #[test]
352    pub fn test_insert_one() {
353        let mut map = UnorderedMap::new(b"m");
354        assert_eq!(None, map.insert(&1, &2));
355        assert_eq!(2, map.insert(&1, &3).unwrap());
356    }
357
358    #[test]
359    pub fn test_insert() {
360        let mut map = UnorderedMap::new(b"m");
361        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(0);
362        for _ in 0..500 {
363            let key = rng.gen::<u64>();
364            let value = rng.gen::<u64>();
365            map.insert(&key, &value);
366        }
367    }
368
369    #[test]
370    pub fn test_insert_remove() {
371        let mut map = UnorderedMap::new(b"m");
372        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(1);
373        let mut keys = vec![];
374        let mut key_to_value = HashMap::new();
375        for _ in 0..100 {
376            let key = rng.gen::<u64>();
377            let value = rng.gen::<u64>();
378            keys.push(key);
379            key_to_value.insert(key, value);
380            map.insert(&key, &value);
381        }
382        keys.shuffle(&mut rng);
383        for key in keys {
384            let actual = map.remove(&key).unwrap();
385            assert_eq!(actual, key_to_value[&key]);
386        }
387    }
388
389    #[test]
390    pub fn test_remove_last_reinsert() {
391        let mut map = UnorderedMap::new(b"m");
392        let key1 = 1u64;
393        let value1 = 2u64;
394        map.insert(&key1, &value1);
395        let key2 = 3u64;
396        let value2 = 4u64;
397        map.insert(&key2, &value2);
398
399        let actual_value2 = map.remove(&key2).unwrap();
400        assert_eq!(actual_value2, value2);
401
402        let actual_insert_value2 = map.insert(&key2, &value2);
403        assert_eq!(actual_insert_value2, None);
404    }
405
406    #[test]
407    pub fn test_insert_override_remove() {
408        let mut map = UnorderedMap::new(b"m");
409        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(2);
410        let mut keys = vec![];
411        let mut key_to_value = HashMap::new();
412        for _ in 0..100 {
413            let key = rng.gen::<u64>();
414            let value = rng.gen::<u64>();
415            keys.push(key);
416            key_to_value.insert(key, value);
417            map.insert(&key, &value);
418        }
419        keys.shuffle(&mut rng);
420        for key in &keys {
421            let value = rng.gen::<u64>();
422            let actual = map.insert(key, &value).unwrap();
423            assert_eq!(actual, key_to_value[key]);
424            key_to_value.insert(*key, value);
425        }
426        keys.shuffle(&mut rng);
427        for key in keys {
428            let actual = map.remove(&key).unwrap();
429            assert_eq!(actual, key_to_value[&key]);
430        }
431    }
432
433    #[test]
434    pub fn test_get_non_existent() {
435        let mut map = UnorderedMap::new(b"m");
436        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(3);
437        let mut key_to_value = HashMap::new();
438        for _ in 0..500 {
439            let key = rng.gen::<u64>() % 20_000;
440            let value = rng.gen::<u64>();
441            key_to_value.insert(key, value);
442            map.insert(&key, &value);
443        }
444        for _ in 0..500 {
445            let key = rng.gen::<u64>() % 20_000;
446            assert_eq!(map.get(&key), key_to_value.get(&key).cloned());
447        }
448    }
449
450    #[test]
451    pub fn test_to_vec() {
452        let mut map = UnorderedMap::new(b"m");
453        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(4);
454        let mut key_to_value = HashMap::new();
455        for _ in 0..400 {
456            let key = rng.gen::<u64>();
457            let value = rng.gen::<u64>();
458            key_to_value.insert(key, value);
459            map.insert(&key, &value);
460        }
461        let actual = HashMap::from_iter(map.to_vec());
462        assert_eq!(actual, key_to_value);
463    }
464
465    #[test]
466    pub fn test_clear() {
467        let mut map = UnorderedMap::new(b"m");
468        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(5);
469        for _ in 0..10 {
470            for _ in 0..=(rng.gen::<u64>() % 20 + 1) {
471                let key = rng.gen::<u64>();
472                let value = rng.gen::<u64>();
473                map.insert(&key, &value);
474            }
475            assert!(!map.to_vec().is_empty());
476            map.clear();
477            assert!(map.to_vec().is_empty());
478        }
479    }
480
481    #[test]
482    pub fn test_keys_values() {
483        let mut map = UnorderedMap::new(b"m");
484        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(4);
485        let mut key_to_value = HashMap::new();
486        for _ in 0..400 {
487            let key = rng.gen::<u64>();
488            let value = rng.gen::<u64>();
489            key_to_value.insert(key, value);
490            map.insert(&key, &value);
491        }
492        let actual: HashMap<u64, u64> = HashMap::from_iter(map.to_vec());
493        assert_eq!(
494            actual.keys().collect::<HashSet<_>>(),
495            key_to_value.keys().collect::<HashSet<_>>()
496        );
497        assert_eq!(
498            actual.values().collect::<HashSet<_>>(),
499            key_to_value.values().collect::<HashSet<_>>()
500        );
501    }
502
503    #[test]
504    pub fn test_iter() {
505        let mut map = UnorderedMap::new(b"m");
506        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(4);
507        let mut key_to_value = HashMap::new();
508        for _ in 0..400 {
509            let key = rng.gen::<u64>();
510            let value = rng.gen::<u64>();
511            key_to_value.insert(key, value);
512            map.insert(&key, &value);
513        }
514        let actual: HashMap<u64, u64> = map.iter().collect();
515        assert_eq!(actual, key_to_value);
516    }
517
518    #[test]
519    pub fn test_iter_nth() {
520        static DES_COUNT: AtomicUsize = AtomicUsize::new(0);
521
522        #[derive(BorshSerialize)]
523        struct DeserializeCounter(u64);
524
525        impl BorshDeserialize for DeserializeCounter {
526            fn deserialize_reader<R: std::io::Read>(reader: &mut R) -> std::io::Result<Self> {
527                DES_COUNT.fetch_add(1, Ordering::SeqCst);
528                u64::deserialize_reader(reader).map(DeserializeCounter)
529            }
530        }
531
532        let mut map = UnorderedMap::new(b"m");
533
534        for i in 0..10 {
535            map.insert(&i, &DeserializeCounter(i));
536        }
537        assert_eq!(DES_COUNT.load(Ordering::SeqCst), 0);
538
539        let collected: Vec<u64> = map.iter().skip(5).take(4).map(|(_, v)| v.0).collect();
540        // 4 or 5 is accepted because pre 1.65 Rust skip loaded an extra value.
541        assert!((4..=5).contains(&DES_COUNT.load(Ordering::SeqCst)));
542        assert_eq!(&collected, &[5, 6, 7, 8]);
543
544        DES_COUNT.store(0, Ordering::SeqCst);
545        let collected: Vec<u64> = map.values().skip(5).take(4).map(|v| v.0).collect();
546        // 4 or 5 is accepted because pre 1.65 Rust skip loaded an extra value.
547        assert!((4..=5).contains(&DES_COUNT.load(Ordering::SeqCst)));
548        assert_eq!(&collected, &[5, 6, 7, 8]);
549    }
550
551    #[test]
552    pub fn test_extend() {
553        let mut map = UnorderedMap::new(b"m");
554        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(4);
555        let mut key_to_value = HashMap::new();
556        for _ in 0..100 {
557            let key = rng.gen::<u64>();
558            let value = rng.gen::<u64>();
559            key_to_value.insert(key, value);
560            map.insert(&key, &value);
561        }
562        for _ in 0..10 {
563            let mut tmp = vec![];
564            for _ in 0..=(rng.gen::<u64>() % 20 + 1) {
565                let key = rng.gen::<u64>();
566                let value = rng.gen::<u64>();
567                tmp.push((key, value));
568            }
569            key_to_value.extend(tmp.iter().cloned());
570            map.extend(tmp.iter().cloned());
571        }
572
573        let actual: HashMap<u64, u64> = map.iter().collect();
574        assert_eq!(actual, key_to_value);
575    }
576
577    #[test]
578    fn test_debug() {
579        let mut map = UnorderedMap::new(b"m");
580        map.insert(&1u64, &100u64);
581        map.insert(&3u64, &300u64);
582        map.insert(&2u64, &200u64);
583
584        if cfg!(feature = "expensive-debug") {
585            assert_eq!(
586                format!("{:?}", map),
587                "UnorderedMap { key_index_prefix: [109, 105], keys: [1, 3, 2], values: [100, 300, 200] }"
588            );
589        } else {
590            assert_eq!(
591                format!("{:?}", map),
592                "UnorderedMap { key_index_prefix: [109, 105], \
593                keys: Vector { len: 3, prefix: [109, 107] }, \
594                values: Vector { len: 3, prefix: [109, 118] } }"
595            );
596        }
597    }
598}
near_sdk/collections/unordered_map/mod.rs

near_sdk/collections/unordered_map/
mod.rs
