From c6c01cc1f0eb85a008d20beb70c19d19d1fca398 Mon Sep 17 00:00:00 2001 From: Peter Goodspeed-Niklaus Date: Thu, 6 May 2021 16:19:04 +0200 Subject: [PATCH 1/2] Add `BoundedBTreeSet` Part of https://github.com/paritytech/substrate/issues/8719 --- .../support/src/storage/bounded_btree_set.rs | 408 ++++++++++++++++++ frame/support/src/storage/mod.rs | 2 + 2 files changed, 410 insertions(+) create mode 100644 frame/support/src/storage/bounded_btree_set.rs diff --git a/frame/support/src/storage/bounded_btree_set.rs b/frame/support/src/storage/bounded_btree_set.rs new file mode 100644 index 0000000000000..6d562e299a870 --- /dev/null +++ b/frame/support/src/storage/bounded_btree_set.rs @@ -0,0 +1,408 @@ +// 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 `BTreeSet`. + +use sp_std::{ + borrow::Borrow, collections::btree_set::BTreeSet, convert::TryFrom, fmt, marker::PhantomData, + ops::Deref, +}; +use crate::{ + storage::StorageDecodeLength, + traits::{Get, MaxEncodedLen}, +}; +use codec::{Encode, Decode}; + +/// A bounded set 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 [`BTreeSet`] for more details. +/// +/// Unlike a standard `BTreeSet`, there is a static, enforced upper limit to the number of items +/// in the set. All internal operations ensure this bound is respected. +#[derive(Encode, Decode)] +pub struct BoundedBTreeSet(BTreeSet, PhantomData); + +impl BoundedBTreeSet +where + S: Get, +{ + /// Get the bound of the type in `usize`. + pub fn bound() -> usize { + S::get() as usize + } +} + +impl BoundedBTreeSet +where + T: Ord, + S: Get, +{ + /// Create a new `BoundedBTreeSet`. + /// + /// Does not allocate. + pub fn new() -> Self { + BoundedBTreeSet(BTreeSet::new(), PhantomData) + } + + /// Create `Self` from a primitive `BTreeSet` without any checks. + unsafe fn unchecked_from(set: BTreeSet) -> Self { + Self(set, Default::default()) + } + + /// Create `Self` from a primitive `BTreeSet` 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_set::BTreeSet; + /// # use frame_support::{parameter_types, storage::bounded_btree_set::BoundedBTreeSet}; + /// parameter_types! { + /// pub const Size: u32 = 5; + /// } + /// let mut set = BTreeSet::new(); + /// set.insert("foo"); + /// set.insert("bar"); + /// let bounded_map = unsafe {BoundedBTreeSet::<_, Size>::force_from(set, "demo")}; + /// ``` + pub unsafe fn force_from(set: BTreeSet, scope: Scope) -> Self + where + Scope: Into>, + { + if set.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(set) + } + + /// Consume self, and return the inner `BTreeSet`. + /// + /// 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) -> BTreeSet { + 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 BTreeSet)) -> Option { + mutate(&mut self.0); + (self.0.len() <= Self::bound()).then(move || self) + } + + // Clears the set, removing all elements. + pub fn clear(&mut self) { + self.0.clear() + } + + /// Exactly the same semantics as [`BTreeSet::insert`], but returns an `Err` (and is a noop) if the + /// new length of the set exceeds `S`. + pub fn try_insert(&mut self, item: T) -> Result<(), ()> { + if self.len() < Self::bound() { + self.0.insert(item); + Ok(()) + } else { + Err(()) + } + } + + /// Remove an item from the set, returning whether it was previously in the set. + /// + /// The item may be any borrowed form of the set's item type, but the ordering on the borrowed + /// form _must_ match the ordering on the item type. + pub fn remove(&mut self, item: &Q) -> bool + where + T: Borrow, + Q: Ord + ?Sized, + { + self.0.remove(item) + } + + /// Removes and returns the value in the set, if any, that is equal to the given one. + /// + /// The value may be any borrowed form of the set's value type, but the ordering on the borrowed + /// form _must_ match the ordering on the value type. + pub fn take(&mut self, value: &Q) -> Option + where + T: Borrow + Ord, + Q: Ord + ?Sized, + { + self.0.take(value) + } +} + +impl Default for BoundedBTreeSet +where + T: Ord, + S: Get, +{ + fn default() -> Self { + Self::new() + } +} + +impl Clone for BoundedBTreeSet +where + BTreeSet: Clone, +{ + fn clone(&self) -> Self { + BoundedBTreeSet(self.0.clone(), PhantomData) + } +} + +#[cfg(feature = "std")] +impl fmt::Debug for BoundedBTreeSet +where + BTreeSet: fmt::Debug, + S: Get, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("BoundedBTreeSet").field(&self.0).field(&Self::bound()).finish() + } +} + +impl PartialEq for BoundedBTreeSet +where + BTreeSet: PartialEq, +{ + fn eq(&self, other: &Self) -> bool { + self.0 == other.0 + } +} + +impl Eq for BoundedBTreeSet where BTreeSet: Eq {} + +impl PartialEq> for BoundedBTreeSet +where + BTreeSet: PartialEq, +{ + fn eq(&self, other: &BTreeSet) -> bool { + self.0 == *other + } +} + +impl PartialOrd for BoundedBTreeSet +where + BTreeSet: PartialOrd, +{ + fn partial_cmp(&self, other: &Self) -> Option { + self.0.partial_cmp(&other.0) + } +} + +impl Ord for BoundedBTreeSet +where + BTreeSet: Ord, +{ + fn cmp(&self, other: &Self) -> sp_std::cmp::Ordering { + self.0.cmp(&other.0) + } +} + +impl IntoIterator for BoundedBTreeSet { + type Item = T; + type IntoIter = sp_std::collections::btree_set::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + +impl MaxEncodedLen for BoundedBTreeSet +where + T: MaxEncodedLen, + S: Get, +{ + fn max_encoded_len() -> usize { + Self::bound() + .saturating_mul(T::max_encoded_len()) + .saturating_add(codec::Compact(S::get()).encoded_size()) + } +} + +impl Deref for BoundedBTreeSet +where + T: Ord, +{ + type Target = BTreeSet; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl AsRef> for BoundedBTreeSet +where + T: Ord, +{ + fn as_ref(&self) -> &BTreeSet { + &self.0 + } +} + +impl From> for BTreeSet +where + T: Ord, +{ + fn from(set: BoundedBTreeSet) -> Self { + set.0 + } +} + +impl TryFrom> for BoundedBTreeSet +where + T: Ord, + S: Get, +{ + type Error = (); + + fn try_from(value: BTreeSet) -> Result { + (value.len() <= Self::bound()).then(move || BoundedBTreeSet(value, PhantomData)).ok_or(()) + } +} + +impl codec::DecodeLength for BoundedBTreeSet { + fn len(self_encoded: &[u8]) -> Result { + // `BoundedBTreeSet` is stored just a `BTreeSet`, 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 BoundedBTreeSet {} + +impl codec::EncodeLike> for BoundedBTreeSet where + BTreeSet: 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), BoundedBTreeSet> } + crate::generate_storage_alias! { + Prefix, + FooDoubleMap => DoubleMap<(u32, Twox128), (u32, Twox128), BoundedBTreeSet> + } + + fn map_from_keys(keys: &[T]) -> BTreeSet + where + T: Ord + Copy, + { + keys.iter().copied().collect() + } + + fn boundedmap_from_keys(keys: &[T]) -> BoundedBTreeSet + where + T: 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/mod.rs b/frame/support/src/storage/mod.rs index adcf44a64620e..508edcf86f78b 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_set; 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_set::BoundedBTreeSet {} } impl StorageAppend for Vec {} From 6725a7e1f423d0c876ac7799f3ddfc83458d6c08 Mon Sep 17 00:00:00 2001 From: Peter Goodspeed-Niklaus Date: Fri, 7 May 2021 11:01:40 +0200 Subject: [PATCH 2/2] fix copy-pasta errors Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com> --- frame/support/src/storage/bounded_btree_set.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/frame/support/src/storage/bounded_btree_set.rs b/frame/support/src/storage/bounded_btree_set.rs index 6d562e299a870..586ecca4c85e4 100644 --- a/frame/support/src/storage/bounded_btree_set.rs +++ b/frame/support/src/storage/bounded_btree_set.rs @@ -80,7 +80,7 @@ where /// let mut set = BTreeSet::new(); /// set.insert("foo"); /// set.insert("bar"); - /// let bounded_map = unsafe {BoundedBTreeSet::<_, Size>::force_from(set, "demo")}; + /// let bounded_set = unsafe {BoundedBTreeSet::<_, Size>::force_from(set, "demo")}; /// ``` pub unsafe fn force_from(set: BTreeSet, scope: Scope) -> Self where @@ -89,7 +89,7 @@ where if set.len() > Self::bound() { log::warn!( target: crate::LOG_TARGET, - "length of a bounded btreemap in scope {} is not respected.", + "length of a bounded btreeset in scope {} is not respected.", scope.into().unwrap_or("UNKNOWN"), ); } @@ -302,8 +302,7 @@ impl StorageDecodeLength for BoundedBTreeSet {} impl codec::EncodeLike> for BoundedBTreeSet where BTreeSet: Encode -{ -} +{} #[cfg(test)] pub mod test {