Implement DecodeWithMemTracking for BoundedBTreeSet

bounded-collections/src/bounded_btree_set.rs

@@ -359,22 +359,65 @@ where
 macro_rules! codec_impl {
 	($codec:ident) => {
 		use super::*;
-		use $codec::{Compact, Decode, DecodeLength, Encode, EncodeLike, Error, Input, MaxEncodedLen};
+		use $codec::{
+			Compact, Decode, DecodeLength, DecodeWithMemTracking, Encode, EncodeLike, Error, Input, MaxEncodedLen,
+		};
+
+		// Struct which allows prepending the compact after reading from an input.
+		pub(crate) struct PrependCompactInput<'a, I> {
+			pub encoded_len: &'a [u8],
+			pub read: usize,
+			pub inner: &'a mut I,
+		}
+
+		impl<'a, I: Input> Input for PrependCompactInput<'a, I> {
+			fn remaining_len(&mut self) -> Result<Option<usize>, Error> {
+				let remaining_compact = self.encoded_len.len().saturating_sub(self.read);
+				Ok(self.inner.remaining_len()?.map(|len| len.saturating_add(remaining_compact)))
+			}
+
+			fn read(&mut self, into: &mut [u8]) -> Result<(), Error> {
+				if into.is_empty() {
+					return Ok(());
+				}
+
+				let remaining_compact = self.encoded_len.len().saturating_sub(self.read);
+				if remaining_compact > 0 {
+					let to_read = into.len().min(remaining_compact);
+					into[..to_read].copy_from_slice(&self.encoded_len[self.read..][..to_read]);
+					self.read += to_read;
+
+					if to_read < into.len() {
+						// Buffer not full, keep reading the inner.
+						self.inner.read(&mut into[to_read..])
+					} else {
+						// Buffer was filled by the compact.
+						Ok(())
+					}
+				} else {
+					// Prepended compact has been read, just read from inner.
+					self.inner.read(into)
+				}
+			}
+		}
+
 		impl<T, S> Decode for BoundedBTreeSet<T, S>
 		where
 			T: Decode + Ord,
 			S: Get<u32>,
 		{
 			fn decode<I: Input>(input: &mut I) -> Result<Self, Error> {
-				// Same as the underlying implementation for `Decode` on `BTreeSet`, except we fail early if
-				// the len is too big.
-				let len: u32 = <Compact<u32>>::decode(input)?.into();
-				if len > S::get() {
+				// Fail early if the len is too big. This is a compact u32 which we will later put back.
+				let compact = <Compact<u32>>::decode(input)?;
+				if compact.0 > S::get() {
 					return Err("BoundedBTreeSet exceeds its limit".into());
 				}
-				input.descend_ref()?;
-				let inner = Result::from_iter((0..len).map(|_| Decode::decode(input)))?;
-				input.ascend_ref();
+				// Reconstruct the original input by prepending the length we just read, then delegate the decoding to BTreeMap.
+				let inner = BTreeSet::decode(&mut PrependCompactInput {
+					encoded_len: compact.encode().as_ref(),
+					read: 0,
+					inner: input,
+				})?;
 				Ok(Self(inner, PhantomData))
 			}
 
@@ -405,6 +448,13 @@ macro_rules! codec_impl {
 		}
 
 		impl<T, S> EncodeLike<BTreeSet<T>> for BoundedBTreeSet<T, S> where BTreeSet<T>: Encode {}
+
+		impl<T, S> DecodeWithMemTracking for BoundedBTreeSet<T, S>
+		where
+			T: Decode + Ord,
+			S: Get<u32>,
+		{
+		}
 	};
 }