1use core::marker::PhantomData;
2
3use super::{
4 SetMapper, StorageClearable, StorageMapper, StorageMapperFromAddress,
5 set_mapper::{self},
6 source::{CurrentStorage, StorageAddress},
7};
8use crate::{
9 api::StorageMapperApi,
10 codec::{NestedDecode, NestedEncode, TopDecode, TopEncode},
11 contract_base::ErrorHelper,
12 storage::{self, StorageKey},
13 types::ManagedAddress,
14};
15
16const MAPPED_STORAGE_VALUE_IDENTIFIER: &[u8] = b".storage";
17type Keys<'a, SA, A, T> = set_mapper::Iter<'a, SA, A, T>;
18
19pub struct MapStorageMapper<SA, K, V, A = CurrentStorage>
116where
117 SA: StorageMapperApi,
118 K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
119 A: StorageAddress<SA>,
120 V: StorageMapper<SA> + StorageClearable,
121{
122 _phantom_api: PhantomData<SA>,
123 base_key: StorageKey<SA>,
124 keys_set: SetMapper<SA, K, A>,
125 _phantom_value: PhantomData<V>,
126}
127
128impl<SA, K, V> StorageMapper<SA> for MapStorageMapper<SA, K, V, CurrentStorage>
129where
130 SA: StorageMapperApi,
131 K: TopEncode + TopDecode + NestedEncode + NestedDecode,
132 V: StorageMapper<SA> + StorageClearable,
133{
134 fn new(base_key: StorageKey<SA>) -> Self {
135 Self {
136 _phantom_api: PhantomData,
137 base_key: base_key.clone(),
138 keys_set: SetMapper::new(base_key),
139 _phantom_value: PhantomData,
140 }
141 }
142}
143
144impl<SA, K, V> StorageClearable for MapStorageMapper<SA, K, V, CurrentStorage>
145where
146 SA: StorageMapperApi,
147 K: TopEncode + TopDecode + NestedEncode + NestedDecode,
148 V: StorageMapper<SA> + StorageClearable,
149{
150 fn clear(&mut self) {
151 for mut value in self.values() {
152 value.clear();
153 }
154 self.keys_set.clear();
155 }
156}
157
158impl<SA, K, V> MapStorageMapper<SA, K, V, CurrentStorage>
159where
160 SA: StorageMapperApi,
161 K: TopEncode + TopDecode + NestedEncode + NestedDecode,
162 V: StorageMapper<SA> + StorageClearable,
163{
164 pub fn insert_default(&mut self, k: K) -> bool {
170 self.keys_set.insert(k)
171 }
172
173 pub fn remove(&mut self, k: &K) -> bool {
179 if self.keys_set.remove(k) {
180 self.get_mapped_storage_value(k).clear();
181 return true;
182 }
183 false
184 }
185}
186
187impl<SA, K, V> StorageMapperFromAddress<SA> for MapStorageMapper<SA, K, V, ManagedAddress<SA>>
188where
189 SA: StorageMapperApi,
190 K: TopEncode + TopDecode + NestedEncode + NestedDecode,
191 V: StorageMapper<SA> + StorageClearable,
192{
193 fn new_from_address(address: ManagedAddress<SA>, base_key: StorageKey<SA>) -> Self {
194 MapStorageMapper {
195 _phantom_api: PhantomData,
196 base_key: base_key.clone(),
197 keys_set: SetMapper::new_from_address(address, base_key),
198 _phantom_value: PhantomData,
199 }
200 }
201}
202
203impl<SA, A, K, V> MapStorageMapper<SA, K, V, A>
204where
205 SA: StorageMapperApi,
206 A: StorageAddress<SA>,
207 K: TopEncode + TopDecode + NestedEncode + NestedDecode,
208 V: StorageMapper<SA> + StorageClearable,
209{
210 fn build_named_key(&self, name: &[u8], key: &K) -> StorageKey<SA> {
211 let mut named_key = self.base_key.clone();
212 named_key.append_bytes(name);
213 named_key.append_item(key);
214 named_key
215 }
216
217 fn get_mapped_storage_value(&self, key: &K) -> V {
218 let key = self.build_named_key(MAPPED_STORAGE_VALUE_IDENTIFIER, key);
219 <V as storage::mappers::StorageMapper<SA>>::new(key)
220 }
221
222 pub fn get(&self, k: &K) -> Option<V> {
224 if self.keys_set.contains(k) {
225 return Some(self.get_mapped_storage_value(k));
226 }
227 None
228 }
229
230 pub fn keys(&self) -> Keys<'_, SA, A, K> {
231 self.keys_set.iter()
232 }
233
234 pub fn is_empty(&self) -> bool {
236 self.keys_set.is_empty()
237 }
238
239 pub fn len(&self) -> usize {
241 self.keys_set.len()
242 }
243
244 pub fn contains_key(&self, k: &K) -> bool {
246 self.keys_set.contains(k)
247 }
248
249 pub fn entry(&mut self, key: K) -> Entry<'_, SA, A, K, V> {
251 if self.contains_key(&key) {
252 Entry::Occupied(OccupiedEntry {
253 key,
254 map: self,
255 _marker: PhantomData,
256 })
257 } else {
258 Entry::Vacant(VacantEntry {
259 key,
260 map: self,
261 _marker: PhantomData,
262 })
263 }
264 }
265
266 pub fn values(&self) -> Values<'_, SA, A, K, V> {
269 Values::new(self)
270 }
271
272 pub fn iter(&self) -> Iter<'_, SA, A, K, V> {
275 Iter::new(self)
276 }
277}
278
279impl<'a, SA, A, K, V> IntoIterator for &'a MapStorageMapper<SA, K, V, A>
280where
281 SA: StorageMapperApi,
282 A: StorageAddress<SA>,
283 K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
284 V: StorageMapper<SA> + StorageClearable,
285{
286 type Item = (K, V);
287
288 type IntoIter = Iter<'a, SA, A, K, V>;
289
290 fn into_iter(self) -> Self::IntoIter {
291 self.iter()
292 }
293}
294
295pub struct Iter<'a, SA, A, K, V>
296where
297 SA: StorageMapperApi,
298 K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
299 A: StorageAddress<SA>,
300 V: StorageMapper<SA> + StorageClearable,
301{
302 key_iter: Keys<'a, SA, A, K>,
303 hash_map: &'a MapStorageMapper<SA, K, V, A>,
304}
305
306impl<'a, SA, A, K, V> Iter<'a, SA, A, K, V>
307where
308 SA: StorageMapperApi,
309 A: StorageAddress<SA>,
310 K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
311 V: StorageMapper<SA> + StorageClearable,
312{
313 fn new(hash_map: &'a MapStorageMapper<SA, K, V, A>) -> Iter<'a, SA, A, K, V> {
314 Iter {
315 key_iter: hash_map.keys(),
316 hash_map,
317 }
318 }
319}
320
321impl<SA, A, K, V> Iterator for Iter<'_, SA, A, K, V>
322where
323 SA: StorageMapperApi,
324 A: StorageAddress<SA>,
325 K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
326 V: StorageMapper<SA> + StorageClearable,
327{
328 type Item = (K, V);
329
330 #[inline]
331 fn next(&mut self) -> Option<(K, V)> {
332 if let Some(key) = self.key_iter.next() {
333 let Some(value) = self.hash_map.get(&key) else {
334 ErrorHelper::<SA>::signal_error_with_message("missing key")
335 };
336 return Some((key, value));
337 }
338 None
339 }
340}
341
342pub struct Values<'a, SA, A, K, V>
343where
344 SA: StorageMapperApi,
345 K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
346 A: StorageAddress<SA>,
347 V: StorageMapper<SA> + StorageClearable,
348{
349 key_iter: Keys<'a, SA, A, K>,
350 hash_map: &'a MapStorageMapper<SA, K, V, A>,
351}
352
353impl<'a, SA, A, K, V> Values<'a, SA, A, K, V>
354where
355 SA: StorageMapperApi,
356 A: StorageAddress<SA>,
357 K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
358 V: StorageMapper<SA> + StorageClearable,
359{
360 fn new(hash_map: &'a MapStorageMapper<SA, K, V, A>) -> Values<'a, SA, A, K, V> {
361 Values {
362 key_iter: hash_map.keys(),
363 hash_map,
364 }
365 }
366}
367
368impl<SA, A, K, V> Iterator for Values<'_, SA, A, K, V>
369where
370 SA: StorageMapperApi,
371 A: StorageAddress<SA>,
372 K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
373 V: StorageMapper<SA> + StorageClearable,
374{
375 type Item = V;
376
377 #[inline]
378 fn next(&mut self) -> Option<V> {
379 if let Some(key) = self.key_iter.next() {
380 let value = self.hash_map.get(&key).unwrap();
381 return Some(value);
382 }
383 None
384 }
385}
386
387pub enum Entry<'a, SA, A, K: 'a, V: 'a>
388where
389 SA: StorageMapperApi,
390 K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
391 A: StorageAddress<SA>,
392 V: StorageMapper<SA> + StorageClearable,
393{
394 Vacant(VacantEntry<'a, SA, A, K, V>),
396
397 Occupied(OccupiedEntry<'a, SA, A, K, V>),
399}
400
401pub struct VacantEntry<'a, SA, A, K: 'a, V: 'a>
404where
405 SA: StorageMapperApi,
406 K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
407 A: StorageAddress<SA>,
408 V: StorageMapper<SA> + StorageClearable,
409{
410 pub(super) key: K,
411 pub(super) map: &'a mut MapStorageMapper<SA, K, V, A>,
412
413 pub(super) _marker: PhantomData<&'a mut (K, V)>,
415}
416
417pub struct OccupiedEntry<'a, SA, A, K: 'a, V: 'a>
420where
421 SA: StorageMapperApi,
422 K: TopEncode + TopDecode + NestedEncode + NestedDecode + 'static,
423 A: StorageAddress<SA>,
424 V: StorageMapper<SA> + StorageClearable,
425{
426 pub(super) key: K,
427 pub(super) map: &'a mut MapStorageMapper<SA, K, V, A>,
428
429 pub(super) _marker: PhantomData<&'a mut (K, V)>,
431}
432
433impl<'a, SA, K, V> Entry<'a, SA, CurrentStorage, K, V>
434where
435 SA: StorageMapperApi,
436 K: TopEncode + TopDecode + NestedEncode + NestedDecode + Clone + 'static,
437 V: StorageMapper<SA> + StorageClearable,
438{
439 pub fn or_insert_default(self) -> OccupiedEntry<'a, SA, CurrentStorage, K, V> {
442 match self {
443 Entry::Occupied(entry) => entry,
444 Entry::Vacant(entry) => entry.insert_default(),
445 }
446 }
447
448 pub fn key(&self) -> &K {
450 match *self {
451 Entry::Occupied(ref entry) => entry.key(),
452 Entry::Vacant(ref entry) => entry.key(),
453 }
454 }
455
456 pub fn and_modify<F>(self, f: F) -> Self
459 where
460 F: FnOnce(&mut V),
461 {
462 match self {
463 Entry::Occupied(mut entry) => {
464 entry.update(f);
465 Entry::Occupied(entry)
466 }
467 Entry::Vacant(entry) => Entry::Vacant(entry),
468 }
469 }
470}
471
472impl<'a, SA, K, V> Entry<'a, SA, CurrentStorage, K, V>
473where
474 SA: StorageMapperApi,
475 K: TopEncode + TopDecode + NestedEncode + NestedDecode + Clone + 'static,
476 V: StorageMapper<SA> + StorageClearable,
477{
478 pub fn or_default(self) -> OccupiedEntry<'a, SA, CurrentStorage, K, V> {
481 match self {
482 Entry::Occupied(entry) => entry,
483 Entry::Vacant(entry) => entry.insert_default(),
484 }
485 }
486}
487
488impl<SA, A, K, V> VacantEntry<'_, SA, A, K, V>
489where
490 SA: StorageMapperApi,
491 A: StorageAddress<SA>,
492 K: TopEncode + TopDecode + NestedEncode + NestedDecode + Clone + 'static,
493 V: StorageMapper<SA> + StorageClearable,
494{
495 pub fn key(&self) -> &K {
498 &self.key
499 }
500}
501
502impl<'a, SA, K, V> VacantEntry<'a, SA, CurrentStorage, K, V>
503where
504 SA: StorageMapperApi,
505 K: TopEncode + TopDecode + NestedEncode + NestedDecode + Clone + 'static,
506 V: StorageMapper<SA> + StorageClearable,
507{
508 pub fn insert_default(self) -> OccupiedEntry<'a, SA, CurrentStorage, K, V> {
511 self.map.insert_default(self.key.clone());
512 OccupiedEntry {
513 key: self.key,
514 map: self.map,
515 _marker: PhantomData,
516 }
517 }
518}
519
520impl<SA, A, K, V> OccupiedEntry<'_, SA, A, K, V>
521where
522 SA: StorageMapperApi,
523 A: StorageAddress<SA>,
524 K: TopEncode + TopDecode + NestedEncode + NestedDecode + Clone + 'static,
525 V: StorageMapper<SA> + StorageClearable,
526{
527 pub fn key(&self) -> &K {
529 &self.key
530 }
531
532 pub fn get(&self) -> V {
534 self.map.get(&self.key).unwrap()
535 }
536}
537
538impl<SA, K, V> OccupiedEntry<'_, SA, CurrentStorage, K, V>
539where
540 SA: StorageMapperApi,
541 K: TopEncode + TopDecode + NestedEncode + NestedDecode + Clone + 'static,
542 V: StorageMapper<SA> + StorageClearable,
543{
544 pub fn update<R, F: FnOnce(&mut V) -> R>(&mut self, f: F) -> R {
548 let mut value = self.get();
549 f(&mut value)
550 }
551
552 pub fn remove(self) {
554 self.map.remove(&self.key);
555 }
556}