multiversx_sc/storage/mappers/
map_mapper.rs

1use core::marker::PhantomData;
2
3use super::{
4    set_mapper::{self},
5    source::{CurrentStorage, StorageAddress},
6    SetMapper, StorageClearable, StorageMapper, StorageMapperFromAddress,
7};
8use crate::{
9    abi::{TypeAbi, TypeAbiFrom, TypeDescriptionContainer, TypeName},
10    api::StorageMapperApi,
11    codec::{
12        multi_encode_iter_or_handle_err, multi_types::MultiValue2, EncodeErrorHandler,
13        NestedDecode, NestedEncode, TopDecode, TopEncode, TopEncodeMulti, TopEncodeMultiOutput,
14    },
15    storage::{storage_clear, storage_set, StorageKey},
16    types::{ManagedAddress, ManagedType, MultiValueEncoded},
17};
18
19const MAPPED_VALUE_IDENTIFIER: &[u8] = b".mapped";
20type Keys<'a, SA, A, T> = set_mapper::Iter<'a, SA, A, T>;
21
22pub struct MapMapper<SA, K, V, A = CurrentStorage>
23where
24    SA: StorageMapperApi,
25    K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
26    A: StorageAddress<SA>,
27    V: TopEncode + TopDecode + 'static,
28{
29    _phantom_api: PhantomData<SA>,
30    address: A,
31    base_key: StorageKey<SA>,
32    keys_set: SetMapper<SA, K, A>,
33    _phantom_value: PhantomData<V>,
34}
35
36impl<SA, K, V> StorageMapper<SA> for MapMapper<SA, K, V, CurrentStorage>
37where
38    SA: StorageMapperApi,
39    K: TopEncode + TopDecode + NestedEncode + NestedDecode,
40    V: TopEncode + TopDecode,
41{
42    fn new(base_key: StorageKey<SA>) -> Self {
43        MapMapper {
44            _phantom_api: PhantomData,
45            address: CurrentStorage,
46            base_key: base_key.clone(),
47            keys_set: SetMapper::new(base_key),
48            _phantom_value: PhantomData,
49        }
50    }
51}
52
53impl<SA, K, V> StorageClearable for MapMapper<SA, K, V, CurrentStorage>
54where
55    SA: StorageMapperApi,
56    K: TopEncode + TopDecode + NestedEncode + NestedDecode,
57    V: TopEncode + TopDecode,
58{
59    fn clear(&mut self) {
60        for key in self.keys_set.iter() {
61            self.clear_mapped_value(&key);
62        }
63        self.keys_set.clear();
64    }
65}
66
67impl<SA, K, V> MapMapper<SA, K, V, CurrentStorage>
68where
69    SA: StorageMapperApi,
70    K: TopEncode + TopDecode + NestedEncode + NestedDecode,
71    V: TopEncode + TopDecode,
72{
73    fn set_mapped_value(&self, key: &K, value: &V) {
74        storage_set(
75            self.build_named_key(MAPPED_VALUE_IDENTIFIER, key).as_ref(),
76            &value,
77        );
78    }
79
80    fn clear_mapped_value(&self, key: &K) {
81        storage_clear(self.build_named_key(MAPPED_VALUE_IDENTIFIER, key).as_ref());
82    }
83
84    /// Sets the value of the entry, and returns the entry's old value.
85    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
86        let old_value = self.get(&k);
87        self.set_mapped_value(&k, &v);
88        self.keys_set.insert(k);
89        old_value
90    }
91
92    /// Takes the value out of the entry, and returns it.
93    pub fn remove(&mut self, k: &K) -> Option<V> {
94        if self.keys_set.remove(k) {
95            let value = self.get_mapped_value(k);
96            self.clear_mapped_value(k);
97            return Some(value);
98        }
99        None
100    }
101}
102
103impl<SA, K, V> StorageMapperFromAddress<SA> for MapMapper<SA, K, V, ManagedAddress<SA>>
104where
105    SA: StorageMapperApi,
106    K: TopEncode + TopDecode + NestedEncode + NestedDecode,
107    V: TopEncode + TopDecode,
108{
109    fn new_from_address(address: ManagedAddress<SA>, base_key: StorageKey<SA>) -> Self {
110        MapMapper {
111            _phantom_api: PhantomData,
112            address: address.clone(),
113            base_key: base_key.clone(),
114            keys_set: SetMapper::new_from_address(address, base_key),
115            _phantom_value: PhantomData,
116        }
117    }
118}
119
120impl<'a, SA, A, K, V> IntoIterator for &'a MapMapper<SA, K, V, A>
121where
122    SA: StorageMapperApi,
123    A: StorageAddress<SA>,
124    K: TopEncode + TopDecode + NestedEncode + NestedDecode,
125    V: TopEncode + TopDecode,
126{
127    type Item = (K, V);
128
129    type IntoIter = Iter<'a, SA, A, K, V>;
130
131    fn into_iter(self) -> Self::IntoIter {
132        self.iter()
133    }
134}
135
136impl<SA, A, K, V> MapMapper<SA, K, V, A>
137where
138    SA: StorageMapperApi,
139    K: TopEncode + TopDecode + NestedEncode + NestedDecode,
140    A: StorageAddress<SA>,
141    V: TopEncode + TopDecode,
142{
143    /// Returns `true` if the map contains a value for the specified key.
144    pub fn contains_key(&self, k: &K) -> bool {
145        self.keys_set.contains(k)
146    }
147
148    fn build_named_key(&self, name: &[u8], key: &K) -> StorageKey<SA> {
149        let mut named_key = self.base_key.clone();
150        named_key.append_bytes(name);
151        named_key.append_item(key);
152        named_key
153    }
154
155    fn get_mapped_value(&self, key: &K) -> V {
156        self.address
157            .address_storage_get(self.build_named_key(MAPPED_VALUE_IDENTIFIER, key).as_ref())
158    }
159
160    /// Gets a reference to the value in the entry.
161    pub fn get(&self, k: &K) -> Option<V> {
162        if self.keys_set.contains(k) {
163            return Some(self.get_mapped_value(k));
164        }
165        None
166    }
167
168    pub fn keys(&self) -> Keys<'_, SA, A, K> {
169        self.keys_set.iter()
170    }
171
172    /// Returns `true` if the map contains no elements.
173    pub fn is_empty(&self) -> bool {
174        self.keys_set.is_empty()
175    }
176
177    /// Returns the number of elements in the map.
178    pub fn len(&self) -> usize {
179        self.keys_set.len()
180    }
181
182    /// Gets the given key's corresponding entry in the map for in-place manipulation.
183    pub fn entry(&mut self, key: K) -> Entry<'_, SA, A, K, V> {
184        if self.contains_key(&key) {
185            Entry::Occupied(OccupiedEntry {
186                key,
187                map: self,
188                _marker: PhantomData,
189            })
190        } else {
191            Entry::Vacant(VacantEntry {
192                key,
193                map: self,
194                _marker: PhantomData,
195            })
196        }
197    }
198
199    /// An iterator visiting all values in arbitrary order.
200    /// The iterator element type is `&'a V`.
201    pub fn values(&self) -> Values<'_, SA, A, K, V> {
202        Values::new(self)
203    }
204
205    /// An iterator visiting all key-value pairs in arbitrary order.
206    /// The iterator element type is `(&'a K, &'a V)`.
207    pub fn iter(&self) -> Iter<'_, SA, A, K, V> {
208        Iter::new(self)
209    }
210}
211
212pub struct Iter<'a, SA, A, K, V>
213where
214    SA: StorageMapperApi,
215    K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
216    A: StorageAddress<SA>,
217    V: TopEncode + TopDecode + 'static,
218{
219    key_iter: Keys<'a, SA, A, K>,
220    hash_map: &'a MapMapper<SA, K, V, A>,
221}
222
223impl<'a, SA, A, K, V> Iter<'a, SA, A, K, V>
224where
225    SA: StorageMapperApi,
226    A: StorageAddress<SA>,
227    K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
228    V: TopEncode + TopDecode + 'static,
229{
230    fn new(hash_map: &'a MapMapper<SA, K, V, A>) -> Iter<'a, SA, A, K, V> {
231        Iter {
232            key_iter: hash_map.keys(),
233            hash_map,
234        }
235    }
236}
237
238impl<SA, A, K, V> Iterator for Iter<'_, SA, A, K, V>
239where
240    SA: StorageMapperApi,
241    A: StorageAddress<SA>,
242    K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
243    V: TopEncode + TopDecode + 'static,
244{
245    type Item = (K, V);
246
247    #[inline]
248    fn next(&mut self) -> Option<(K, V)> {
249        if let Some(key) = self.key_iter.next() {
250            let value = self.hash_map.get(&key).unwrap();
251            return Some((key, value));
252        }
253        None
254    }
255}
256
257pub struct Values<'a, SA, A, K, V>
258where
259    SA: StorageMapperApi,
260    K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
261    A: StorageAddress<SA>,
262    V: TopEncode + TopDecode + 'static,
263{
264    key_iter: Keys<'a, SA, A, K>,
265    hash_map: &'a MapMapper<SA, K, V, A>,
266}
267
268impl<'a, SA, A, K, V> Values<'a, SA, A, K, V>
269where
270    SA: StorageMapperApi,
271    A: StorageAddress<SA>,
272    K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
273    V: TopEncode + TopDecode + 'static,
274{
275    fn new(hash_map: &'a MapMapper<SA, K, V, A>) -> Values<'a, SA, A, K, V> {
276        Values {
277            key_iter: hash_map.keys(),
278            hash_map,
279        }
280    }
281}
282
283impl<SA, A, K, V> Iterator for Values<'_, SA, A, K, V>
284where
285    SA: StorageMapperApi,
286    A: StorageAddress<SA>,
287    K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
288    V: TopEncode + TopDecode + 'static,
289{
290    type Item = V;
291
292    #[inline]
293    fn next(&mut self) -> Option<V> {
294        if let Some(key) = self.key_iter.next() {
295            let value = self.hash_map.get(&key).unwrap();
296            return Some(value);
297        }
298        None
299    }
300}
301
302pub enum Entry<'a, SA, A, K: 'a, V: 'a>
303where
304    SA: StorageMapperApi,
305    K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
306    A: StorageAddress<SA>,
307    V: TopEncode + TopDecode + 'static,
308{
309    /// A vacant entry.
310    Vacant(VacantEntry<'a, SA, A, K, V>),
311
312    /// An occupied entry.
313    Occupied(OccupiedEntry<'a, SA, A, K, V>),
314}
315
316/// A view into a vacant entry in a `MapMapper`.
317/// It is part of the [`Entry`] enum.
318pub struct VacantEntry<'a, SA, A, K: 'a, V: 'a>
319where
320    SA: StorageMapperApi,
321    K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
322    A: StorageAddress<SA>,
323    V: TopEncode + TopDecode + 'static,
324{
325    pub(super) key: K,
326    pub(super) map: &'a mut MapMapper<SA, K, V, A>,
327
328    // Be invariant in `K` and `V`
329    pub(super) _marker: PhantomData<&'a mut (K, V)>,
330}
331
332/// A view into an occupied entry in a `MapMapper`.
333/// It is part of the [`Entry`] enum.
334pub struct OccupiedEntry<'a, SA, A, K: 'a, V: 'a>
335where
336    SA: StorageMapperApi,
337    K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
338    A: StorageAddress<SA>,
339    V: TopEncode + TopDecode + 'static,
340{
341    pub(super) key: K,
342    pub(super) map: &'a mut MapMapper<SA, K, V, A>,
343
344    // Be invariant in `K` and `V`
345    pub(super) _marker: PhantomData<&'a mut (K, V)>,
346}
347
348impl<SA, A, K, V> Entry<'_, SA, A, K, V>
349where
350    SA: StorageMapperApi,
351    A: StorageAddress<SA>,
352    K: TopEncode + TopDecode + NestedEncode + NestedDecode + Clone,
353    V: TopEncode + TopDecode + 'static,
354{
355    /// Returns a reference to this entry's key.
356    pub fn key(&self) -> &K {
357        match *self {
358            Entry::Occupied(ref entry) => entry.key(),
359            Entry::Vacant(ref entry) => entry.key(),
360        }
361    }
362}
363
364impl<'a, SA, K, V> Entry<'a, SA, CurrentStorage, K, V>
365where
366    SA: StorageMapperApi,
367    K: TopEncode + TopDecode + NestedEncode + NestedDecode + Clone,
368    V: TopEncode + TopDecode + 'static,
369{
370    /// Ensures a value is in the entry by inserting the default if empty, and returns
371    /// an `OccupiedEntry`.
372    pub fn or_insert(self, default: V) -> OccupiedEntry<'a, SA, CurrentStorage, K, V> {
373        match self {
374            Entry::Occupied(entry) => entry,
375            Entry::Vacant(entry) => entry.insert(default),
376        }
377    }
378
379    /// Ensures a value is in the entry by inserting the result of the default function if empty,
380    /// and returns an `OccupiedEntry`.
381    pub fn or_insert_with<F: FnOnce() -> V>(
382        self,
383        default: F,
384    ) -> OccupiedEntry<'a, SA, CurrentStorage, K, V> {
385        match self {
386            Entry::Occupied(entry) => entry,
387            Entry::Vacant(entry) => entry.insert(default()),
388        }
389    }
390
391    /// Ensures a value is in the entry by inserting, if empty, the result of the default function.
392    /// This method allows for generating key-derived values for insertion by providing the default
393    /// function a reference to the key that was moved during the `.entry(key)` method call.
394    ///
395    /// The reference to the moved key is provided so that cloning or copying the key is
396    /// unnecessary, unlike with `.or_insert_with(|| ... )`.
397    pub fn or_insert_with_key<F: FnOnce(&K) -> V>(
398        self,
399        default: F,
400    ) -> OccupiedEntry<'a, SA, CurrentStorage, K, V> {
401        match self {
402            Entry::Occupied(entry) => entry,
403            Entry::Vacant(entry) => {
404                let value = default(entry.key());
405                entry.insert(value)
406            }
407        }
408    }
409
410    /// Provides in-place mutable access to an occupied entry before any
411    /// potential inserts into the map.
412    pub fn and_modify<F>(self, f: F) -> Self
413    where
414        F: FnOnce(&mut V),
415    {
416        match self {
417            Entry::Occupied(mut entry) => {
418                entry.update(f);
419                Entry::Occupied(entry)
420            }
421            Entry::Vacant(entry) => Entry::Vacant(entry),
422        }
423    }
424}
425
426impl<'a, SA, K, V: Default> Entry<'a, SA, CurrentStorage, K, V>
427where
428    SA: StorageMapperApi,
429    K: TopEncode + TopDecode + NestedEncode + NestedDecode + Clone,
430    V: TopEncode + TopDecode + 'static,
431{
432    /// Ensures a value is in the entry by inserting the default value if empty,
433    /// and returns an `OccupiedEntry`.
434    pub fn or_default(self) -> OccupiedEntry<'a, SA, CurrentStorage, K, V> {
435        match self {
436            Entry::Occupied(entry) => entry,
437            Entry::Vacant(entry) => entry.insert(Default::default()),
438        }
439    }
440}
441
442impl<SA, A, K, V> VacantEntry<'_, SA, A, K, V>
443where
444    SA: StorageMapperApi,
445    A: StorageAddress<SA>,
446    K: TopEncode + TopDecode + NestedEncode + NestedDecode + Clone,
447    V: TopEncode + TopDecode + 'static,
448{
449    /// Gets a reference to the key that would be used when inserting a value
450    /// through the VacantEntry.
451    pub fn key(&self) -> &K {
452        &self.key
453    }
454}
455
456impl<'a, SA, K, V> VacantEntry<'a, SA, CurrentStorage, K, V>
457where
458    SA: StorageMapperApi,
459    K: TopEncode + TopDecode + NestedEncode + NestedDecode + Clone,
460    V: TopEncode + TopDecode + 'static,
461{
462    /// Sets the value of the entry with the `VacantEntry`'s key,
463    /// and returns an `OccupiedEntry`.
464    pub fn insert(self, value: V) -> OccupiedEntry<'a, SA, CurrentStorage, K, V> {
465        self.map.insert(self.key.clone(), value);
466        OccupiedEntry {
467            key: self.key,
468            map: self.map,
469            _marker: PhantomData,
470        }
471    }
472}
473
474impl<SA, A, K, V> OccupiedEntry<'_, SA, A, K, V>
475where
476    SA: StorageMapperApi,
477    A: StorageAddress<SA>,
478    K: TopEncode + TopDecode + NestedEncode + NestedDecode + Clone,
479    V: TopEncode + TopDecode + 'static,
480{
481    /// Gets a reference to the key in the entry.
482    pub fn key(&self) -> &K {
483        &self.key
484    }
485
486    /// Gets the value in the entry.
487    pub fn get(&self) -> V {
488        self.map.get(&self.key).unwrap()
489    }
490}
491
492impl<SA, K, V> OccupiedEntry<'_, SA, CurrentStorage, K, V>
493where
494    SA: StorageMapperApi,
495    K: TopEncode + TopDecode + NestedEncode + NestedDecode + Clone,
496    V: TopEncode + TopDecode + 'static,
497{
498    /// Take ownership of the key and value from the map.
499    pub fn remove_entry(self) -> (K, V) {
500        let value = self.map.remove(&self.key).unwrap();
501        (self.key, value)
502    }
503
504    /// Syntactic sugar, to more compactly express a get, update and set in one line.
505    /// Takes whatever lies in storage, apples the given closure and saves the final value back to storage.
506    /// Propagates the return value of the given function.
507    pub fn update<R, F: FnOnce(&mut V) -> R>(&mut self, f: F) -> R {
508        let mut value = self.get();
509        let result = f(&mut value);
510        self.map.insert(self.key.clone(), value);
511        result
512    }
513
514    /// Sets the value of the entry with the `OccupiedEntry`'s key,
515    /// and returns the entry's old value.
516    pub fn insert(self, value: V) -> V {
517        self.map.insert(self.key, value).unwrap()
518    }
519
520    /// Takes the value of the entry out of the map, and returns it.
521    pub fn remove(self) -> V {
522        self.map.remove(&self.key).unwrap()
523    }
524}
525
526/// Behaves like a MultiResultVec<MultiValue2<K, V>> when an endpoint result.
527impl<SA, K, V> TopEncodeMulti for MapMapper<SA, K, V, CurrentStorage>
528where
529    SA: StorageMapperApi,
530    K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
531    V: TopEncode + TopDecode + 'static,
532{
533    fn multi_encode_or_handle_err<O, H>(&self, output: &mut O, h: H) -> Result<(), H::HandledErr>
534    where
535        O: TopEncodeMultiOutput,
536        H: EncodeErrorHandler,
537    {
538        let iter = self.iter().map(MultiValue2::<K, V>::from);
539        multi_encode_iter_or_handle_err(iter, output, h)
540    }
541}
542
543impl<SA, K, V> TypeAbiFrom<MapMapper<SA, K, V, CurrentStorage>>
544    for MultiValueEncoded<SA, MultiValue2<K, V>>
545where
546    SA: StorageMapperApi,
547    K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
548    V: TopEncode + TopDecode + 'static,
549{
550}
551
552impl<SA, K, V> TypeAbiFrom<Self> for MapMapper<SA, K, V, CurrentStorage>
553where
554    SA: StorageMapperApi,
555    K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
556    V: TopEncode + TopDecode + 'static,
557{
558}
559
560/// Behaves like a MultiResultVec<MultiValue<K, V>> when an endpoint result.
561impl<SA, K, V> TypeAbi for MapMapper<SA, K, V, CurrentStorage>
562where
563    SA: StorageMapperApi,
564    K: TopEncode + TopDecode + NestedEncode + NestedDecode + TypeAbi + 'static,
565    V: TopEncode + TopDecode + TypeAbi + 'static,
566{
567    type Unmanaged = Self;
568
569    fn type_name() -> TypeName {
570        MultiValueEncoded::<SA, MultiValue2<K, V>>::type_name()
571    }
572
573    fn type_name_rust() -> TypeName {
574        MultiValueEncoded::<SA, MultiValue2<K, V>>::type_name_rust()
575    }
576
577    fn provide_type_descriptions<TDC: TypeDescriptionContainer>(accumulator: &mut TDC) {
578        K::provide_type_descriptions(accumulator);
579        V::provide_type_descriptions(accumulator);
580    }
581
582    fn is_variadic() -> bool {
583        true
584    }
585}