Add BoundedBTreeSet (#8750)

* Add `BoundedBTreeSet` Part of https://github.com/paritytech/substrate/issues/8719 * fix copy-pasta errors Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com> Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>
2026-04-27 01:07:57 +00:00 · 2021-05-10 09:49:45 +02:00
parent efe2267ec7
commit 655ebc95fb
2 changed files with 409 additions and 0 deletions
@@ -0,0 +1,407 @@
+// 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<T, S>(BTreeSet<T>, PhantomData<S>);
+
+impl<T, S> BoundedBTreeSet<T, S>
+where
+	S: Get<u32>,
+{
+	/// Get the bound of the type in `usize`.
+	pub fn bound() -> usize {
+		S::get() as usize
+	}
+}
+
+impl<T, S> BoundedBTreeSet<T, S>
+where
+	T: Ord,
+	S: Get<u32>,
+{
+	/// 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<T>) -> 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_set = unsafe {BoundedBTreeSet::<_, Size>::force_from(set, "demo")};
+	/// ```
+	pub unsafe fn force_from<Scope>(set: BTreeSet<T>, scope: Scope) -> Self
+	where
+		Scope: Into<Option<&'static str>>,
+	{
+		if set.len() > Self::bound() {
+			log::warn!(
+				target: crate::LOG_TARGET,
+				"length of a bounded btreeset 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<T> {
+		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<T>)) -> Option<Self> {
+		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<Q>(&mut self, item: &Q) -> bool
+	where
+		T: Borrow<Q>,
+		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<Q>(&mut self, value: &Q) -> Option<T>
+	where
+		T: Borrow<Q> + Ord,
+		Q: Ord + ?Sized,
+	{
+		self.0.take(value)
+	}
+}
+
+impl<T, S> Default for BoundedBTreeSet<T, S>
+where
+	T: Ord,
+	S: Get<u32>,
+{
+	fn default() -> Self {
+		Self::new()
+	}
+}
+
+impl<T, S> Clone for BoundedBTreeSet<T, S>
+where
+	BTreeSet<T>: Clone,
+{
+	fn clone(&self) -> Self {
+		BoundedBTreeSet(self.0.clone(), PhantomData)
+	}
+}
+
+#[cfg(feature = "std")]
+impl<T, S> fmt::Debug for BoundedBTreeSet<T, S>
+where
+	BTreeSet<T>: fmt::Debug,
+	S: Get<u32>,
+{
+	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+		f.debug_tuple("BoundedBTreeSet").field(&self.0).field(&Self::bound()).finish()
+	}
+}
+
+impl<T, S> PartialEq for BoundedBTreeSet<T, S>
+where
+	BTreeSet<T>: PartialEq,
+{
+	fn eq(&self, other: &Self) -> bool {
+		self.0 == other.0
+	}
+}
+
+impl<T, S> Eq for BoundedBTreeSet<T, S> where BTreeSet<T>: Eq {}
+
+impl<T, S> PartialEq<BTreeSet<T>> for BoundedBTreeSet<T, S>
+where
+	BTreeSet<T>: PartialEq,
+{
+	fn eq(&self, other: &BTreeSet<T>) -> bool {
+		self.0 == *other
+	}
+}
+
+impl<T, S> PartialOrd for BoundedBTreeSet<T, S>
+where
+	BTreeSet<T>: PartialOrd,
+{
+	fn partial_cmp(&self, other: &Self) -> Option<sp_std::cmp::Ordering> {
+		self.0.partial_cmp(&other.0)
+	}
+}
+
+impl<T, S> Ord for BoundedBTreeSet<T, S>
+where
+	BTreeSet<T>: Ord,
+{
+	fn cmp(&self, other: &Self) -> sp_std::cmp::Ordering {
+		self.0.cmp(&other.0)
+	}
+}
+
+impl<T, S> IntoIterator for BoundedBTreeSet<T, S> {
+	type Item = T;
+	type IntoIter = sp_std::collections::btree_set::IntoIter<T>;
+
+	fn into_iter(self) -> Self::IntoIter {
+		self.0.into_iter()
+	}
+}
+
+impl<T, S> MaxEncodedLen for BoundedBTreeSet<T, S>
+where
+	T: MaxEncodedLen,
+	S: Get<u32>,
+{
+	fn max_encoded_len() -> usize {
+		Self::bound()
+			.saturating_mul(T::max_encoded_len())
+			.saturating_add(codec::Compact(S::get()).encoded_size())
+	}
+}
+
+impl<T, S> Deref for BoundedBTreeSet<T, S>
+where
+	T: Ord,
+{
+	type Target = BTreeSet<T>;
+
+	fn deref(&self) -> &Self::Target {
+		&self.0
+	}
+}
+
+impl<T, S> AsRef<BTreeSet<T>> for BoundedBTreeSet<T, S>
+where
+	T: Ord,
+{
+	fn as_ref(&self) -> &BTreeSet<T> {
+		&self.0
+	}
+}
+
+impl<T, S> From<BoundedBTreeSet<T, S>> for BTreeSet<T>
+where
+	T: Ord,
+{
+	fn from(set: BoundedBTreeSet<T, S>) -> Self {
+		set.0
+	}
+}
+
+impl<T, S> TryFrom<BTreeSet<T>> for BoundedBTreeSet<T, S>
+where
+	T: Ord,
+	S: Get<u32>,
+{
+	type Error = ();
+
+	fn try_from(value: BTreeSet<T>) -> Result<Self, Self::Error> {
+		(value.len() <= Self::bound()).then(move || BoundedBTreeSet(value, PhantomData)).ok_or(())
+	}
+}
+
+impl<T, S> codec::DecodeLength for BoundedBTreeSet<T, S> {
+	fn len(self_encoded: &[u8]) -> Result<usize, codec::Error> {
+		// `BoundedBTreeSet<T, S>` is stored just a `BTreeSet<T>`, which is stored as a
+		// `Compact<u32>` with its length followed by an iteration of its items. We can just use
+		// the underlying implementation.
+		<BTreeSet<T> as codec::DecodeLength>::len(self_encoded)
+	}
+}
+
+impl<T, S> StorageDecodeLength for BoundedBTreeSet<T, S> {}
+
+impl<T, S> codec::EncodeLike<BTreeSet<T>> for BoundedBTreeSet<T, S> where
+	BTreeSet<T>: 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<BoundedBTreeSet<u32, Seven>> }
+	crate::generate_storage_alias! { Prefix, FooMap => Map<(u32, Twox128), BoundedBTreeSet<u32, Seven>> }
+	crate::generate_storage_alias! {
+		Prefix,
+		FooDoubleMap => DoubleMap<(u32, Twox128), (u32, Twox128), BoundedBTreeSet<u32, Seven>>
+	}
+
+	fn map_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>
+	where
+		T: Ord + Copy,
+		S: Get<u32>,
+	{
+		map_from_keys(keys).try_into().unwrap()
+	}
+
+	#[test]
+	fn decode_len_works() {
+		TestExternalities::default().execute_with(|| {
+			let bounded = boundedmap_from_keys::<u32, Seven>(&[1, 2, 3]);
+			Foo::put(bounded);
+			assert_eq!(Foo::decode_len().unwrap(), 3);
+		});
+
+		TestExternalities::default().execute_with(|| {
+			let bounded = boundedmap_from_keys::<u32, Seven>(&[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::<u32, Seven>(&[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::<u32, Four>(&[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::<u32, Seven>(&[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::<u32, Seven>(&[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::<u32, Seven>(&[1, 2, 3, 4, 5, 6]);
+		assert_eq!(bounded, map_from_keys(&[1, 2, 3, 4, 5, 6]));
+	}
+}
@@ -27,6 +27,7 @@ pub use sp_runtime::TransactionOutcome;
 pub mod unhashed;
 pub mod hashed;
 pub mod bounded_btree_map;
+pub mod bounded_btree_set;
 pub mod bounded_vec;
 pub mod child;
 #[doc(hidden)]
@@ -816,6 +817,7 @@ mod private {
 	impl<Hash: Encode> Sealed for Digest<Hash> {}
 	impl<T, S> Sealed for BoundedVec<T, S> {}
 	impl<K, V, S> Sealed for bounded_btree_map::BoundedBTreeMap<K, V, S> {}
+	impl<T, S> Sealed for bounded_btree_set::BoundedBTreeSet<T, S> {}
 }

 impl<T: Encode> StorageAppend<T> for Vec<T> {}