From ddf455a77ebccf7314b14c88c8eda927b84c60aa Mon Sep 17 00:00:00 2001 From: playX18 <158266309+playX18@users.noreply.github.com> Date: Tue, 13 Aug 2024 17:56:00 +0700 Subject: [PATCH] feat(common): introduce AuxiliaryStorageWrap to remove boilerplate code (#4145) --- common/src/auxiliary/mailbox.rs | 91 ++---------- common/src/auxiliary/mod.rs | 246 ++++++++++++++++++++++++++++++++ 2 files changed, 259 insertions(+), 78 deletions(-) diff --git a/common/src/auxiliary/mailbox.rs b/common/src/auxiliary/mailbox.rs index 0eb3820e5a1..119a20f171b 100644 --- a/common/src/auxiliary/mailbox.rs +++ b/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); - fn contains_keys(key1: &Self::Key1, key2: &Self::Key2) -> bool { - MAILBOX_STORAGE.with_borrow(|map| map.contains_keys(key1, key2)) + fn with_storage(f: F) -> R + where + F: FnOnce(&DoubleBTreeMap) -> R, + { + MAILBOX_STORAGE.with_borrow(f) } - fn get(key1: &Self::Key1, key2: &Self::Key2) -> Option { - 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>( - _key1: Self::Key1, - _key2: Self::Key2, - _f: F, - ) -> R { - unimplemented!() - } - - fn mutate_values 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 { - MAILBOX_STORAGE.with_borrow_mut(|map| map.remove(key1, key2)) - } - - fn clear_prefix(_first_key: Self::Key1) { - unimplemented!() + fn with_storage_mut(f: F) -> R + where + F: FnOnce(&mut DoubleBTreeMap) -> 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)> for MailboxStorageWrap { - type Key = ProgramId; - - type DrainIter = IteratorWrap< - IntoIter)>, - (MailboxedMessage, Interval), - GetSecondPos, - >; - - type Iter = IteratorWrap< - IntoIter)>, - (MailboxedMessage, Interval), - 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 { diff --git a/common/src/auxiliary/mod.rs b/common/src/auxiliary/mod.rs index c3a3866b009..18098a8b5d9 100644 --- a/common/src/auxiliary/mod.rs +++ b/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 Default for DoubleBTreeMap { 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: F) -> R + where + F: FnOnce(&DoubleBTreeMap) -> R; + + fn with_storage_mut(f: F) -> R + where + F: FnOnce(&mut DoubleBTreeMap) -> R; +} + +impl DoubleMapStorage for T { + type Key1 = T::Key1; + type Key2 = T::Key2; + type Value = T::Value; + + fn get(key1: &Self::Key1, key2: &Self::Key2) -> Option { + 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>( + 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>( + key1: Self::Key1, + key2: Self::Key2, + f: F, + ) -> Option { + 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 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 { + T::with_storage_mut(|map| map.remove(key1, key2)) + } +} + +impl IterableByKeyMap for T { + type Key = T::Key1; + + type DrainIter = IteratorWrap, T::Value, GetSecondPos>; + + type Iter = IteratorWrap, 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 KeyIterableByKeyMap for T { + type Key1 = T::Key1; + type Key2 = T::Key2; + type DrainIter = IteratorWrap, T::Key2, GetFirstPos>; + type Iter = IteratorWrap, 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 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: F) -> R + where + F: FnOnce(&BTreeMap) -> R; + + fn with_storage_mut(f: F) -> R + where + F: FnOnce(&mut BTreeMap) -> R; +} + +impl 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 { + 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>(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>(key: Self::Key, f: F) -> Option { + T::with_storage_mut(|map| map.get_mut(&key).map(f)) + } + + fn mutate_values 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 { + T::with_storage_mut(|map| map.remove(&key)) + } +} + +impl Counted for T { + type Length = usize; + fn len() -> Self::Length { + T::with_storage(|map| map.len()) + } +}