feat(common): introduce AuxiliaryStorageWrap to remove boilerplate co…

…de (#4145)

common/src/auxiliary/mailbox.rs

@@ -18,14 +18,11 @@
 
 //! Auxiliary implementation of the mailbox.
 
+use super::AuxiliaryDoubleStorageWrap;
 use crate::{
     auxiliary::DoubleBTreeMap,
-    storage::{
-        CountedByKey, DoubleMapStorage, GetSecondPos, Interval, IterableByKeyMap, IteratorWrap,
-        MailboxError, MailboxImpl, MailboxKeyGen,
-    },
+    storage::{Interval, MailboxError, MailboxImpl, MailboxKeyGen},
 };
-use alloc::collections::btree_map::IntoIter;
 use core::cell::RefCell;
 use gear_core::{
     ids::{MessageId, ProgramId},
@@ -55,87 +52,25 @@ std::thread_local! {
 /// `Mailbox` double storage map manager.
 pub struct MailboxStorageWrap;
 
-impl DoubleMapStorage for MailboxStorageWrap {
+impl AuxiliaryDoubleStorageWrap for MailboxStorageWrap {
     type Key1 = ProgramId;
     type Key2 = MessageId;
     type Value = (MailboxedMessage, Interval<BlockNumber>);
 
-    fn contains_keys(key1: &Self::Key1, key2: &Self::Key2) -> bool {
-        MAILBOX_STORAGE.with_borrow(|map| map.contains_keys(key1, key2))
+    fn with_storage<F, R>(f: F) -> R
+    where
+        F: FnOnce(&DoubleBTreeMap<Self::Key1, Self::Key2, Self::Value>) -> R,
+    {
+        MAILBOX_STORAGE.with_borrow(f)
     }
 
-    fn get(key1: &Self::Key1, key2: &Self::Key2) -> Option<Self::Value> {
-        MAILBOX_STORAGE.with_borrow(|map| map.get(key1, key2).cloned())
-    }
-
-    fn insert(key1: Self::Key1, key2: Self::Key2, value: Self::Value) {
-        MAILBOX_STORAGE.with_borrow_mut(|map| map.insert(key1, key2, value));
-    }
-
-    fn mutate<R, F: FnOnce(&mut Option<Self::Value>) -> R>(
-        _key1: Self::Key1,
-        _key2: Self::Key2,
-        _f: F,
-    ) -> R {
-        unimplemented!()
-    }
-
-    fn mutate_values<F: FnMut(Self::Value) -> Self::Value>(_f: F) {
-        unimplemented!()
-    }
-
-    fn remove(key1: Self::Key1, key2: Self::Key2) {
-        Self::take(key1, key2);
-    }
-
-    fn clear() {
-        MAILBOX_STORAGE.with_borrow_mut(|map| map.clear())
-    }
-
-    fn take(key1: Self::Key1, key2: Self::Key2) -> Option<Self::Value> {
-        MAILBOX_STORAGE.with_borrow_mut(|map| map.remove(key1, key2))
-    }
-
-    fn clear_prefix(_first_key: Self::Key1) {
-        unimplemented!()
+    fn with_storage_mut<F, R>(f: F) -> R
+    where
+        F: FnOnce(&mut DoubleBTreeMap<Self::Key1, Self::Key2, Self::Value>) -> R,
+    {
+        MAILBOX_STORAGE.with_borrow_mut(f)
     }
 }
-
-impl CountedByKey for MailboxStorageWrap {
-    type Key = ProgramId;
-    type Length = usize;
-
-    fn len(key: &Self::Key) -> Self::Length {
-        MAILBOX_STORAGE.with_borrow(|map| map.count_key(key))
-    }
-}
-
-impl IterableByKeyMap<(MailboxedMessage, Interval<BlockNumber>)> for MailboxStorageWrap {
-    type Key = ProgramId;
-
-    type DrainIter = IteratorWrap<
-        IntoIter<MessageId, (MailboxedMessage, Interval<BlockNumber>)>,
-        (MailboxedMessage, Interval<BlockNumber>),
-        GetSecondPos,
-    >;
-
-    type Iter = IteratorWrap<
-        IntoIter<MessageId, (MailboxedMessage, Interval<BlockNumber>)>,
-        (MailboxedMessage, Interval<BlockNumber>),
-        GetSecondPos,
-    >;
-
-    fn drain_key(key: Self::Key) -> Self::DrainIter {
-        MAILBOX_STORAGE
-            .with_borrow_mut(|map| map.drain_key(&key))
-            .into()
-    }
-
-    fn iter_key(key: Self::Key) -> Self::Iter {
-        MAILBOX_STORAGE.with_borrow(|map| map.iter_key(&key)).into()
-    }
-}
-
 /// An implementor of the error returned from calling `Mailbox` trait functions.
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub enum MailboxErrorImpl {

common/src/auxiliary/mod.rs

@@ -23,6 +23,10 @@
 pub mod gas_provider;
 pub mod mailbox;
 
+use crate::storage::{
+    Counted, CountedByKey, DoubleMapStorage, GetFirstPos, GetSecondPos, IterableByKeyMap,
+    IteratorWrap, KeyIterableByKeyMap, MapStorage,
+};
 use alloc::collections::btree_map::{BTreeMap, Entry, IntoIter};
 
 /// Double key `BTreeMap`.
@@ -136,3 +140,245 @@ impl<K1, K2, V> Default for DoubleBTreeMap<K1, K2, V> {
         Self::new()
     }
 }
+
+/// An auxiliary storage wrapper type.
+///
+/// Implements DoubleMapStorage and traits like [`IterableByKeyMap`] for such type automatically.
+pub trait AuxiliaryDoubleStorageWrap {
+    type Key1: Ord + Clone;
+    type Key2: Ord + Clone;
+    type Value: Clone;
+    fn with_storage<F, R>(f: F) -> R
+    where
+        F: FnOnce(&DoubleBTreeMap<Self::Key1, Self::Key2, Self::Value>) -> R;
+
+    fn with_storage_mut<F, R>(f: F) -> R
+    where
+        F: FnOnce(&mut DoubleBTreeMap<Self::Key1, Self::Key2, Self::Value>) -> R;
+}
+
+impl<T: AuxiliaryDoubleStorageWrap> DoubleMapStorage for T {
+    type Key1 = T::Key1;
+    type Key2 = T::Key2;
+    type Value = T::Value;
+
+    fn get(key1: &Self::Key1, key2: &Self::Key2) -> Option<Self::Value> {
+        T::with_storage(|map| map.get(key1, key2).cloned())
+    }
+
+    fn insert(key1: Self::Key1, key2: Self::Key2, value: Self::Value) {
+        T::with_storage_mut(|map| map.insert(key1, key2, value));
+    }
+
+    fn clear() {
+        T::with_storage_mut(|map| map.clear());
+    }
+
+    fn clear_prefix(first_key: Self::Key1) {
+        T::with_storage_mut(|map| {
+            let keys = map.iter_key(&first_key).map(|(k, _)| k.clone());
+            for key in keys {
+                map.remove(first_key.clone(), key);
+            }
+        });
+    }
+
+    fn contains_keys(key1: &Self::Key1, key2: &Self::Key2) -> bool {
+        T::with_storage_mut(|map| map.contains_keys(key1, key2))
+    }
+
+    fn mutate<R, F: FnOnce(&mut Option<Self::Value>) -> R>(
+        key1: Self::Key1,
+        key2: Self::Key2,
+        f: F,
+    ) -> R {
+        T::with_storage_mut(|map| {
+            let inner_map = map.inner.entry(key1).or_default();
+            match inner_map.entry(key2) {
+                Entry::Occupied(mut occupied) => {
+                    let mut value = Some(occupied.get().clone());
+                    let result = f(&mut value);
+                    if let Some(value) = value {
+                        *occupied.get_mut() = value;
+                    } else {
+                        occupied.remove();
+                    }
+
+                    result
+                }
+
+                Entry::Vacant(vacant) => {
+                    let mut value = None;
+                    let result = f(&mut value);
+                    if let Some(value) = value {
+                        vacant.insert(value);
+                    }
+                    result
+                }
+            }
+        })
+    }
+
+    fn mutate_exists<R, F: FnOnce(&mut Self::Value) -> R>(
+        key1: Self::Key1,
+        key2: Self::Key2,
+        f: F,
+    ) -> Option<R> {
+        T::with_storage_mut(|map| {
+            if let Some(inner_map) = map.inner.get_mut(&key1) {
+                if let Some(value) = inner_map.get_mut(&key2) {
+                    return Some(f(value));
+                }
+            }
+
+            None
+        })
+    }
+
+    fn mutate_values<F: FnMut(Self::Value) -> Self::Value>(mut f: F) {
+        T::with_storage_mut(|map| {
+            for (_, inner_map) in map.inner.iter_mut() {
+                for (_, value) in inner_map.iter_mut() {
+                    *value = f(value.clone());
+                }
+            }
+        });
+    }
+
+    fn remove(key1: Self::Key1, key2: Self::Key2) {
+        Self::take(key1, key2);
+    }
+
+    fn take(key1: Self::Key1, key2: Self::Key2) -> Option<Self::Value> {
+        T::with_storage_mut(|map| map.remove(key1, key2))
+    }
+}
+
+impl<T: AuxiliaryDoubleStorageWrap> IterableByKeyMap<T::Value> for T {
+    type Key = T::Key1;
+
+    type DrainIter = IteratorWrap<IntoIter<T::Key2, T::Value>, T::Value, GetSecondPos>;
+
+    type Iter = IteratorWrap<IntoIter<T::Key2, T::Value>, T::Value, GetSecondPos>;
+
+    fn drain_key(key: Self::Key) -> Self::DrainIter {
+        T::with_storage_mut(|map| map.drain_key(&key)).into()
+    }
+
+    fn iter_key(key: Self::Key) -> Self::Iter {
+        T::with_storage(|map| map.iter_key(&key)).into()
+    }
+}
+
+impl<T: AuxiliaryDoubleStorageWrap> KeyIterableByKeyMap for T {
+    type Key1 = T::Key1;
+    type Key2 = T::Key2;
+    type DrainIter = IteratorWrap<IntoIter<T::Key2, T::Value>, T::Key2, GetFirstPos>;
+    type Iter = IteratorWrap<IntoIter<T::Key2, T::Value>, T::Key2, GetFirstPos>;
+
+    fn drain_prefix_keys(key: Self::Key1) -> Self::DrainIter {
+        T::with_storage_mut(|map| map.drain_key(&key).into())
+    }
+
+    fn iter_prefix_keys(key: Self::Key1) -> Self::Iter {
+        T::with_storage(|map| map.iter_key(&key)).into()
+    }
+}
+
+impl<T: AuxiliaryDoubleStorageWrap> CountedByKey for T {
+    type Key = T::Key1;
+    type Length = usize;
+
+    fn len(key: &Self::Key) -> Self::Length {
+        T::with_storage(|map| map.count_key(key))
+    }
+}
+
+pub trait AuxiliaryStorageWrap {
+    type Key: Clone + Ord;
+    type Value: Clone;
+
+    fn with_storage<F, R>(f: F) -> R
+    where
+        F: FnOnce(&BTreeMap<Self::Key, Self::Value>) -> R;
+
+    fn with_storage_mut<F, R>(f: F) -> R
+    where
+        F: FnOnce(&mut BTreeMap<Self::Key, Self::Value>) -> R;
+}
+
+impl<T: AuxiliaryStorageWrap> MapStorage for T {
+    type Key = T::Key;
+    type Value = T::Value;
+
+    fn clear() {
+        T::with_storage_mut(|map| map.clear());
+    }
+
+    fn contains_key(key: &Self::Key) -> bool {
+        T::with_storage(|map| map.contains_key(key))
+    }
+
+    fn get(key: &Self::Key) -> Option<Self::Value> {
+        T::with_storage(|map| map.get(key).cloned())
+    }
+
+    fn insert(key: Self::Key, value: Self::Value) {
+        T::with_storage_mut(|map| map.insert(key, value));
+    }
+
+    fn mutate<R, F: FnOnce(&mut Option<Self::Value>) -> R>(key: Self::Key, f: F) -> R {
+        T::with_storage_mut(|map| match map.entry(key) {
+            Entry::Occupied(mut occupied) => {
+                let mut value = Some(occupied.get().clone());
+
+                let result = f(&mut value);
+                if let Some(value) = value.take() {
+                    *occupied.get_mut() = value;
+                } else {
+                    occupied.remove();
+                }
+
+                result
+            }
+
+            Entry::Vacant(vacant) => {
+                let mut value = None;
+
+                let result = f(&mut value);
+
+                if let Some(value) = value.take() {
+                    vacant.insert(value);
+                }
+
+                result
+            }
+        })
+    }
+
+    fn mutate_exists<R, F: FnOnce(&mut Self::Value) -> R>(key: Self::Key, f: F) -> Option<R> {
+        T::with_storage_mut(|map| map.get_mut(&key).map(f))
+    }
+
+    fn mutate_values<F: FnMut(Self::Value) -> Self::Value>(mut f: F) {
+        T::with_storage_mut(|map| {
+            map.iter_mut()
+                .for_each(|(_, value)| *value = f(value.clone()))
+        });
+    }
+
+    fn remove(key: Self::Key) {
+        Self::take(key);
+    }
+
+    fn take(key: Self::Key) -> Option<Self::Value> {
+        T::with_storage_mut(|map| map.remove(&key))
+    }
+}
+
+impl<T: AuxiliaryStorageWrap> Counted for T {
+    type Length = usize;
+    fn len() -> Self::Length {
+        T::with_storage(|map| map.len())
+    }
+}