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 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 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 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 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 pub fn is_empty(&self) -> bool {
174 self.keys_set.is_empty()
175 }
176
177 pub fn len(&self) -> usize {
179 self.keys_set.len()
180 }
181
182 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 pub fn values(&self) -> Values<'_, SA, A, K, V> {
202 Values::new(self)
203 }
204
205 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 Vacant(VacantEntry<'a, SA, A, K, V>),
311
312 Occupied(OccupiedEntry<'a, SA, A, K, V>),
314}
315
316pub 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 pub(super) _marker: PhantomData<&'a mut (K, V)>,
330}
331
332pub 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 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 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 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 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 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 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 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 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 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 pub fn key(&self) -> &K {
483 &self.key
484 }
485
486 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 pub fn remove_entry(self) -> (K, V) {
500 let value = self.map.remove(&self.key).unwrap();
501 (self.key, value)
502 }
503
504 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 pub fn insert(self, value: V) -> V {
517 self.map.insert(self.key, value).unwrap()
518 }
519
520 pub fn remove(self) -> V {
522 self.map.remove(&self.key).unwrap()
523 }
524}
525
526impl<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
560impl<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}