unc_sdk/store/iterable_map/mod.rs
1// This suppresses the depreciation warnings for uses of UnorderedSet in this module
2#![allow(deprecated)]
3
4mod entry;
5mod impls;
6mod iter;
7
8use std::borrow::Borrow;
9use std::{fmt, mem};
10
11use borsh::{BorshDeserialize, BorshSerialize};
12
13use unc_sdk_macros::unc;
14
15use crate::store::key::{Sha256, ToKey};
16use crate::{env, IntoStorageKey};
17
18use crate::store::Vector;
19pub use entry::{Entry, OccupiedEntry, VacantEntry};
20
21pub use self::iter::{Drain, Iter, IterMut, Keys, Values, ValuesMut};
22use super::{LookupMap, ERR_INCONSISTENT_STATE, ERR_NOT_EXIST};
23
24/// A lazily loaded storage map that stores its content directly on the storage trie.
25/// This structure is similar to [`unc_sdk::store::LookupMap`](crate::store::LookupMap), except
26/// that it stores the keys so that [`IterableMap`] can be iterable.
27///
28/// This map stores the values under a hash of the map's `prefix` and [`BorshSerialize`] of the key
29/// using the map's [`ToKey`] implementation.
30///
31/// The default hash function for [`IterableMap`] is [`Sha256`] which uses a syscall
32/// (or host function) built into the UNC runtime to hash the key. To use a custom function,
33/// use [`with_hasher`]. Alternative builtin hash functions can be found at
34/// [`unc_sdk::store::key`](crate::store::key).
35///
36///
37/// # Examples
38/// ```
39/// use unc_sdk::store::IterableMap;
40///
41/// // Initializes a map, the generic types can be inferred to `IterableMap<String, u8, Sha256>`
42/// // The `b"a"` parameter is a prefix for the storage keys of this data structure.
43/// let mut map = IterableMap::new(b"a");
44///
45/// map.insert("test".to_string(), 7u8);
46/// assert!(map.contains_key("test"));
47/// assert_eq!(map.get("test"), Some(&7u8));
48///
49/// let prev = std::mem::replace(map.get_mut("test").unwrap(), 5u8);
50/// assert_eq!(prev, 7u8);
51/// assert_eq!(map["test"], 5u8);
52/// ```
53///
54/// [`IterableMap`] also implements an [`Entry API`](Self::entry), which allows
55/// for more complex methods of getting, setting, updating and removing keys and
56/// their values:
57///
58/// ```
59/// use unc_sdk::store::IterableMap;
60///
61/// // type inference lets us omit an explicit type signature (which
62/// // would be `IterableMap<String, u8>` in this example).
63/// let mut player_stats = IterableMap::new(b"m");
64///
65/// fn random_stat_buff() -> u8 {
66/// // could actually return some random value here - let's just return
67/// // some fixed value for now
68/// 42
69/// }
70///
71/// // insert a key only if it doesn't already exist
72/// player_stats.entry("health".to_string()).or_insert(100);
73///
74/// // insert a key using a function that provides a new value only if it
75/// // doesn't already exist
76/// player_stats.entry("defence".to_string()).or_insert_with(random_stat_buff);
77///
78/// // update a key, guarding against the key possibly not being set
79/// let stat = player_stats.entry("attack".to_string()).or_insert(100);
80/// *stat += random_stat_buff();
81/// ```
82///
83/// [`with_hasher`]: Self::with_hasher
84#[unc(inside_uncsdk)]
85pub struct IterableMap<K, V, H = Sha256>
86where
87 K: BorshSerialize + Ord,
88 V: BorshSerialize,
89 H: ToKey,
90{
91 // NOTE: It's important that the `keys` collection is one that's optimized for iteration, e.g.
92 // not skipping empty/unoccupied entries white trying to get to the next element.
93 // See https://github.com/unc/unc-sdk-rs/issues/1134 to understand the difference between
94 // `store::UnorderedMap` and `store::IterableMap`.
95
96 // ser/de is independent of `K` ser/de, `BorshSerialize`/`BorshDeserialize` bounds removed
97 #[borsh(bound(serialize = "", deserialize = ""))]
98 keys: Vector<K>,
99 // ser/de is independent of `K`, `V`, `H` ser/de, `BorshSerialize`/`BorshDeserialize` bounds removed
100 #[borsh(bound(serialize = "", deserialize = ""))]
101 values: LookupMap<K, ValueAndIndex<V>, H>,
102}
103
104#[unc(inside_uncsdk)]
105struct ValueAndIndex<V> {
106 value: V,
107 key_index: u32,
108}
109
110impl<K, V, H> Drop for IterableMap<K, V, H>
111where
112 K: BorshSerialize + Ord,
113 V: BorshSerialize,
114 H: ToKey,
115{
116 fn drop(&mut self) {
117 self.flush()
118 }
119}
120
121impl<K, V, H> fmt::Debug for IterableMap<K, V, H>
122where
123 K: BorshSerialize + Ord + BorshDeserialize + fmt::Debug,
124 V: BorshSerialize,
125 H: ToKey,
126{
127 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
128 f.debug_struct("IterableMap")
129 .field("keys", &self.keys)
130 .field("values", &self.values)
131 .finish()
132 }
133}
134
135impl<K, V> IterableMap<K, V, Sha256>
136where
137 K: BorshSerialize + Ord,
138 V: BorshSerialize,
139{
140 /// Create a new iterable map. Use `prefix` as a unique prefix for keys.
141 ///
142 /// This prefix can be anything that implements [`IntoStorageKey`]. The prefix is used when
143 /// storing and looking up values in storage to ensure no collisions with other collections.
144 ///
145 /// # Examples
146 ///
147 /// ```
148 /// use unc_sdk::store::IterableMap;
149 ///
150 /// let mut map: IterableMap<String, u8> = IterableMap::new(b"b");
151 /// ```
152 #[inline]
153 pub fn new<S>(prefix: S) -> Self
154 where
155 S: IntoStorageKey,
156 {
157 Self::with_hasher(prefix)
158 }
159}
160
161impl<K, V, H> IterableMap<K, V, H>
162where
163 K: BorshSerialize + Ord,
164 V: BorshSerialize,
165 H: ToKey,
166{
167 /// Initialize a [`IterableMap`] with a custom hash function.
168 ///
169 /// # Example
170 /// ```
171 /// use unc_sdk::store::{IterableMap, key::Keccak256};
172 ///
173 /// let map = IterableMap::<String, String, Keccak256>::with_hasher(b"m");
174 /// ```
175 pub fn with_hasher<S>(prefix: S) -> Self
176 where
177 S: IntoStorageKey,
178 {
179 let mut vec_key = prefix.into_storage_key();
180 let map_key = [vec_key.as_slice(), b"m"].concat();
181 vec_key.push(b'v');
182 Self { keys: Vector::new(vec_key), values: LookupMap::with_hasher(map_key) }
183 }
184
185 /// Return the amount of elements inside of the map.
186 ///
187 /// # Example
188 /// ```
189 /// use unc_sdk::store::IterableMap;
190 ///
191 /// let mut map: IterableMap<String, u8> = IterableMap::new(b"b");
192 /// assert_eq!(map.len(), 0);
193 /// map.insert("a".to_string(), 1);
194 /// map.insert("b".to_string(), 2);
195 /// assert_eq!(map.len(), 2);
196 /// ```
197 pub fn len(&self) -> u32 {
198 self.keys.len()
199 }
200
201 /// Returns true if there are no elements inside of the map.
202 ///
203 /// # Example
204 /// ```
205 /// use unc_sdk::store::IterableMap;
206 ///
207 /// let mut map: IterableMap<String, u8> = IterableMap::new(b"b");
208 /// assert!(map.is_empty());
209 /// map.insert("a".to_string(), 1);
210 /// assert!(!map.is_empty());
211 /// ```
212 pub fn is_empty(&self) -> bool {
213 self.keys.is_empty()
214 }
215
216 /// Clears the map, removing all key-value pairs. Keeps the allocated memory
217 /// for reuse.
218 ///
219 /// # Examples
220 ///
221 /// ```
222 /// use unc_sdk::store::IterableMap;
223 ///
224 /// let mut map: IterableMap<String, u8> = IterableMap::new(b"b");
225 /// map.insert("a".to_string(), 1);
226 ///
227 /// map.clear();
228 ///
229 /// assert!(map.is_empty());
230 /// ```
231 pub fn clear(&mut self)
232 where
233 K: BorshDeserialize + Clone,
234 V: BorshDeserialize,
235 {
236 for k in self.keys.drain(..) {
237 // Set instead of remove to avoid loading the value from storage.
238 self.values.set(k, None);
239 }
240 }
241
242 /// An iterator visiting all key-value pairs in arbitrary order.
243 /// The iterator element type is `(&'a K, &'a V)`.
244 ///
245 /// # Examples
246 ///
247 /// ```
248 /// use unc_sdk::store::IterableMap;
249 ///
250 /// let mut map = IterableMap::new(b"m");
251 /// map.insert("a".to_string(), 1);
252 /// map.insert("b".to_string(), 2);
253 /// map.insert("c".to_string(), 3);
254 ///
255 /// for (key, val) in map.iter() {
256 /// println!("key: {} val: {}", key, val);
257 /// }
258 /// ```
259 pub fn iter(&self) -> Iter<K, V, H>
260 where
261 K: BorshDeserialize,
262 {
263 Iter::new(self)
264 }
265
266 /// An iterator visiting all key-value pairs in arbitrary order,
267 /// with exclusive references to the values.
268 /// The iterator element type is `(&'a K, &'a mut V)`.
269 ///
270 /// # Examples
271 ///
272 /// ```
273 /// use unc_sdk::store::IterableMap;
274 ///
275 /// let mut map = IterableMap::new(b"m");
276 /// map.insert("a".to_string(), 1);
277 /// map.insert("b".to_string(), 2);
278 /// map.insert("c".to_string(), 3);
279 ///
280 /// // Update all values
281 /// for (_, val) in map.iter_mut() {
282 /// *val *= 2;
283 /// }
284 ///
285 /// for (key, val) in &map {
286 /// println!("key: {} val: {}", key, val);
287 /// }
288 /// ```
289 pub fn iter_mut(&mut self) -> IterMut<K, V, H>
290 where
291 K: BorshDeserialize,
292 {
293 IterMut::new(self)
294 }
295
296 /// An iterator visiting all keys in arbitrary order.
297 /// The iterator element type is `&'a K`.
298 ///
299 /// # Examples
300 ///
301 /// ```
302 /// use unc_sdk::store::IterableMap;
303 ///
304 /// let mut map = IterableMap::new(b"m");
305 /// map.insert("a".to_string(), 1);
306 /// map.insert("b".to_string(), 2);
307 /// map.insert("c".to_string(), 3);
308 ///
309 /// for key in map.keys() {
310 /// println!("{}", key);
311 /// }
312 /// ```
313 pub fn keys(&self) -> Keys<K>
314 where
315 K: BorshDeserialize,
316 {
317 Keys::new(self)
318 }
319
320 /// An iterator visiting all values in arbitrary order.
321 /// The iterator element type is `&'a V`.
322 ///
323 /// # Examples
324 ///
325 /// ```
326 /// use unc_sdk::store::IterableMap;
327 ///
328 /// let mut map = IterableMap::new(b"m");
329 /// map.insert("a".to_string(), 1);
330 /// map.insert("b".to_string(), 2);
331 /// map.insert("c".to_string(), 3);
332 ///
333 /// for val in map.values() {
334 /// println!("{}", val);
335 /// }
336 /// ```
337 pub fn values(&self) -> Values<K, V, H>
338 where
339 K: BorshDeserialize,
340 {
341 Values::new(self)
342 }
343
344 /// A mutable iterator visiting all values in arbitrary order.
345 /// The iterator element type is `&'a mut V`.
346 ///
347 /// # Examples
348 ///
349 /// ```
350 /// use unc_sdk::store::IterableMap;
351 ///
352 /// let mut map = IterableMap::new(b"m");
353 /// map.insert("a".to_string(), 1);
354 /// map.insert("b".to_string(), 2);
355 /// map.insert("c".to_string(), 3);
356 ///
357 /// for val in map.values_mut() {
358 /// *val = *val + 10;
359 /// }
360 ///
361 /// for val in map.values() {
362 /// println!("{}", val);
363 /// }
364 /// ```
365 pub fn values_mut(&mut self) -> ValuesMut<K, V, H>
366 where
367 K: BorshDeserialize,
368 {
369 ValuesMut::new(self)
370 }
371
372 /// Clears the map, returning all key-value pairs as an iterator.
373 ///
374 /// This will clear all values, even if only some key/value pairs are yielded.
375 ///
376 /// # Examples
377 ///
378 /// ```
379 /// use unc_sdk::store::IterableMap;
380 ///
381 /// let mut a = IterableMap::new(b"m");
382 /// a.insert(1, "a".to_string());
383 /// a.insert(2, "b".to_string());
384 ///
385 /// for (k, v) in a.drain().take(1) {
386 /// assert!(k == 1 || k == 2);
387 /// assert!(&v == "a" || &v == "b");
388 /// }
389 ///
390 /// assert!(a.is_empty());
391 /// ```
392 pub fn drain(&mut self) -> Drain<K, V, H>
393 where
394 K: BorshDeserialize,
395 {
396 Drain::new(self)
397 }
398}
399
400impl<K, V, H> IterableMap<K, V, H>
401where
402 K: BorshSerialize + Ord + Clone,
403 V: BorshSerialize + BorshDeserialize,
404 H: ToKey,
405{
406 /// Returns a reference to the value corresponding to the key.
407 ///
408 /// The key may be any borrowed form of the map's key type, but
409 /// [`BorshSerialize`] and [`ToOwned<Owned = K>`](ToOwned) on the borrowed form *must* match
410 /// those for the key type.
411 ///
412 /// # Examples
413 ///
414 /// ```
415 /// use unc_sdk::store::IterableMap;
416 ///
417 /// let mut map: IterableMap<String, u8> = IterableMap::new(b"b");
418 /// assert!(map.insert("test".to_string(), 5u8).is_none());
419 /// assert_eq!(map.get("test"), Some(&5));
420 /// ```
421 pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
422 where
423 K: Borrow<Q>,
424 Q: BorshSerialize + ToOwned<Owned = K>,
425 {
426 self.values.get(k).map(|v| &v.value)
427 }
428
429 /// Returns a mutable reference to the value corresponding to the key.
430 ///
431 /// The key may be any borrowed form of the map's key type, but
432 /// [`BorshSerialize`] and [`ToOwned<Owned = K>`](ToOwned) on the borrowed form *must* match
433 /// those for the key type.
434 ///
435 /// # Examples
436 ///
437 /// ```
438 /// use unc_sdk::store::IterableMap;
439 ///
440 /// let mut map: IterableMap<String, u8> = IterableMap::new(b"b");
441 /// assert!(map.insert("test".to_string(), 5u8).is_none());
442 ///
443 /// *map.get_mut("test").unwrap() = 6;
444 /// assert_eq!(map["test"], 6);
445 /// ```
446 pub fn get_mut<Q: ?Sized>(&mut self, k: &Q) -> Option<&mut V>
447 where
448 K: Borrow<Q>,
449 Q: BorshSerialize + ToOwned<Owned = K>,
450 {
451 self.values.get_mut(k).map(|v| &mut v.value)
452 }
453
454 /// Inserts a key-value pair into the map.
455 ///
456 /// If the map did not have this key present, [`None`] is returned.
457 ///
458 /// If the map did have this key present, the value is updated, and the old
459 /// value is returned. The key is not updated, though; this matters for
460 /// types that can be `==` without being identical.
461 ///
462 /// # Examples
463 ///
464 /// ```
465 /// use unc_sdk::store::IterableMap;
466 ///
467 /// let mut map: IterableMap<String, u8> = IterableMap::new(b"b");
468 /// assert!(map.is_empty());
469 ///
470 /// map.insert("a".to_string(), 1);
471 ///
472 /// assert!(!map.is_empty());
473 /// assert_eq!(map.values().collect::<Vec<_>>(), [&1]);
474 /// ```
475 pub fn insert(&mut self, k: K, value: V) -> Option<V>
476 where
477 K: Clone + BorshDeserialize,
478 {
479 // Check if value is in map to replace first
480 let entry = self.values.get_mut_inner(&k);
481 if let Some(existing) = entry.value_mut() {
482 return Some(mem::replace(&mut existing.value, value));
483 }
484
485 // At this point, we know that the key-value doesn't exist in the map, add key to bucket.
486 self.keys.push(k);
487 let key_index = self.keys.len() - 1;
488 entry.replace(Some(ValueAndIndex { value, key_index }));
489 None
490 }
491
492 /// Returns `true` if the map contains a value for the specified key.
493 ///
494 /// The key may be any borrowed form of the map's key type, but
495 /// [`BorshSerialize`] and [`ToOwned<Owned = K>`](ToOwned) on the borrowed form *must* match
496 /// those for the key type.
497 ///
498 /// # Examples
499 ///
500 /// ```
501 /// use unc_sdk::store::IterableMap;
502 ///
503 /// let mut map: IterableMap<String, u8> = IterableMap::new(b"b");
504 /// map.insert("test".to_string(), 7u8);
505 ///
506 /// assert!(map.contains_key("test"));
507 /// ```
508 pub fn contains_key<Q: ?Sized>(&self, k: &Q) -> bool
509 where
510 K: Borrow<Q>,
511 Q: BorshSerialize + ToOwned<Owned = K> + Ord,
512 {
513 self.values.contains_key(k)
514 }
515
516 /// Removes a key from the map, returning the value at the key if the key
517 /// was previously in the map.
518 ///
519 /// The key may be any borrowed form of the map's key type, but
520 /// [`BorshSerialize`] and [`ToOwned<Owned = K>`](ToOwned) on the borrowed form *must* match
521 /// those for the key type.
522 ///
523 /// # Performance
524 ///
525 /// When elements are removed, the underlying vector of keys is rearranged by means of swapping
526 /// an obsolete key with the last element in the list and deleting that. Note that that requires
527 /// updating the `values` map due to the fact that it holds `keys` vector indices.
528 ///
529 /// # Examples
530 ///
531 /// ```
532 /// use unc_sdk::store::IterableMap;
533 ///
534 /// let mut map: IterableMap<String, u8> = IterableMap::new(b"b");
535 /// map.insert("test".to_string(), 7u8);
536 /// assert_eq!(map.len(), 1);
537 ///
538 /// map.remove("test");
539 ///
540 /// assert_eq!(map.len(), 0);
541 /// ```
542 pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
543 where
544 K: Borrow<Q> + BorshDeserialize,
545 Q: BorshSerialize + ToOwned<Owned = K>,
546 {
547 self.remove_entry(k).map(|(_, v)| v)
548 }
549
550 /// Removes a key from the map, returning the stored key and value if the
551 /// key was previously in the map.
552 ///
553 /// The key may be any borrowed form of the map's key type, but
554 /// [`BorshSerialize`] and [`ToOwned<Owned = K>`](ToOwned) on the borrowed form *must* match
555 /// those for the key type.
556 ///
557 /// # Performance
558 ///
559 /// When elements are removed, the underlying vector of keys is rearranged by means of swapping
560 /// an obsolete key with the last element in the list and deleting that. Note that that requires
561 /// updating the `values` map due to the fact that it holds `keys` vector indices.
562 ///
563 /// # Examples
564 ///
565 /// ```
566 /// use unc_sdk::store::IterableMap;
567 ///
568 /// let mut map = IterableMap::new(b"m");
569 /// map.insert(1, "a".to_string());
570 /// assert_eq!(map.remove(&1), Some("a".to_string()));
571 /// assert_eq!(map.remove(&1), None);
572 /// ```
573 pub fn remove_entry<Q: ?Sized>(&mut self, k: &Q) -> Option<(K, V)>
574 where
575 K: BorshDeserialize + Clone,
576 Q: BorshSerialize + ToOwned<Owned = K>,
577 {
578 // Remove value
579 let old_value = self.values.remove(&k.to_owned())?;
580
581 // Remove key with index if value exists
582 let last_index = self.keys.len() - 1;
583 let key = self.keys.swap_remove(old_value.key_index);
584
585 Self::remove_entry_helper(&self.keys, &mut self.values, old_value.key_index, last_index);
586
587 // Return removed value
588 Some((key, old_value.value))
589 }
590
591 fn remove_entry_helper(
592 keys: &Vector<K>,
593 values: &mut LookupMap<K, ValueAndIndex<V>, H>,
594 key_index: u32,
595 last_index: u32,
596 ) where
597 K: BorshDeserialize + Clone,
598 {
599 match key_index {
600 // If it's the last/only element - do nothing.
601 x if x == last_index => {}
602 // Otherwise update it's index.
603 _ => {
604 let swapped_key =
605 keys.get(key_index).unwrap_or_else(|| env::panic_str(ERR_INCONSISTENT_STATE));
606 let value = values
607 .get_mut(swapped_key)
608 .unwrap_or_else(|| env::panic_str(ERR_INCONSISTENT_STATE));
609 value.key_index = key_index;
610 }
611 }
612 }
613
614 /// Gets the given key's corresponding entry in the map for in-place manipulation.
615 ///
616 /// # Performance
617 /// Note that due to the fact that we need to potentially re-arrange `keys` and update `values`,
618 /// `Entry` API actually operates on those two collections as opposed to an actual `Entry`
619 /// ```
620 /// use unc_sdk::store::IterableMap;
621 ///
622 /// let mut count = IterableMap::new(b"m");
623 ///
624 /// for ch in [7, 2, 4, 7, 4, 1, 7] {
625 /// let counter = count.entry(ch).or_insert(0);
626 /// *counter += 1;
627 /// }
628 ///
629 /// assert_eq!(count[&4], 2);
630 /// assert_eq!(count[&7], 3);
631 /// assert_eq!(count[&1], 1);
632 /// assert_eq!(count.get(&8), None);
633 /// ```
634 pub fn entry(&mut self, key: K) -> Entry<K, V, H>
635 where
636 K: Clone,
637 {
638 Entry::new(key, &mut self.keys, &mut self.values)
639 }
640}
641
642impl<K, V, H> IterableMap<K, V, H>
643where
644 K: BorshSerialize + Ord,
645 V: BorshSerialize,
646 H: ToKey,
647{
648 /// Flushes the intermediate values of the map before this is called when the structure is
649 /// [`Drop`]ed. This will write all modified values to storage but keep all cached values
650 /// in memory.
651 pub fn flush(&mut self) {
652 self.keys.flush();
653 self.values.flush();
654 }
655}
656
657#[cfg(not(target_arch = "wasm32"))]
658#[cfg(test)]
659mod tests {
660 use super::IterableMap;
661 use crate::test_utils::test_env::setup_free;
662 use arbitrary::{Arbitrary, Unstructured};
663 use borsh::{to_vec, BorshDeserialize};
664 use rand::RngCore;
665 use rand::SeedableRng;
666 use std::collections::HashMap;
667
668 #[test]
669 fn basic_functionality() {
670 let mut map = IterableMap::new(b"b");
671 assert!(map.is_empty());
672 assert!(map.insert("test".to_string(), 5u8).is_none());
673 assert_eq!(map.get("test"), Some(&5));
674 assert_eq!(map.len(), 1);
675
676 *map.get_mut("test").unwrap() = 6;
677 assert_eq!(map["test"], 6);
678
679 assert_eq!(map.remove("test"), Some(6));
680 assert_eq!(map.len(), 0);
681 }
682
683 #[test]
684 fn entry_api() {
685 let mut map = IterableMap::new(b"b");
686 {
687 let test_entry = map.entry("test".to_string());
688 assert_eq!(test_entry.key(), "test");
689 let entry_ref = test_entry.or_insert(8u8);
690 *entry_ref += 1;
691 }
692 assert_eq!(map["test"], 9);
693
694 // Try getting entry of filled value
695 let value = map.entry("test".to_string()).and_modify(|v| *v += 3).or_default();
696 assert_eq!(*value, 12);
697 }
698
699 #[test]
700 fn map_iterator() {
701 let mut map = IterableMap::new(b"b");
702
703 map.insert(0u8, 0u8);
704 map.insert(1, 1);
705 map.insert(2, 2);
706 map.insert(3, 3);
707 map.remove(&1);
708
709 let iter = map.iter();
710 assert_eq!(iter.len(), 3);
711 assert_eq!(iter.collect::<Vec<_>>(), [(&0, &0), (&3, &3), (&2, &2)]);
712
713 let iter = map.iter_mut().rev();
714 assert_eq!(iter.collect::<Vec<_>>(), [(&2, &mut 2), (&3, &mut 3), (&0, &mut 0)]);
715
716 let mut iter = map.iter();
717 assert_eq!(iter.nth(2), Some((&2, &2)));
718 // Check fused iterator assumption that each following one will be None
719 assert_eq!(iter.next(), None);
720
721 // Double all values
722 map.values_mut().for_each(|v| {
723 *v *= 2;
724 });
725 assert_eq!(map.values().collect::<Vec<_>>(), [&0, &6, &4]);
726
727 // Collect all keys
728 assert_eq!(map.keys().collect::<Vec<_>>(), [&0, &3, &2]);
729 }
730
731 #[derive(Arbitrary, Debug)]
732 enum Op {
733 Insert(u8, u8),
734 Remove(u8),
735 Flush,
736 Restore,
737 Get(u8),
738 }
739
740 #[test]
741 fn arbitrary() {
742 setup_free();
743
744 let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(0);
745 let mut buf = vec![0; 4096];
746 for _ in 0..512 {
747 // Clear storage in-between runs
748 crate::mock::with_mocked_blockchain(|b| b.take_storage());
749 rng.fill_bytes(&mut buf);
750
751 let mut um = IterableMap::new(b"l");
752 let mut hm = HashMap::new();
753 let u = Unstructured::new(&buf);
754 if let Ok(ops) = Vec::<Op>::arbitrary_take_rest(u) {
755 for op in ops {
756 match op {
757 Op::Insert(k, v) => {
758 let r1 = um.insert(k, v);
759 let r2 = hm.insert(k, v);
760 assert_eq!(r1, r2)
761 }
762 Op::Remove(k) => {
763 let r1 = um.remove(&k);
764 let r2 = hm.remove(&k);
765 assert_eq!(r1, r2)
766 }
767 Op::Flush => {
768 um.flush();
769 }
770 Op::Restore => {
771 let serialized = to_vec(&um).unwrap();
772 um = IterableMap::deserialize(&mut serialized.as_slice()).unwrap();
773 }
774 Op::Get(k) => {
775 let r1 = um.get(&k);
776 let r2 = hm.get(&k);
777 assert_eq!(r1, r2)
778 }
779 }
780 }
781 }
782 }
783 }
784}
785
786// Hashbrown-like tests.
787#[cfg(test)]
788mod test_map {
789 use super::Entry::{Occupied, Vacant};
790 use crate::store::IterableMap;
791 use borsh::{BorshDeserialize, BorshSerialize};
792 use rand::{rngs::SmallRng, Rng, SeedableRng};
793 use std::cell::RefCell;
794 use std::usize;
795 use std::vec::Vec;
796
797 #[test]
798 fn test_insert() {
799 let mut m = IterableMap::new(b"b");
800 assert_eq!(m.len(), 0);
801 assert!(m.insert(1, 2).is_none());
802 assert_eq!(m.len(), 1);
803 assert!(m.insert(2, 4).is_none());
804 assert_eq!(m.len(), 2);
805 assert_eq!(*m.get(&1).unwrap(), 2);
806 assert_eq!(*m.get(&2).unwrap(), 4);
807 }
808
809 thread_local! { static DROP_VECTOR: RefCell<Vec<i32>> = const { RefCell::new(Vec::new()) }}
810
811 #[derive(Hash, PartialEq, Eq, BorshSerialize, BorshDeserialize, PartialOrd, Ord)]
812 struct Droppable {
813 k: usize,
814 }
815
816 impl Droppable {
817 fn new(k: usize) -> Droppable {
818 DROP_VECTOR.with(|slot| {
819 slot.borrow_mut()[k] += 1;
820 });
821
822 Droppable { k }
823 }
824 }
825
826 impl Drop for Droppable {
827 fn drop(&mut self) {
828 DROP_VECTOR.with(|slot| {
829 slot.borrow_mut()[self.k] -= 1;
830 });
831 }
832 }
833
834 impl Clone for Droppable {
835 fn clone(&self) -> Self {
836 Droppable::new(self.k)
837 }
838 }
839
840 #[test]
841 fn test_drops() {
842 DROP_VECTOR.with(|slot| {
843 *slot.borrow_mut() = vec![0; 200];
844 });
845
846 {
847 let mut m = IterableMap::new(b"b");
848
849 DROP_VECTOR.with(|v| {
850 for i in 0..200 {
851 assert_eq!(v.borrow()[i], 0);
852 }
853 });
854
855 for i in 0..100 {
856 let d1 = Droppable::new(i);
857 let d2 = Droppable::new(i + 100);
858 m.insert(d1, d2);
859 }
860
861 DROP_VECTOR.with(|v| {
862 for i in 0..100 {
863 assert_eq!(v.borrow()[i], 2);
864 }
865 });
866
867 for i in 0..50 {
868 let k = Droppable::new(i);
869 let v = m.remove(&k);
870
871 assert!(v.is_some());
872
873 DROP_VECTOR.with(|v| {
874 assert_eq!(v.borrow()[i], 2);
875 assert_eq!(v.borrow()[i + 100], 1);
876 });
877 }
878
879 DROP_VECTOR.with(|v| {
880 for i in 0..50 {
881 assert_eq!(v.borrow()[i], 1);
882 assert_eq!(v.borrow()[i + 100], 0);
883 }
884
885 for i in 50..100 {
886 assert_eq!(v.borrow()[i], 2);
887 assert_eq!(v.borrow()[i + 100], 1);
888 }
889 });
890 }
891
892 DROP_VECTOR.with(|v| {
893 for i in 0..200 {
894 assert_eq!(v.borrow()[i], 0);
895 }
896 });
897 }
898
899 #[test]
900 fn test_into_iter_drops() {
901 DROP_VECTOR.with(|v| {
902 *v.borrow_mut() = vec![0; 200];
903 });
904
905 let hm = {
906 let mut hm = IterableMap::new(b"b");
907
908 DROP_VECTOR.with(|v| {
909 for i in 0..200 {
910 assert_eq!(v.borrow()[i], 0);
911 }
912 });
913
914 for i in 0..100 {
915 let d1 = Droppable::new(i);
916 let d2 = Droppable::new(i + 100);
917 hm.insert(d1, d2);
918 }
919
920 DROP_VECTOR.with(|v| {
921 for i in 0..100 {
922 assert_eq!(v.borrow()[i], 2);
923 }
924 for i in 101..200 {
925 assert_eq!(v.borrow()[i], 1);
926 }
927 });
928
929 hm
930 };
931
932 {
933 let mut half = hm.into_iter().take(50);
934
935 DROP_VECTOR.with(|v| {
936 for i in 0..100 {
937 assert_eq!(v.borrow()[i], 2);
938 }
939 for i in 101..200 {
940 assert_eq!(v.borrow()[i], 1);
941 }
942 });
943
944 #[allow(let_underscore_drop)] // kind-of a false positive
945 for _ in half.by_ref() {}
946
947 DROP_VECTOR.with(|v| {
948 let nk = (0..100).filter(|&i| v.borrow()[i] == 2).count();
949
950 let nv = (0..100).filter(|&i| v.borrow()[i + 100] == 1).count();
951
952 assert_eq!(nk, 100);
953 assert_eq!(nv, 100);
954 });
955 };
956
957 drop(hm);
958
959 DROP_VECTOR.with(|v| {
960 for i in 0..200 {
961 assert_eq!(v.borrow()[i], 0);
962 }
963 });
964 }
965
966 #[test]
967 fn test_empty_remove() {
968 let mut m: IterableMap<i32, bool> = IterableMap::new(b"b");
969 assert_eq!(m.remove(&0), None);
970 }
971
972 #[test]
973 fn test_empty_entry() {
974 let mut m: IterableMap<i32, bool> = IterableMap::new(b"b");
975 match m.entry(0) {
976 Occupied(_) => panic!(),
977 Vacant(_) => {}
978 }
979 assert!(*m.entry(0).or_insert(true));
980 assert_eq!(m.len(), 1);
981 }
982
983 #[test]
984 fn test_empty_iter() {
985 let mut m: IterableMap<i32, bool> = IterableMap::new(b"b");
986 assert_eq!(m.drain().next(), None);
987 assert_eq!(m.keys().next(), None);
988 assert_eq!(m.values().next(), None);
989 assert_eq!(m.values_mut().next(), None);
990 assert_eq!(m.iter().next(), None);
991 assert_eq!(m.iter_mut().next(), None);
992 assert_eq!(m.len(), 0);
993 assert!(m.is_empty());
994 assert_eq!(m.into_iter().next(), None);
995 }
996
997 #[test]
998 #[cfg_attr(miri, ignore)] // FIXME: takes too long
999 fn test_lots_of_insertions() {
1000 let mut m = IterableMap::new(b"b");
1001
1002 // Try this a few times to make sure we never screw up the IterableMap's
1003 // internal state.
1004 for _ in 0..10 {
1005 assert!(m.is_empty());
1006
1007 for i in 1..1001 {
1008 assert!(m.insert(i, i).is_none());
1009
1010 for j in 1..=i {
1011 let r = m.get(&j);
1012 assert_eq!(r, Some(&j));
1013 }
1014
1015 for j in i + 1..1001 {
1016 let r = m.get(&j);
1017 assert_eq!(r, None);
1018 }
1019 }
1020
1021 for i in 1001..2001 {
1022 assert!(!m.contains_key(&i));
1023 }
1024
1025 // remove forwards
1026 for i in 1..1001 {
1027 assert!(m.remove(&i).is_some());
1028
1029 for j in 1..=i {
1030 assert!(!m.contains_key(&j));
1031 }
1032
1033 for j in i + 1..1001 {
1034 assert!(m.contains_key(&j));
1035 }
1036 }
1037
1038 for i in 1..1001 {
1039 assert!(!m.contains_key(&i));
1040 }
1041
1042 for i in 1..1001 {
1043 assert!(m.insert(i, i).is_none());
1044 }
1045
1046 // remove backwards
1047 for i in (1..1001).rev() {
1048 assert!(m.remove(&i).is_some());
1049
1050 for j in i..1001 {
1051 assert!(!m.contains_key(&j));
1052 }
1053
1054 for j in 1..i {
1055 assert!(m.contains_key(&j));
1056 }
1057 }
1058 }
1059 }
1060
1061 #[test]
1062 fn test_find_mut() {
1063 let mut m = IterableMap::new(b"b");
1064 assert!(m.insert(1, 12).is_none());
1065 assert!(m.insert(2, 8).is_none());
1066 assert!(m.insert(5, 14).is_none());
1067 let new = 100;
1068 match m.get_mut(&5) {
1069 None => panic!(),
1070 Some(x) => *x = new,
1071 }
1072 assert_eq!(m.get(&5), Some(&new));
1073 }
1074
1075 #[test]
1076 fn test_insert_overwrite() {
1077 let mut m = IterableMap::new(b"b");
1078 assert!(m.insert(1, 2).is_none());
1079 assert_eq!(*m.get(&1).unwrap(), 2);
1080 assert!(m.insert(1, 3).is_some());
1081 assert_eq!(*m.get(&1).unwrap(), 3);
1082 }
1083
1084 #[test]
1085 fn test_is_empty() {
1086 let mut m = IterableMap::new(b"b");
1087 assert!(m.insert(1, 2).is_none());
1088 assert!(!m.is_empty());
1089 assert!(m.remove(&1).is_some());
1090 assert!(m.is_empty());
1091 }
1092
1093 #[test]
1094 fn test_remove() {
1095 let mut m = IterableMap::new(b"b");
1096 m.insert(1, 2);
1097 assert_eq!(m.remove(&1), Some(2));
1098 assert_eq!(m.remove(&1), None);
1099 }
1100
1101 #[test]
1102 fn test_remove_entry() {
1103 let mut m = IterableMap::new(b"b");
1104 m.insert(1, 2);
1105 assert_eq!(m.remove_entry(&1), Some((1, 2)));
1106 assert_eq!(m.remove(&1), None);
1107 }
1108
1109 #[test]
1110 fn test_iterate() {
1111 let mut m = IterableMap::new(b"b");
1112 for i in 0..32 {
1113 assert!(m.insert(i, i * 2).is_none());
1114 }
1115 assert_eq!(m.len(), 32);
1116
1117 let mut observed: u32 = 0;
1118
1119 for (k, v) in &m {
1120 assert_eq!(*v, *k * 2);
1121 observed |= 1 << *k;
1122 }
1123 assert_eq!(observed, 0xFFFF_FFFF);
1124 }
1125
1126 #[test]
1127 fn test_find() {
1128 let mut m = IterableMap::new(b"b");
1129 assert!(m.get(&1).is_none());
1130 m.insert(1, 2);
1131 match m.get(&1) {
1132 None => panic!(),
1133 Some(v) => assert_eq!(*v, 2),
1134 }
1135 }
1136
1137 #[test]
1138 fn test_show() {
1139 let mut map = IterableMap::new(b"b");
1140 let empty: IterableMap<i32, i32> = IterableMap::new(b"c");
1141
1142 map.insert(1, 2);
1143 map.insert(3, 4);
1144
1145 let map_str = format!("{:?}", map);
1146
1147 assert_eq!(map_str, "IterableMap { keys: Vector { len: 2, prefix: [98, 118] }, values: LookupMap { prefix: [98, 109] } }");
1148 assert_eq!(format!("{:?}", empty), "IterableMap { keys: Vector { len: 0, prefix: [99, 118] }, values: LookupMap { prefix: [99, 109] } }");
1149 }
1150
1151 #[test]
1152 fn test_size_hint() {
1153 let mut map = IterableMap::new(b"b");
1154
1155 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)];
1156
1157 for v in xs {
1158 map.insert(v.0, v.1);
1159 }
1160
1161 let mut iter = map.iter();
1162
1163 for _ in iter.by_ref().take(3) {}
1164
1165 assert_eq!(iter.size_hint(), (3, Some(3)));
1166 }
1167
1168 #[test]
1169 fn test_iter_len() {
1170 let mut map = IterableMap::new(b"b");
1171
1172 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)];
1173
1174 for v in xs {
1175 map.insert(v.0, v.1);
1176 }
1177
1178 let mut iter = map.iter();
1179
1180 for _ in iter.by_ref().take(3) {}
1181
1182 assert_eq!(iter.len(), 3);
1183 }
1184
1185 #[test]
1186 fn test_mut_size_hint() {
1187 let mut map = IterableMap::new(b"b");
1188
1189 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)];
1190
1191 for v in xs {
1192 map.insert(v.0, v.1);
1193 }
1194
1195 let mut iter = map.iter_mut();
1196
1197 for _ in iter.by_ref().take(3) {}
1198
1199 assert_eq!(iter.size_hint(), (3, Some(3)));
1200 }
1201
1202 #[test]
1203 fn test_iter_mut_len() {
1204 let mut map = IterableMap::new(b"b");
1205
1206 let xs = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)];
1207
1208 for v in xs {
1209 map.insert(v.0, v.1);
1210 }
1211
1212 let mut iter = map.iter_mut();
1213
1214 for _ in iter.by_ref().take(3) {}
1215
1216 assert_eq!(iter.len(), 3);
1217 }
1218
1219 #[test]
1220 fn test_index() {
1221 let mut map = IterableMap::new(b"b");
1222
1223 map.insert(1, 2);
1224 map.insert(2, 1);
1225 map.insert(3, 4);
1226
1227 assert_eq!(map[&2], 1);
1228 }
1229
1230 #[test]
1231 #[should_panic]
1232 #[allow(clippy::unnecessary_operation)]
1233 fn test_index_nonexistent() {
1234 let mut map = IterableMap::new(b"b");
1235
1236 map.insert(1, 2);
1237 map.insert(2, 1);
1238 map.insert(3, 4);
1239
1240 #[allow(clippy::no_effect)] // false positive lint
1241 map[&4];
1242 }
1243
1244 #[test]
1245 fn test_entry() {
1246 let mut map = IterableMap::new(b"b");
1247
1248 let xs = [(1, 10), (2, 20), (3, 30), (4, 40), (5, 50), (6, 60)];
1249
1250 for v in xs {
1251 map.insert(v.0, v.1);
1252 }
1253
1254 // Existing key (insert)
1255 match map.entry(1) {
1256 Vacant(_) => unreachable!(),
1257 Occupied(mut view) => {
1258 assert_eq!(view.get(), &10);
1259 assert_eq!(view.insert(100), 10);
1260 }
1261 }
1262 assert_eq!(map.get(&1).unwrap(), &100);
1263 assert_eq!(map.len(), 6);
1264
1265 // Existing key (update)
1266 match map.entry(2) {
1267 Vacant(_) => unreachable!(),
1268 Occupied(mut view) => {
1269 let v = view.get_mut();
1270 let new_v = (*v) * 10;
1271 *v = new_v;
1272 }
1273 }
1274 assert_eq!(map.get(&2).unwrap(), &200);
1275 assert_eq!(map.len(), 6);
1276
1277 // Existing key (take)
1278 match map.entry(3) {
1279 Vacant(_) => unreachable!(),
1280 Occupied(view) => {
1281 assert_eq!(view.remove(), 30);
1282 }
1283 }
1284 assert_eq!(map.get(&3), None);
1285 assert_eq!(map.len(), 5);
1286
1287 // Inexistent key (insert)
1288 match map.entry(10) {
1289 Occupied(_) => unreachable!(),
1290 Vacant(view) => {
1291 assert_eq!(*view.insert(1000), 1000);
1292 }
1293 }
1294 assert_eq!(map.get(&10).unwrap(), &1000);
1295 assert_eq!(map.len(), 6);
1296 }
1297
1298 #[test]
1299 fn test_entry_take_doesnt_corrupt() {
1300 fn check(m: &IterableMap<i32, ()>) {
1301 for k in m.keys() {
1302 assert!(m.contains_key(k), "{} is in keys() but not in the map?", k);
1303 }
1304 }
1305
1306 let mut m = IterableMap::new(b"b");
1307
1308 let mut rng = {
1309 let seed = u64::from_le_bytes(*b"testseed");
1310 SmallRng::seed_from_u64(seed)
1311 };
1312
1313 // Populate the map with some items.
1314 for _ in 0..50 {
1315 let x = rng.gen_range(-10..10);
1316 m.insert(x, ());
1317 }
1318
1319 for _ in 0..1000 {
1320 let x = rng.gen_range(-10..10);
1321 match m.entry(x) {
1322 Vacant(_) => {}
1323 Occupied(e) => {
1324 e.remove();
1325 }
1326 }
1327
1328 check(&m);
1329 }
1330 }
1331
1332 #[test]
1333 fn test_extend_ref_kv_tuple() {
1334 let mut a = IterableMap::new(b"b");
1335 a.insert(0, 0);
1336
1337 let for_iter: Vec<(i32, i32)> = (0..100).map(|i| (i, i)).collect();
1338 a.extend(for_iter);
1339
1340 assert_eq!(a.len(), 100);
1341
1342 for item in 0..100 {
1343 assert_eq!(a[&item], item);
1344 }
1345 }
1346
1347 #[test]
1348 fn test_occupied_entry_key() {
1349 let mut a = IterableMap::new(b"b");
1350 let key = "hello there";
1351 let value = "value goes here";
1352 assert!(a.is_empty());
1353 a.insert(key.to_string(), value.to_string());
1354 assert_eq!(a.len(), 1);
1355 assert_eq!(a[key], value);
1356
1357 match a.entry(key.to_string()) {
1358 Vacant(_) => panic!(),
1359 Occupied(e) => assert_eq!(key, *e.key()),
1360 }
1361 assert_eq!(a.len(), 1);
1362 assert_eq!(a[key], value);
1363 }
1364
1365 #[test]
1366 fn test_vacant_entry_key() {
1367 let mut a = IterableMap::new(b"b");
1368 let key = "hello there";
1369 let value = "value goes here".to_string();
1370
1371 assert!(a.is_empty());
1372 match a.entry(key.to_string()) {
1373 Occupied(_) => panic!(),
1374 Vacant(e) => {
1375 assert_eq!(key, *e.key());
1376 e.insert(value.clone());
1377 }
1378 }
1379 assert_eq!(a.len(), 1);
1380 assert_eq!(a[key], value);
1381 }
1382}