Create new trait for non-dedup storage decode

substrate/frame/support/src/storage/mod.rs

@@ -167,6 +167,32 @@ pub trait StorageValue<T: FullCodec> {
 	{
 		T::decode_len(&Self::hashed_key())
 	}
+
+	/// Read the length of the storage value without decoding the entire value.
+	///
+	/// `T` is required to implement [`StorageDecodeNonDedupLength`].
+	///
+	/// If the value does not exists or it fails to decode the length, `None` is returned.
+	/// Otherwise `Some(len)` is returned.
+	///
+	/// # Warning
+	///
+	/// - The value returned is the non-deduplicated length of the underlying Vector in storage.
+	/// This means that any duplicate items are included.
+	///
+	/// - `None` does not mean that `get()` does not return a value. The default value is completely
+	/// ignored by this function.
+	///
+	/// # Example
+	#[doc = docify::embed!("src/storage/mod.rs", btree_set_decode_non_dedup_len)]
+	/// This demonstrates how `decode_non_dedup_len` will count even the duplicate values
+	/// in the storage (in this case, the number `4` is counted twice).
+	fn decode_non_dedup_len() -> Option<usize>
+	where
+		T: StorageDecodeNonDedupLength,
+	{
+		T::decode_non_dedup_len(&Self::hashed_key())
+	}
 }
 
 /// A non-continuous container type.
@@ -1400,8 +1426,7 @@ pub trait StoragePrefixedMap<Value: FullCodec> {
 /// This trait is sealed.
 pub trait StorageAppend<Item: Encode>: private::Sealed {}
 
-/// Marker trait that will be implemented for types that support to decode their length in an
-/// efficient way. It is expected that the length is at the beginning of the encoded object
+/// It is expected that the length is at the beginning of the encoded object
 /// and that the length is a `Compact<u32>`.
 ///
 /// This trait is sealed.
@@ -1421,6 +1446,24 @@ pub trait StorageDecodeLength: private::Sealed + codec::DecodeLength {
 	}
 }
 
+/// It is expected that the length is at the beginning of the encoded object
+/// and that the length is a `Compact<u32>`.
+///
+/// # Note
+/// The length returned by this trait is not deduplicated, i.e. it is the length of the
+/// underlying stored Vec.
+///
+/// This trait is sealed.
+pub trait StorageDecodeNonDedupLength: private::Sealed + codec::DecodeLength {
+	/// Decode the length of the storage value at `key`.
+	///
+	/// This function assumes that the length is at the beginning of the encoded object
+	/// and is a `Compact<u32>`.
+	///
+	/// Returns `None` if the storage value does not exist or the decoding failed.
+	fn decode_non_dedup_len(key: &[u8]) -> Option<usize>;
+}
+
 /// Provides `Sealed` trait to prevent implementing trait `StorageAppend` & `StorageDecodeLength`
 /// & `EncodeLikeTuple` outside of this crate.
 mod private {
@@ -1471,7 +1514,22 @@ impl<T: Encode> StorageAppend<T> for Vec<T> {}
 impl<T: Encode> StorageDecodeLength for Vec<T> {}
 
 impl<T: Encode> StorageAppend<T> for BTreeSet<T> {}
-impl<T: Encode> StorageDecodeLength for BTreeSet<T> {}
+impl<T: Encode> StorageDecodeNonDedupLength for BTreeSet<T> {
+	fn decode_non_dedup_len(key: &[u8]) -> Option<usize> {
+		let mut data = [0u8; 5];
+		let len = sp_io::storage::read(key, &mut data, 0)?;
+		let len = data.len().min(len as usize);
+		<Self as codec::DecodeLength>::len(&data[..len]).ok()
+	}
+}
+
+// Blanket implementation StorageDecodeNonDedupLength
+// for all types that are StorageDecodeLength.
+impl<T: StorageDecodeLength> StorageDecodeNonDedupLength for T {
+	fn decode_non_dedup_len(key: &[u8]) -> Option<usize> {
+		T::decode_len(key)
+	}
+}
 
 /// We abuse the fact that SCALE does not put any marker into the encoding, i.e. we only encode the
 /// internal vec and we can append to this vec. We have a test that ensures that if the `Digest`
@@ -2026,7 +2084,24 @@ mod test {
 			FooSet::append(6);
 			FooSet::append(7);
 
-			assert_eq!(FooSet::decode_len().unwrap(), 7);
+			assert_eq!(FooSet::decode_non_dedup_len().unwrap(), 7);
+		});
+	}
+
+	#[docify::export]
+	#[test]
+	fn btree_set_decode_non_dedup_len() {
+		#[crate::storage_alias]
+		type Store = StorageValue<Prefix, BTreeSet<u32>>;
+
+		TestExternalities::default().execute_with(|| {
+			Store::append(4);
+			Store::append(4); // duplicate value
+			Store::append(5);
+
+			let length_with_dup_items = 3;
+
+			assert_eq!(Store::decode_non_dedup_len().unwrap(), length_with_dup_items);
 		});
 	}
 }