some improvements to bounded vec API (#10590)

* some improvements to bounded vec

* revert license tweak

* more tests

* fix

* Update frame/support/src/storage/bounded_vec.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* add the same stuff for btree map and set as well

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

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<u32>,
 {
+	/// Create `Self` from `t` without any checks.
+	fn unchecked_from(t: BTreeMap<K, V>) -> Self {
+		Self(t, Default::default())
+	}
+
 	/// Create a new `BoundedBTreeMap`.
 	///
 	/// Does not allocate.
@@ -183,16 +191,23 @@ where
 	}
 }
 
-impl<K, V, S> PartialEq for BoundedBTreeMap<K, V, S>
+impl<K, V, S1, S2> PartialEq<BoundedBTreeMap<K, V, S1>> for BoundedBTreeMap<K, V, S2>
 where
 	BTreeMap<K, V>: PartialEq,
+	S1: Get<u32>,
+	S2: Get<u32>,
 {
-	fn eq(&self, other: &Self) -> bool {
-		self.0 == other.0
+	fn eq(&self, other: &BoundedBTreeMap<K, V, S1>) -> bool {
+		S1::get() == S2::get() && self.0 == other.0
 	}
 }
 
-impl<K, V, S> Eq for BoundedBTreeMap<K, V, S> where BTreeMap<K, V>: Eq {}
+impl<K, V, S> Eq for BoundedBTreeMap<K, V, S>
+where
+	BTreeMap<K, V>: Eq,
+	S: Get<u32>,
+{
+}
 
 impl<K, V, S> PartialEq<BTreeMap<K, V>> for BoundedBTreeMap<K, V, S>
 where
@@ -206,6 +221,7 @@ where
 impl<K, V, S> PartialOrd for BoundedBTreeMap<K, V, S>
 where
 	BTreeMap<K, V>: PartialOrd,
+	S: Get<u32>,
 {
 	fn partial_cmp(&self, other: &Self) -> Option<sp_std::cmp::Ordering> {
 		self.0.partial_cmp(&other.0)
@@ -215,6 +231,7 @@ where
 impl<K, V, S> Ord for BoundedBTreeMap<K, V, S>
 where
 	BTreeMap<K, V>: Ord,
+	S: Get<u32>,
 {
 	fn cmp(&self, other: &Self) -> sp_std::cmp::Ordering {
 		self.0.cmp(&other.0)
@@ -302,6 +319,23 @@ impl<K, V, S> codec::EncodeLike<BTreeMap<K, V>> for BoundedBTreeMap<K, V, S> whe
 {
 }
 
+impl<I, K, V, Bound> TryCollect<BoundedBTreeMap<K, V, Bound>> for I
+where
+	K: Ord,
+	I: ExactSizeIterator + Iterator<Item = (K, V)>,
+	Bound: Get<u32>,
+{
+	type Error = &'static str;
+
+	fn try_collect(self) -> Result<BoundedBTreeMap<K, V, Bound>, Self::Error> {
+		if self.len() > Bound::get() as usize {
+			Err("iterator length too big")
+		} else {
+			Ok(BoundedBTreeMap::<K, V, Bound>::unchecked_from(self.collect::<BTreeMap<K, V>>()))
+		}
+	}
+}
+
 #[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::<u32, ConstU32<5>>(&[1, 2, 3, 4]);
+		let b2: BoundedBTreeMap<u32, (), ConstU32<5>> =
+			b1.iter().map(|(k, v)| (k + 1, *v)).try_collect().unwrap();
+		assert_eq!(b2.into_iter().map(|(k, _)| k).collect::<Vec<_>>(), vec![2, 3, 4, 5]);
+
+		// can also be collected into a collection of length 4.
+		let b2: BoundedBTreeMap<u32, (), ConstU32<4>> =
+			b1.iter().map(|(k, v)| (k + 1, *v)).try_collect().unwrap();
+		assert_eq!(b2.into_iter().map(|(k, _)| k).collect::<Vec<_>>(), vec![2, 3, 4, 5]);
+
+		// can be mutated further into iterators that are `ExactSizedIterator`.
+		let b2: BoundedBTreeMap<u32, (), ConstU32<5>> =
+			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<_>>(), vec![2, 3]);
+
+		let b2: BoundedBTreeMap<u32, (), ConstU32<5>> =
+			b1.iter().map(|(k, v)| (k + 1, *v)).take(2).try_collect().unwrap();
+		assert_eq!(b2.into_iter().map(|(k, _)| k).collect::<Vec<_>>(), vec![2, 3]);
+
+		// but these worn't work
+		let b2: Result<BoundedBTreeMap<u32, (), ConstU32<3>>, _> =
+			b1.iter().map(|(k, v)| (k + 1, *v)).try_collect();
+		assert!(b2.is_err());
+
+		let b2: Result<BoundedBTreeMap<u32, (), ConstU32<1>>, _> =
+			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::<u32, ConstU32<7>>(&[1, 2]);
+		let b2 = boundedmap_from_keys::<u32, ConstU32<7>>(&[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::<u32, B1>(&[1, 2]);
+		let b2 = boundedmap_from_keys::<u32, B2>(&[1, 2]);
+		assert_eq!(b1, b2);
+	}
 }

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<u32>,
 {
+	/// Create `Self` from `t` without any checks.
+	fn unchecked_from(t: BTreeSet<T>) -> Self {
+		Self(t, Default::default())
+	}
+
 	/// Create a new `BoundedBTreeSet`.
 	///
 	/// Does not allocate.
@@ -168,20 +176,28 @@ where
 	}
 }
 
-impl<T, S> PartialEq for BoundedBTreeSet<T, S>
+impl<T, S1, S2> PartialEq<BoundedBTreeSet<T, S1>> for BoundedBTreeSet<T, S2>
 where
 	BTreeSet<T>: PartialEq,
+	S1: Get<u32>,
+	S2: Get<u32>,
 {
-	fn eq(&self, other: &Self) -> bool {
-		self.0 == other.0
+	fn eq(&self, other: &BoundedBTreeSet<T, S1>) -> bool {
+		S1::get() == S2::get() && self.0 == other.0
 	}
 }
 
-impl<T, S> Eq for BoundedBTreeSet<T, S> where BTreeSet<T>: Eq {}
+impl<T, S> Eq for BoundedBTreeSet<T, S>
+where
+	BTreeSet<T>: Eq,
+	S: Get<u32>,
+{
+}
 
 impl<T, S> PartialEq<BTreeSet<T>> for BoundedBTreeSet<T, S>
 where
 	BTreeSet<T>: PartialEq,
+	S: Get<u32>,
 {
 	fn eq(&self, other: &BTreeSet<T>) -> bool {
 		self.0 == *other
@@ -191,6 +207,7 @@ where
 impl<T, S> PartialOrd for BoundedBTreeSet<T, S>
 where
 	BTreeSet<T>: PartialOrd,
+	S: Get<u32>,
 {
 	fn partial_cmp(&self, other: &Self) -> Option<sp_std::cmp::Ordering> {
 		self.0.partial_cmp(&other.0)
@@ -200,6 +217,7 @@ where
 impl<T, S> Ord for BoundedBTreeSet<T, S>
 where
 	BTreeSet<T>: Ord,
+	S: Get<u32>,
 {
 	fn cmp(&self, other: &Self) -> sp_std::cmp::Ordering {
 		self.0.cmp(&other.0)
@@ -283,6 +301,23 @@ impl<T, S> StorageDecodeLength for BoundedBTreeSet<T, S> {}
 
 impl<T, S> codec::EncodeLike<BTreeSet<T>> for BoundedBTreeSet<T, S> where BTreeSet<T>: Encode {}
 
+impl<I, T, Bound> TryCollect<BoundedBTreeSet<T, Bound>> for I
+where
+	T: Ord,
+	I: ExactSizeIterator + Iterator<Item = T>,
+	Bound: Get<u32>,
+{
+	type Error = &'static str;
+
+	fn try_collect(self) -> Result<BoundedBTreeSet<T, Bound>, Self::Error> {
+		if self.len() > Bound::get() as usize {
+			Err("iterator length too big")
+		} else {
+			Ok(BoundedBTreeSet::<T, Bound>::unchecked_from(self.collect::<BTreeSet<T>>()))
+		}
+	}
+}
+
 #[cfg(test)]
 pub mod test {
 	use super::*;
@@ -298,31 +333,31 @@ pub mod test {
 		FooDoubleMap => DoubleMap<(u32, Twox128), (u32, Twox128), BoundedBTreeSet<u32, ConstU32<7>>>
 	}
 
-	fn map_from_keys<T>(keys: &[T]) -> BTreeSet<T>
+	fn set_from_keys<T>(keys: &[T]) -> BTreeSet<T>
 	where
 		T: Ord + Copy,
 	{
 		keys.iter().copied().collect()
 	}
 
-	fn boundedmap_from_keys<T, S>(keys: &[T]) -> BoundedBTreeSet<T, S>
+	fn boundedset_from_keys<T, S>(keys: &[T]) -> BoundedBTreeSet<T, S>
 	where
 		T: Ord + Copy,
 		S: Get<u32>,
 	{
-		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::<u32, ConstU32<7>>(&[1, 2, 3]);
+			let bounded = boundedset_from_keys::<u32, ConstU32<7>>(&[1, 2, 3]);
 			Foo::put(bounded);
 			assert_eq!(Foo::decode_len().unwrap(), 3);
 		});
 
 		TestExternalities::default().execute_with(|| {
-			let bounded = boundedmap_from_keys::<u32, ConstU32<7>>(&[1, 2, 3]);
+			let bounded = boundedset_from_keys::<u32, ConstU32<7>>(&[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::<u32, ConstU32<7>>(&[1, 2, 3]);
+			let bounded = boundedset_from_keys::<u32, ConstU32<7>>(&[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::<u32, ConstU32<4>>(&[1, 2, 3]);
+		let mut bounded = boundedset_from_keys::<u32, ConstU32<4>>(&[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::<u32, ConstU32<7>>(&[1, 2, 3]);
+		let bounded = boundedset_from_keys::<u32, ConstU32<7>>(&[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::<u32, ConstU32<7>>(&[1, 2, 3, 4, 5, 6]);
+		let bounded = boundedset_from_keys::<u32, ConstU32<7>>(&[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::<u32, ConstU32<7>>(&[1, 2, 3, 4, 5, 6]);
-		assert_eq!(bounded, map_from_keys(&[1, 2, 3, 4, 5, 6]));
+		let bounded = boundedset_from_keys::<u32, ConstU32<7>>(&[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::<u32, ConstU32<5>>(&[1, 2, 3, 4]);
+		let b2: BoundedBTreeSet<u32, ConstU32<5>> = b1.iter().map(|k| k + 1).try_collect().unwrap();
+		assert_eq!(b2.into_iter().collect::<Vec<_>>(), vec![2, 3, 4, 5]);
+
+		// can also be collected into a collection of length 4.
+		let b2: BoundedBTreeSet<u32, ConstU32<4>> = b1.iter().map(|k| k + 1).try_collect().unwrap();
+		assert_eq!(b2.into_iter().collect::<Vec<_>>(), vec![2, 3, 4, 5]);
+
+		// can be mutated further into iterators that are `ExactSizedIterator`.
+		let b2: BoundedBTreeSet<u32, ConstU32<5>> =
+			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<_>>(), vec![2, 3]);
+
+		let b2: BoundedBTreeSet<u32, ConstU32<5>> =
+			b1.iter().map(|k| k + 1).take(2).try_collect().unwrap();
+		assert_eq!(b2.into_iter().collect::<Vec<_>>(), vec![2, 3]);
+
+		// but these worn't work
+		let b2: Result<BoundedBTreeSet<u32, ConstU32<3>>, _> =
+			b1.iter().map(|k| k + 1).try_collect();
+		assert!(b2.is_err());
+
+		let b2: Result<BoundedBTreeSet<u32, ConstU32<1>>, _> =
+			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::<u32, ConstU32<7>>(&[1, 2]);
+		let b2 = boundedset_from_keys::<u32, ConstU32<7>>(&[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::<u32, B1>(&[1, 2]);
+		let b2 = boundedset_from_keys::<u32, B2>(&[1, 2]);
+		assert_eq!(b1, b2);
+	}
 }

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<T, S: Get<u32>> From<BoundedVec<T, S>> for Vec<T> {
 }
 
 impl<T, S: Get<u32>> BoundedVec<T, S> {
+	/// 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<T, S> codec::DecodeLength for BoundedVec<T, S> {
 	}
 }
 
-// NOTE: we could also implement this as:
-// impl<T: Value, S1: Get<u32>, S2: Get<u32>> PartialEq<BoundedVec<T, S2>> for BoundedVec<T, S1>
-// to allow comparison of bounded vectors with different bounds.
-impl<T, S> PartialEq for BoundedVec<T, S>
+impl<T, BoundSelf, BoundRhs> PartialEq<BoundedVec<T, BoundRhs>> for BoundedVec<T, BoundSelf>
 where
 	T: PartialEq,
+	BoundSelf: Get<u32>,
+	BoundRhs: Get<u32>,
 {
-	fn eq(&self, rhs: &Self) -> bool {
-		self.0 == rhs.0
+	fn eq(&self, rhs: &BoundedVec<T, BoundRhs>) -> bool {
+		BoundSelf::get() == BoundRhs::get() && self.0 == rhs.0
 	}
 }
 
@@ -318,7 +340,7 @@ impl<T: PartialEq, S: Get<u32>> PartialEq<Vec<T>> for BoundedVec<T, S> {
 	}
 }
 
-impl<T, S> Eq for BoundedVec<T, S> where T: Eq {}
+impl<T, S: Get<u32>> Eq for BoundedVec<T, S> where T: Eq {}
 
 impl<T, S> StorageDecodeLength for BoundedVec<T, S> {}
 
@@ -344,6 +366,22 @@ where
 	}
 }
 
+impl<I, T, Bound> TryCollect<BoundedVec<T, Bound>> for I
+where
+	I: ExactSizeIterator + Iterator<Item = T>,
+	Bound: Get<u32>,
+{
+	type Error = &'static str;
+
+	fn try_collect(self) -> Result<BoundedVec<T, Bound>, Self::Error> {
+		if self.len() > Bound::get() as usize {
+			Err("iterator length too big")
+		} else {
+			Ok(BoundedVec::<T, Bound>::unchecked_from(self.collect::<Vec<T>>()))
+		}
+	}
+}
+
 #[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<u32, ConstU32<5>> = vec![1, 2, 3, 4].try_into().unwrap();
+		let b2: BoundedVec<u32, ConstU32<5>> = 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<u32, ConstU32<4>> = 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<u32, ConstU32<4>> =
+			b1.iter().map(|x| x + 1).rev().try_collect().unwrap();
+		assert_eq!(b2, vec![5, 4, 3, 2]);
+
+		let b2: BoundedVec<u32, ConstU32<4>> =
+			b1.iter().map(|x| x + 1).rev().skip(2).try_collect().unwrap();
+		assert_eq!(b2, vec![3, 2]);
+		let b2: BoundedVec<u32, ConstU32<2>> =
+			b1.iter().map(|x| x + 1).rev().skip(2).try_collect().unwrap();
+		assert_eq!(b2, vec![3, 2]);
+
+		let b2: BoundedVec<u32, ConstU32<4>> =
+			b1.iter().map(|x| x + 1).rev().take(2).try_collect().unwrap();
+		assert_eq!(b2, vec![5, 4]);
+		let b2: BoundedVec<u32, ConstU32<2>> =
+			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<BoundedVec<u32, ConstU32<3>>, _> = b1.iter().map(|x| x + 1).try_collect();
+		assert!(b2.is_err());
+
+		let b2: Result<BoundedVec<u32, ConstU32<1>>, _> =
+			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<u32, ConstU32<7>> = vec![1, 2, 3].try_into().unwrap();
+		let b2: BoundedVec<u32, ConstU32<7>> = 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<u32, B1> = vec![1, 2, 3].try_into().unwrap();
+		let b2: BoundedVec<u32, B2> = vec![1, 2, 3].try_into().unwrap();
+		assert_eq!(b1, b2);
+	}
 }