diff --git a/substrate/frame/support/src/storage/bounded_btree_map.rs b/substrate/frame/support/src/storage/bounded_btree_map.rs index 201015e3a6..ed132adac6 100644 --- a/substrate/frame/support/src/storage/bounded_btree_map.rs +++ b/substrate/frame/support/src/storage/bounded_btree_map.rs @@ -17,7 +17,10 @@ //! Traits, types and structs to support a bounded BTreeMap. -use crate::{storage::StorageDecodeLength, traits::Get}; +use crate::{ + storage::StorageDecodeLength, + traits::{Get, TryCollect}, +}; use codec::{Decode, Encode, MaxEncodedLen}; use sp_std::{ borrow::Borrow, collections::btree_map::BTreeMap, convert::TryFrom, marker::PhantomData, @@ -69,6 +72,11 @@ where K: Ord, S: Get, { + /// Create `Self` from `t` without any checks. + fn unchecked_from(t: BTreeMap) -> Self { + Self(t, Default::default()) + } + /// Create a new `BoundedBTreeMap`. /// /// Does not allocate. @@ -183,16 +191,23 @@ where } } -impl PartialEq for BoundedBTreeMap +impl PartialEq> for BoundedBTreeMap where BTreeMap: PartialEq, + S1: Get, + S2: Get, { - fn eq(&self, other: &Self) -> bool { - self.0 == other.0 + fn eq(&self, other: &BoundedBTreeMap) -> bool { + S1::get() == S2::get() && self.0 == other.0 } } -impl Eq for BoundedBTreeMap where BTreeMap: Eq {} +impl Eq for BoundedBTreeMap +where + BTreeMap: Eq, + S: Get, +{ +} impl PartialEq> for BoundedBTreeMap where @@ -206,6 +221,7 @@ where impl PartialOrd for BoundedBTreeMap where BTreeMap: PartialOrd, + S: Get, { fn partial_cmp(&self, other: &Self) -> Option { self.0.partial_cmp(&other.0) @@ -215,6 +231,7 @@ where impl Ord for BoundedBTreeMap where BTreeMap: Ord, + S: Get, { fn cmp(&self, other: &Self) -> sp_std::cmp::Ordering { self.0.cmp(&other.0) @@ -302,6 +319,23 @@ impl codec::EncodeLike> for BoundedBTreeMap whe { } +impl TryCollect> for I +where + K: Ord, + I: ExactSizeIterator + Iterator, + Bound: Get, +{ + type Error = &'static str; + + fn try_collect(self) -> Result, Self::Error> { + if self.len() > Bound::get() as usize { + Err("iterator length too big") + } else { + Ok(BoundedBTreeMap::::unchecked_from(self.collect::>())) + } + } +} + #[cfg(test)] pub mod test { use super::*; @@ -452,4 +486,53 @@ pub mod test { assert_eq!(zero_key.1, false); assert_eq!(*zero_value, 6); } + + #[test] + fn can_be_collected() { + let b1 = boundedmap_from_keys::>(&[1, 2, 3, 4]); + let b2: BoundedBTreeMap> = + b1.iter().map(|(k, v)| (k + 1, *v)).try_collect().unwrap(); + assert_eq!(b2.into_iter().map(|(k, _)| k).collect::>(), vec![2, 3, 4, 5]); + + // can also be collected into a collection of length 4. + let b2: BoundedBTreeMap> = + b1.iter().map(|(k, v)| (k + 1, *v)).try_collect().unwrap(); + assert_eq!(b2.into_iter().map(|(k, _)| k).collect::>(), vec![2, 3, 4, 5]); + + // can be mutated further into iterators that are `ExactSizedIterator`. + let b2: BoundedBTreeMap> = + b1.iter().map(|(k, v)| (k + 1, *v)).rev().skip(2).try_collect().unwrap(); + // note that the binary tree will re-sort this, so rev() is not really seen + assert_eq!(b2.into_iter().map(|(k, _)| k).collect::>(), vec![2, 3]); + + let b2: BoundedBTreeMap> = + b1.iter().map(|(k, v)| (k + 1, *v)).take(2).try_collect().unwrap(); + assert_eq!(b2.into_iter().map(|(k, _)| k).collect::>(), vec![2, 3]); + + // but these worn't work + let b2: Result>, _> = + b1.iter().map(|(k, v)| (k + 1, *v)).try_collect(); + assert!(b2.is_err()); + + let b2: Result>, _> = + b1.iter().map(|(k, v)| (k + 1, *v)).skip(2).try_collect(); + assert!(b2.is_err()); + } + + #[test] + fn eq_works() { + // of same type + let b1 = boundedmap_from_keys::>(&[1, 2]); + let b2 = boundedmap_from_keys::>(&[1, 2]); + assert_eq!(b1, b2); + + // of different type, but same value and bound. + crate::parameter_types! { + B1: u32 = 7; + B2: u32 = 7; + } + let b1 = boundedmap_from_keys::(&[1, 2]); + let b2 = boundedmap_from_keys::(&[1, 2]); + assert_eq!(b1, b2); + } } diff --git a/substrate/frame/support/src/storage/bounded_btree_set.rs b/substrate/frame/support/src/storage/bounded_btree_set.rs index f51439f04d..7d543549c6 100644 --- a/substrate/frame/support/src/storage/bounded_btree_set.rs +++ b/substrate/frame/support/src/storage/bounded_btree_set.rs @@ -17,7 +17,10 @@ //! Traits, types and structs to support a bounded `BTreeSet`. -use crate::{storage::StorageDecodeLength, traits::Get}; +use crate::{ + storage::StorageDecodeLength, + traits::{Get, TryCollect}, +}; use codec::{Decode, Encode, MaxEncodedLen}; use sp_std::{ borrow::Borrow, collections::btree_set::BTreeSet, convert::TryFrom, marker::PhantomData, @@ -68,6 +71,11 @@ where T: Ord, S: Get, { + /// Create `Self` from `t` without any checks. + fn unchecked_from(t: BTreeSet) -> Self { + Self(t, Default::default()) + } + /// Create a new `BoundedBTreeSet`. /// /// Does not allocate. @@ -168,20 +176,28 @@ where } } -impl PartialEq for BoundedBTreeSet +impl PartialEq> for BoundedBTreeSet where BTreeSet: PartialEq, + S1: Get, + S2: Get, { - fn eq(&self, other: &Self) -> bool { - self.0 == other.0 + fn eq(&self, other: &BoundedBTreeSet) -> bool { + S1::get() == S2::get() && self.0 == other.0 } } -impl Eq for BoundedBTreeSet where BTreeSet: Eq {} +impl Eq for BoundedBTreeSet +where + BTreeSet: Eq, + S: Get, +{ +} impl PartialEq> for BoundedBTreeSet where BTreeSet: PartialEq, + S: Get, { fn eq(&self, other: &BTreeSet) -> bool { self.0 == *other @@ -191,6 +207,7 @@ where impl PartialOrd for BoundedBTreeSet where BTreeSet: PartialOrd, + S: Get, { fn partial_cmp(&self, other: &Self) -> Option { self.0.partial_cmp(&other.0) @@ -200,6 +217,7 @@ where impl Ord for BoundedBTreeSet where BTreeSet: Ord, + S: Get, { fn cmp(&self, other: &Self) -> sp_std::cmp::Ordering { self.0.cmp(&other.0) @@ -283,6 +301,23 @@ impl StorageDecodeLength for BoundedBTreeSet {} impl codec::EncodeLike> for BoundedBTreeSet where BTreeSet: Encode {} +impl TryCollect> for I +where + T: Ord, + I: ExactSizeIterator + Iterator, + Bound: Get, +{ + type Error = &'static str; + + fn try_collect(self) -> Result, Self::Error> { + if self.len() > Bound::get() as usize { + Err("iterator length too big") + } else { + Ok(BoundedBTreeSet::::unchecked_from(self.collect::>())) + } + } +} + #[cfg(test)] pub mod test { use super::*; @@ -298,31 +333,31 @@ pub mod test { FooDoubleMap => DoubleMap<(u32, Twox128), (u32, Twox128), BoundedBTreeSet>> } - fn map_from_keys(keys: &[T]) -> BTreeSet + fn set_from_keys(keys: &[T]) -> BTreeSet where T: Ord + Copy, { keys.iter().copied().collect() } - fn boundedmap_from_keys(keys: &[T]) -> BoundedBTreeSet + fn boundedset_from_keys(keys: &[T]) -> BoundedBTreeSet where T: Ord + Copy, S: Get, { - map_from_keys(keys).try_into().unwrap() + set_from_keys(keys).try_into().unwrap() } #[test] fn decode_len_works() { TestExternalities::default().execute_with(|| { - let bounded = boundedmap_from_keys::>(&[1, 2, 3]); + let bounded = boundedset_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]); + let bounded = boundedset_from_keys::>(&[1, 2, 3]); FooMap::insert(1, bounded); assert_eq!(FooMap::decode_len(1).unwrap(), 3); assert!(FooMap::decode_len(0).is_none()); @@ -330,7 +365,7 @@ pub mod test { }); TestExternalities::default().execute_with(|| { - let bounded = boundedmap_from_keys::>(&[1, 2, 3]); + let bounded = boundedset_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()); @@ -341,17 +376,17 @@ pub mod test { #[test] fn try_insert_works() { - let mut bounded = boundedmap_from_keys::>(&[1, 2, 3]); + let mut bounded = boundedset_from_keys::>(&[1, 2, 3]); bounded.try_insert(0).unwrap(); - assert_eq!(*bounded, map_from_keys(&[1, 0, 2, 3])); + assert_eq!(*bounded, set_from_keys(&[1, 0, 2, 3])); assert!(bounded.try_insert(9).is_err()); - assert_eq!(*bounded, map_from_keys(&[1, 0, 2, 3])); + assert_eq!(*bounded, set_from_keys(&[1, 0, 2, 3])); } #[test] fn deref_coercion_works() { - let bounded = boundedmap_from_keys::>(&[1, 2, 3]); + let bounded = boundedset_from_keys::>(&[1, 2, 3]); // these methods come from deref-ed vec. assert_eq!(bounded.len(), 3); assert!(bounded.iter().next().is_some()); @@ -360,7 +395,7 @@ pub mod test { #[test] fn try_mutate_works() { - let bounded = boundedmap_from_keys::>(&[1, 2, 3, 4, 5, 6]); + let bounded = boundedset_from_keys::>(&[1, 2, 3, 4, 5, 6]); let bounded = bounded .try_mutate(|v| { v.insert(7); @@ -376,8 +411,8 @@ pub mod test { #[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])); + let bounded = boundedset_from_keys::>(&[1, 2, 3, 4, 5, 6]); + assert_eq!(bounded, set_from_keys(&[1, 2, 3, 4, 5, 6])); } #[test] @@ -433,4 +468,51 @@ pub mod test { assert_eq!(zero_item.0, 0); assert_eq!(zero_item.1, false); } + + #[test] + fn can_be_collected() { + let b1 = boundedset_from_keys::>(&[1, 2, 3, 4]); + let b2: BoundedBTreeSet> = b1.iter().map(|k| k + 1).try_collect().unwrap(); + assert_eq!(b2.into_iter().collect::>(), vec![2, 3, 4, 5]); + + // can also be collected into a collection of length 4. + let b2: BoundedBTreeSet> = b1.iter().map(|k| k + 1).try_collect().unwrap(); + assert_eq!(b2.into_iter().collect::>(), vec![2, 3, 4, 5]); + + // can be mutated further into iterators that are `ExactSizedIterator`. + let b2: BoundedBTreeSet> = + b1.iter().map(|k| k + 1).rev().skip(2).try_collect().unwrap(); + // note that the binary tree will re-sort this, so rev() is not really seen + assert_eq!(b2.into_iter().collect::>(), vec![2, 3]); + + let b2: BoundedBTreeSet> = + b1.iter().map(|k| k + 1).take(2).try_collect().unwrap(); + assert_eq!(b2.into_iter().collect::>(), vec![2, 3]); + + // but these worn't work + let b2: Result>, _> = + b1.iter().map(|k| k + 1).try_collect(); + assert!(b2.is_err()); + + let b2: Result>, _> = + b1.iter().map(|k| k + 1).skip(2).try_collect(); + assert!(b2.is_err()); + } + + #[test] + fn eq_works() { + // of same type + let b1 = boundedset_from_keys::>(&[1, 2]); + let b2 = boundedset_from_keys::>(&[1, 2]); + assert_eq!(b1, b2); + + // of different type, but same value and bound. + crate::parameter_types! { + B1: u32 = 7; + B2: u32 = 7; + } + let b1 = boundedset_from_keys::(&[1, 2]); + let b2 = boundedset_from_keys::(&[1, 2]); + assert_eq!(b1, b2); + } } diff --git a/substrate/frame/support/src/storage/bounded_vec.rs b/substrate/frame/support/src/storage/bounded_vec.rs index 5ea7a62c79..bdac8f23d7 100644 --- a/substrate/frame/support/src/storage/bounded_vec.rs +++ b/substrate/frame/support/src/storage/bounded_vec.rs @@ -20,7 +20,7 @@ use crate::{ storage::{StorageDecodeLength, StorageTryAppend}, - traits::Get, + traits::{Get, TryCollect}, WeakBoundedVec, }; use codec::{Decode, Encode, EncodeLike, MaxEncodedLen}; @@ -146,11 +146,34 @@ impl> From> for Vec { } impl> BoundedVec { + /// Pre-allocate `capacity` items in self. + /// + /// If `capacity` is greater than [`Self::bound`], then the minimum of the two is used. + pub fn with_bounded_capacity(capacity: usize) -> Self { + let capacity = capacity.min(Self::bound()); + Self(Vec::with_capacity(capacity), Default::default()) + } + + /// Allocate self with the maximum possible capacity. + pub fn with_max_capacity() -> Self { + Self::with_bounded_capacity(Self::bound()) + } + /// Get the bound of the type in `usize`. pub fn bound() -> usize { S::get() as usize } + /// Same as `Vec::resize`, but if `size` is more than [`Self::bound`], then [`Self::bound`] is + /// used. + pub fn bounded_resize(&mut self, size: usize, value: T) + where + T: Clone, + { + let size = size.min(Self::bound()); + self.0.resize(size, value); + } + /// 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 @@ -300,15 +323,14 @@ impl codec::DecodeLength for BoundedVec { } } -// NOTE: we could also implement this as: -// impl, S2: Get> PartialEq> for BoundedVec -// to allow comparison of bounded vectors with different bounds. -impl PartialEq for BoundedVec +impl PartialEq> for BoundedVec where T: PartialEq, + BoundSelf: Get, + BoundRhs: Get, { - fn eq(&self, rhs: &Self) -> bool { - self.0 == rhs.0 + fn eq(&self, rhs: &BoundedVec) -> bool { + BoundSelf::get() == BoundRhs::get() && self.0 == rhs.0 } } @@ -318,7 +340,7 @@ impl> PartialEq> for BoundedVec { } } -impl Eq for BoundedVec where T: Eq {} +impl> Eq for BoundedVec where T: Eq {} impl StorageDecodeLength for BoundedVec {} @@ -344,6 +366,22 @@ where } } +impl TryCollect> for I +where + I: ExactSizeIterator + Iterator, + Bound: Get, +{ + type Error = &'static str; + + fn try_collect(self) -> Result, Self::Error> { + if self.len() > Bound::get() as usize { + Err("iterator length too big") + } else { + Ok(BoundedVec::::unchecked_from(self.collect::>())) + } + } +} + #[cfg(test)] pub mod test { use super::*; @@ -452,4 +490,59 @@ pub mod test { Err("BoundedVec exceeds its limit".into()), ); } + + #[test] + fn can_be_collected() { + let b1: BoundedVec> = vec![1, 2, 3, 4].try_into().unwrap(); + let b2: BoundedVec> = b1.iter().map(|x| x + 1).try_collect().unwrap(); + assert_eq!(b2, vec![2, 3, 4, 5]); + + // can also be collected into a collection of length 4. + let b2: BoundedVec> = b1.iter().map(|x| x + 1).try_collect().unwrap(); + assert_eq!(b2, vec![2, 3, 4, 5]); + + // can be mutated further into iterators that are `ExactSizedIterator`. + let b2: BoundedVec> = + b1.iter().map(|x| x + 1).rev().try_collect().unwrap(); + assert_eq!(b2, vec![5, 4, 3, 2]); + + let b2: BoundedVec> = + b1.iter().map(|x| x + 1).rev().skip(2).try_collect().unwrap(); + assert_eq!(b2, vec![3, 2]); + let b2: BoundedVec> = + b1.iter().map(|x| x + 1).rev().skip(2).try_collect().unwrap(); + assert_eq!(b2, vec![3, 2]); + + let b2: BoundedVec> = + b1.iter().map(|x| x + 1).rev().take(2).try_collect().unwrap(); + assert_eq!(b2, vec![5, 4]); + let b2: BoundedVec> = + b1.iter().map(|x| x + 1).rev().take(2).try_collect().unwrap(); + assert_eq!(b2, vec![5, 4]); + + // but these worn't work + let b2: Result>, _> = b1.iter().map(|x| x + 1).try_collect(); + assert!(b2.is_err()); + + let b2: Result>, _> = + b1.iter().map(|x| x + 1).rev().take(2).try_collect(); + assert!(b2.is_err()); + } + + #[test] + fn eq_works() { + // of same type + let b1: BoundedVec> = vec![1, 2, 3].try_into().unwrap(); + let b2: BoundedVec> = vec![1, 2, 3].try_into().unwrap(); + assert_eq!(b1, b2); + + // of different type, but same value and bound. + crate::parameter_types! { + B1: u32 = 7; + B2: u32 = 7; + } + let b1: BoundedVec = vec![1, 2, 3].try_into().unwrap(); + let b2: BoundedVec = vec![1, 2, 3].try_into().unwrap(); + assert_eq!(b1, b2); + } } diff --git a/substrate/frame/support/src/traits.rs b/substrate/frame/support/src/traits.rs index d0719063ef..85a0698759 100644 --- a/substrate/frame/support/src/traits.rs +++ b/substrate/frame/support/src/traits.rs @@ -54,7 +54,8 @@ pub use misc::{ ConstU32, ConstU64, ConstU8, EnsureInherentsAreFirst, EqualPrivilegeOnly, EstimateCallFee, ExecuteBlock, ExtrinsicCall, Get, GetBacking, GetDefault, HandleLifetime, IsSubType, IsType, Len, OffchainWorker, OnKilledAccount, OnNewAccount, PreimageProvider, PreimageRecipient, - PrivilegeCmp, SameOrOther, Time, TryDrop, UnixTime, WrapperKeepOpaque, WrapperOpaque, + PrivilegeCmp, SameOrOther, Time, TryCollect, TryDrop, UnixTime, WrapperKeepOpaque, + WrapperOpaque, }; mod stored_map; diff --git a/substrate/frame/support/src/traits/misc.rs b/substrate/frame/support/src/traits/misc.rs index c68561367e..df21624c3c 100644 --- a/substrate/frame/support/src/traits/misc.rs +++ b/substrate/frame/support/src/traits/misc.rs @@ -23,6 +23,16 @@ use scale_info::{build::Fields, meta_type, Path, Type, TypeInfo, TypeParameter}; use sp_runtime::{traits::Block as BlockT, DispatchError}; use sp_std::{cmp::Ordering, prelude::*}; +/// Try and collect into a collection `C`. +pub trait TryCollect { + type Error; + /// Consume self and try to collect the results into `C`. + /// + /// This is useful in preventing the undesirable `.collect().try_into()` call chain on + /// collections that need to be converted into a bounded type (e.g. `BoundedVec`). + fn try_collect(self) -> Result; +} + /// Anything that can have a `::len()` method. pub trait Len { /// Return the length of data type.