diff --git a/frame/support/src/storage/bounded_btree_map.rs b/frame/support/src/storage/bounded_btree_map.rs new file mode 100644 index 0000000000000..7fd0d175fda99 --- /dev/null +++ b/frame/support/src/storage/bounded_btree_map.rs @@ -0,0 +1,421 @@ +// This file is part of Substrate. + +// Copyright (C) 2017-2021 Parity Technologies (UK) Ltd. +// SPDX-License-Identifier: Apache-2.0 + +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Traits, types and structs to support a bounded BTreeMap. + +use sp_std::{ + borrow::Borrow, collections::btree_map::BTreeMap, convert::TryFrom, fmt, marker::PhantomData, + ops::Deref, +}; +use crate::{ + storage::StorageDecodeLength, + traits::{Get, MaxEncodedLen}, +}; +use codec::{Encode, Decode}; + +/// A bounded map based on a B-Tree. +/// +/// B-Trees represent a fundamental compromise between cache-efficiency and actually minimizing +/// the amount of work performed in a search. See [`BTreeMap`] for more details. +/// +/// Unlike a standard `BTreeMap`, there is a static, enforced upper limit to the number of items +/// in the map. All internal operations ensure this bound is respected. +#[derive(Encode, Decode)] +pub struct BoundedBTreeMap(BTreeMap, PhantomData); + +impl BoundedBTreeMap +where + S: Get, +{ + /// Get the bound of the type in `usize`. + pub fn bound() -> usize { + S::get() as usize + } +} + +impl BoundedBTreeMap +where + K: Ord, + S: Get, +{ + /// Create a new `BoundedBTreeMap`. + /// + /// Does not allocate. + pub fn new() -> Self { + BoundedBTreeMap(BTreeMap::new(), PhantomData) + } + + /// Create `Self` from a primitive `BTreeMap` without any checks. + unsafe fn unchecked_from(map: BTreeMap) -> Self { + Self(map, Default::default()) + } + + /// Create `Self` from a primitive `BTreeMap` without any checks. + /// + /// Logs warnings if the bound is not being respected. The scope is mentioned in the log message + /// to indicate where overflow is happening. + /// + /// # Example + /// + /// ``` + /// # use sp_std::collections::btree_map::BTreeMap; + /// # use frame_support::{parameter_types, storage::bounded_btree_map::BoundedBTreeMap}; + /// parameter_types! { + /// pub const Size: u32 = 5; + /// } + /// let mut map = BTreeMap::new(); + /// map.insert("foo", 1); + /// map.insert("bar", 2); + /// let bounded_map = unsafe {BoundedBTreeMap::<_, _, Size>::force_from(map, "demo")}; + /// ``` + pub unsafe fn force_from(map: BTreeMap, scope: Scope) -> Self + where + Scope: Into>, + { + if map.len() > Self::bound() { + log::warn!( + target: crate::LOG_TARGET, + "length of a bounded btreemap in scope {} is not respected.", + scope.into().unwrap_or("UNKNOWN"), + ); + } + + Self::unchecked_from(map) + } + + /// Consume self, and return the inner `BTreeMap`. + /// + /// This is useful when a mutating API of the inner type is desired, and closure-based mutation + /// such as provided by [`try_mutate`][Self::try_mutate] is inconvenient. + pub fn into_inner(self) -> BTreeMap { + debug_assert!(self.0.len() <= Self::bound()); + self.0 + } + + /// Consumes self and mutates self via the given `mutate` function. + /// + /// If the outcome of mutation is within bounds, `Some(Self)` is returned. Else, `None` is + /// returned. + /// + /// This is essentially a *consuming* shorthand [`Self::into_inner`] -> `...` -> + /// [`Self::try_from`]. + pub fn try_mutate(mut self, mut mutate: impl FnMut(&mut BTreeMap)) -> Option { + mutate(&mut self.0); + (self.0.len() <= Self::bound()).then(move || self) + } + + // Clears the map, removing all elements. + pub fn clear(&mut self) { + self.0.clear() + } + + /// Return a mutable reference to the value corresponding to the key. + /// + /// The key may be any borrowed form of the map's key type, but the ordering on the borrowed + /// form _must_ match the ordering on the key type. + pub fn get_mut(&mut self, key: &Q) -> Option<&mut V> + where + K: Borrow, + Q: Ord + ?Sized, + { + self.0.get_mut(key) + } + + /// Exactly the same semantics as [`BTreeMap::insert`], but returns an `Err` (and is a noop) if the + /// new length of the map exceeds `S`. + pub fn try_insert(&mut self, key: K, value: V) -> Result<(), ()> { + if self.len() < Self::bound() { + self.0.insert(key, value); + Ok(()) + } else { + Err(()) + } + } + + /// Remove a key from the map, returning the value at the key if the key was previously in the map. + /// + /// The key may be any borrowed form of the map's key type, but the ordering on the borrowed + /// form _must_ match the ordering on the key type. + pub fn remove(&mut self, key: &Q) -> Option + where + K: Borrow, + Q: Ord + ?Sized, + { + self.0.remove(key) + } + + /// Remove a key from the map, returning the value at the key if the key was previously in the map. + /// + /// The key may be any borrowed form of the map's key type, but the ordering on the borrowed + /// form _must_ match the ordering on the key type. + pub fn remove_entry(&mut self, key: &Q) -> Option<(K, V)> + where + K: Borrow, + Q: Ord + ?Sized, + { + self.0.remove_entry(key) + } +} + +impl Default for BoundedBTreeMap +where + K: Ord, + S: Get, +{ + fn default() -> Self { + Self::new() + } +} + +impl Clone for BoundedBTreeMap +where + BTreeMap: Clone, +{ + fn clone(&self) -> Self { + BoundedBTreeMap(self.0.clone(), PhantomData) + } +} + +#[cfg(feature = "std")] +impl fmt::Debug for BoundedBTreeMap +where + BTreeMap: fmt::Debug, + S: Get, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("BoundedBTreeMap").field(&self.0).field(&Self::bound()).finish() + } +} + +impl PartialEq for BoundedBTreeMap +where + BTreeMap: PartialEq, +{ + fn eq(&self, other: &Self) -> bool { + self.0 == other.0 + } +} + +impl Eq for BoundedBTreeMap where BTreeMap: Eq {} + +impl PartialEq> for BoundedBTreeMap +where + BTreeMap: PartialEq, +{ + fn eq(&self, other: &BTreeMap) -> bool { + self.0 == *other + } +} + +impl PartialOrd for BoundedBTreeMap +where + BTreeMap: PartialOrd, +{ + fn partial_cmp(&self, other: &Self) -> Option { + self.0.partial_cmp(&other.0) + } +} + +impl Ord for BoundedBTreeMap +where + BTreeMap: Ord, +{ + fn cmp(&self, other: &Self) -> sp_std::cmp::Ordering { + self.0.cmp(&other.0) + } +} + +impl IntoIterator for BoundedBTreeMap { + type Item = (K, V); + type IntoIter = sp_std::collections::btree_map::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + +impl MaxEncodedLen for BoundedBTreeMap +where + K: MaxEncodedLen, + V: MaxEncodedLen, + S: Get, +{ + fn max_encoded_len() -> usize { + Self::bound() + .saturating_mul(K::max_encoded_len().saturating_add(V::max_encoded_len())) + .saturating_add(codec::Compact(S::get()).encoded_size()) + } +} + +impl Deref for BoundedBTreeMap +where + K: Ord, +{ + type Target = BTreeMap; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl AsRef> for BoundedBTreeMap +where + K: Ord, +{ + fn as_ref(&self) -> &BTreeMap { + &self.0 + } +} + +impl From> for BTreeMap +where + K: Ord, +{ + fn from(map: BoundedBTreeMap) -> Self { + map.0 + } +} + +impl TryFrom> for BoundedBTreeMap +where + K: Ord, + S: Get, +{ + type Error = (); + + fn try_from(value: BTreeMap) -> Result { + (value.len() <= Self::bound()).then(move || BoundedBTreeMap(value, PhantomData)).ok_or(()) + } +} + +impl codec::DecodeLength for BoundedBTreeMap { + fn len(self_encoded: &[u8]) -> Result { + // `BoundedBTreeMap` is stored just a `BTreeMap`, which is stored as a + // `Compact` with its length followed by an iteration of its items. We can just use + // the underlying implementation. + as codec::DecodeLength>::len(self_encoded) + } +} + +impl StorageDecodeLength for BoundedBTreeMap {} + +impl codec::EncodeLike> for BoundedBTreeMap where + BTreeMap: Encode +{ +} + +#[cfg(test)] +pub mod test { + use super::*; + use sp_io::TestExternalities; + use sp_std::convert::TryInto; + use crate::Twox128; + + crate::parameter_types! { + pub const Seven: u32 = 7; + pub const Four: u32 = 4; + } + + crate::generate_storage_alias! { Prefix, Foo => Value> } + crate::generate_storage_alias! { Prefix, FooMap => Map<(u32, Twox128), BoundedBTreeMap> } + crate::generate_storage_alias! { + Prefix, + FooDoubleMap => DoubleMap<(u32, Twox128), (u32, Twox128), BoundedBTreeMap> + } + + fn map_from_keys(keys: &[K]) -> BTreeMap + where + K: Ord + Copy, + { + keys.iter().copied().zip(std::iter::repeat(())).collect() + } + + fn boundedmap_from_keys(keys: &[K]) -> BoundedBTreeMap + where + K: Ord + Copy, + S: Get, + { + map_from_keys(keys).try_into().unwrap() + } + + #[test] + fn decode_len_works() { + TestExternalities::default().execute_with(|| { + let bounded = boundedmap_from_keys::(&[1, 2, 3]); + Foo::put(bounded); + assert_eq!(Foo::decode_len().unwrap(), 3); + }); + + TestExternalities::default().execute_with(|| { + let bounded = boundedmap_from_keys::(&[1, 2, 3]); + FooMap::insert(1, bounded); + assert_eq!(FooMap::decode_len(1).unwrap(), 3); + assert!(FooMap::decode_len(0).is_none()); + assert!(FooMap::decode_len(2).is_none()); + }); + + TestExternalities::default().execute_with(|| { + let bounded = boundedmap_from_keys::(&[1, 2, 3]); + FooDoubleMap::insert(1, 1, bounded); + assert_eq!(FooDoubleMap::decode_len(1, 1).unwrap(), 3); + assert!(FooDoubleMap::decode_len(2, 1).is_none()); + assert!(FooDoubleMap::decode_len(1, 2).is_none()); + assert!(FooDoubleMap::decode_len(2, 2).is_none()); + }); + } + + #[test] + fn try_insert_works() { + let mut bounded = boundedmap_from_keys::(&[1, 2, 3]); + bounded.try_insert(0, ()).unwrap(); + assert_eq!(*bounded, map_from_keys(&[1, 0, 2, 3])); + + assert!(bounded.try_insert(9, ()).is_err()); + assert_eq!(*bounded, map_from_keys(&[1, 0, 2, 3])); + } + + #[test] + fn deref_coercion_works() { + let bounded = boundedmap_from_keys::(&[1, 2, 3]); + // these methods come from deref-ed vec. + assert_eq!(bounded.len(), 3); + assert!(bounded.iter().next().is_some()); + assert!(!bounded.is_empty()); + } + + #[test] + fn try_mutate_works() { + let bounded = boundedmap_from_keys::(&[1, 2, 3, 4, 5, 6]); + let bounded = bounded + .try_mutate(|v| { + v.insert(7, ()); + }) + .unwrap(); + assert_eq!(bounded.len(), 7); + assert!(bounded + .try_mutate(|v| { + v.insert(8, ()); + }) + .is_none()); + } + + #[test] + fn btree_map_eq_works() { + let bounded = boundedmap_from_keys::(&[1, 2, 3, 4, 5, 6]); + assert_eq!(bounded, map_from_keys(&[1, 2, 3, 4, 5, 6])); + } +} diff --git a/frame/support/src/storage/bounded_vec.rs b/frame/support/src/storage/bounded_vec.rs index f441ba39b8843..320e8520e2312 100644 --- a/frame/support/src/storage/bounded_vec.rs +++ b/frame/support/src/storage/bounded_vec.rs @@ -58,22 +58,14 @@ impl> BoundedVec { } /// Create `Self` from `t` without any checks. - /// - /// # WARNING - /// - /// Only use when you are sure you know what you are doing. - fn unchecked_from(t: Vec) -> Self { + unsafe fn unchecked_from(t: Vec) -> Self { Self(t, Default::default()) } /// Create `Self` from `t` without any checks. Logs warnings if the bound is not being /// respected. The additional scope can be used to indicate where a potential overflow is /// happening. - /// - /// # WARNING - /// - /// Only use when you are sure you know what you are doing. - pub fn force_from(t: Vec, scope: Option<&'static str>) -> Self { + pub unsafe fn force_from(t: Vec, scope: Option<&'static str>) -> Self { if t.len() > Self::bound() { log::warn!( target: crate::LOG_TARGET, @@ -166,7 +158,8 @@ impl> TryFrom> for BoundedVec { type Error = (); fn try_from(t: Vec) -> Result { if t.len() <= Self::bound() { - Ok(Self::unchecked_from(t)) + // explicit check just above + Ok(unsafe {Self::unchecked_from(t)}) } else { Err(()) } @@ -434,11 +427,11 @@ pub mod test { // append to a non-existing assert!(FooMap::get(2).is_none()); assert_ok!(FooMap::try_append(2, 4)); - assert_eq!(FooMap::get(2).unwrap(), BoundedVec::::unchecked_from(vec![4])); + assert_eq!(FooMap::get(2).unwrap(), unsafe {BoundedVec::::unchecked_from(vec![4])}); assert_ok!(FooMap::try_append(2, 5)); assert_eq!( FooMap::get(2).unwrap(), - BoundedVec::::unchecked_from(vec![4, 5]) + unsafe {BoundedVec::::unchecked_from(vec![4, 5])} ); }); @@ -458,12 +451,12 @@ pub mod test { assert_ok!(FooDoubleMap::try_append(2, 1, 4)); assert_eq!( FooDoubleMap::get(2, 1).unwrap(), - BoundedVec::::unchecked_from(vec![4]) + unsafe {BoundedVec::::unchecked_from(vec![4])} ); assert_ok!(FooDoubleMap::try_append(2, 1, 5)); assert_eq!( FooDoubleMap::get(2, 1).unwrap(), - BoundedVec::::unchecked_from(vec![4, 5]) + unsafe {BoundedVec::::unchecked_from(vec![4, 5])} ); }); } diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs index adcf44a64620e..1eed6f0c4a7f2 100644 --- a/frame/support/src/storage/mod.rs +++ b/frame/support/src/storage/mod.rs @@ -29,6 +29,7 @@ pub use sp_runtime::TransactionOutcome; pub mod unhashed; pub mod hashed; +pub mod bounded_btree_map; pub mod bounded_vec; pub mod child; #[doc(hidden)] @@ -817,6 +818,7 @@ mod private { impl Sealed for Vec {} impl Sealed for Digest {} impl> Sealed for BoundedVec {} + impl Sealed for bounded_btree_map::BoundedBTreeMap {} } impl StorageAppend for Vec {}