Create new trait for non-dedup storage decode (#1932)

- This adds the new trait `StorageDecodeNonDedupLength` and implements
them for `BTreeSet` and its bounded types.
- New unit test has been added to cover the case.  
- See linked
[issue](#126) which
outlines the original issue.

Note that the added trait here doesn't add new logic but improves
semantics.

---------

Co-authored-by: joe petrowski <[email protected]>
Co-authored-by: Kian Paimani <[email protected]>
Co-authored-by: Oliver Tale-Yazdi <[email protected]>
Co-authored-by: command-bot <>

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

@@ -79,5 +79,16 @@ pub mod test {
 			assert!(FooDoubleMap::decode_len(1, 2).is_none());
 			assert!(FooDoubleMap::decode_len(2, 2).is_none());
 		});
+
+		TestExternalities::default().execute_with(|| {
+			let bounded = boundedmap_from_keys::<u32, ConstU32<7>>(&[1, 2, 3]);
+			FooDoubleMap::insert(1, 1, bounded.clone());
+			FooDoubleMap::insert(2, 2, bounded); // duplicate value
+
+			assert_eq!(FooDoubleMap::decode_len(1, 1).unwrap(), 3);
+			assert_eq!(FooDoubleMap::decode_len(2, 2).unwrap(), 3);
+			assert!(FooDoubleMap::decode_len(2, 1).is_none());
+			assert!(FooDoubleMap::decode_len(1, 2).is_none());
+		});
 	}
 }

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

@@ -17,10 +17,10 @@
 
 //! Traits, types and structs to support a bounded `BTreeSet`.
 
-use crate::storage::StorageDecodeLength;
+use frame_support::storage::StorageDecodeNonDedupLength;
 pub use sp_runtime::BoundedBTreeSet;
 
-impl<T, S> StorageDecodeLength for BoundedBTreeSet<T, S> {}
+impl<T, S> StorageDecodeNonDedupLength for BoundedBTreeSet<T, S> {}
 
 #[cfg(test)]
 pub mod test {
@@ -56,28 +56,28 @@ pub mod test {
 	}
 
 	#[test]
-	fn decode_len_works() {
+	fn decode_non_dedup_len_works() {
 		TestExternalities::default().execute_with(|| {
 			let bounded = boundedset_from_keys::<u32, ConstU32<7>>(&[1, 2, 3]);
 			Foo::put(bounded);
-			assert_eq!(Foo::decode_len().unwrap(), 3);
+			assert_eq!(Foo::decode_non_dedup_len().unwrap(), 3);
 		});
 
 		TestExternalities::default().execute_with(|| {
 			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());
-			assert!(FooMap::decode_len(2).is_none());
+			assert_eq!(FooMap::decode_non_dedup_len(1).unwrap(), 3);
+			assert!(FooMap::decode_non_dedup_len(0).is_none());
+			assert!(FooMap::decode_non_dedup_len(2).is_none());
 		});
 
 		TestExternalities::default().execute_with(|| {
 			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());
-			assert!(FooDoubleMap::decode_len(1, 2).is_none());
-			assert!(FooDoubleMap::decode_len(2, 2).is_none());
+			assert_eq!(FooDoubleMap::decode_non_dedup_len(1, 1).unwrap(), 3);
+			assert!(FooDoubleMap::decode_non_dedup_len(2, 1).is_none());
+			assert!(FooDoubleMap::decode_non_dedup_len(1, 2).is_none());
+			assert!(FooDoubleMap::decode_non_dedup_len(2, 2).is_none());
 		});
 	}
 }

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.
@@ -346,6 +372,27 @@ pub trait StorageMap<K: FullEncode, V: FullCodec> {
 		V::decode_len(&Self::hashed_key_for(key))
 	}
 
+	/// Read the length of the storage value without decoding the entire value.
+	///
+	/// `V` 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
+	///
+	///  - `None` does not mean that `get()` does not return a value. The default value is completly
+	/// ignored by this function.
+	///
+	/// - The value returned is the non-deduplicated length of the underlying Vector in storage.This
+	/// means that any duplicate items are included.
+	fn decode_non_dedup_len<KeyArg: EncodeLike<K>>(key: KeyArg) -> Option<usize>
+	where
+		V: StorageDecodeNonDedupLength,
+	{
+		V::decode_non_dedup_len(&Self::hashed_key_for(key))
+	}
+
 	/// Migrate an item with the given `key` from a defunct `OldHasher` to the current hasher.
 	///
 	/// If the key doesn't exist, then it's a no-op. If it does, then it returns its value.
@@ -741,6 +788,27 @@ pub trait StorageDoubleMap<K1: FullEncode, K2: FullEncode, V: FullCodec> {
 		V::decode_len(&Self::hashed_key_for(key1, key2))
 	}
 
+	/// Read the length of the storage value without decoding the entire value under the
+	/// given `key1` and `key2`.
+	///
+	/// `V` 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
+	///
+	/// `None` does not mean that `get()` does not return a value. The default value is completly
+	/// ignored by this function.
+	fn decode_non_dedup_len<KArg1, KArg2>(key1: KArg1, key2: KArg2) -> Option<usize>
+	where
+		KArg1: EncodeLike<K1>,
+		KArg2: EncodeLike<K2>,
+		V: StorageDecodeNonDedupLength,
+	{
+		V::decode_non_dedup_len(&Self::hashed_key_for(key1, key2))
+	}
+
 	/// Migrate an item with the given `key1` and `key2` from defunct `OldHasher1` and
 	/// `OldHasher2` to the current hashers.
 	///
@@ -1400,8 +1468,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 +1488,29 @@ pub trait StorageDecodeLength: private::Sealed + codec::DecodeLength {
 	}
 }
 
+/// It is expected that the length is at the beginning of the encoded objectand 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> {
+		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()
+	}
+}
+
 /// Provides `Sealed` trait to prevent implementing trait `StorageAppend` & `StorageDecodeLength`
 /// & `EncodeLikeTuple` outside of this crate.
 mod private {
@@ -1471,7 +1561,14 @@ 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> {}
+
+// 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 +2123,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);
 		});
 	}
 }

substrate/frame/support/src/storage/types/double_map.rs

@@ -27,6 +27,7 @@ use crate::{
 	StorageHasher, Twox128,
 };
 use codec::{Decode, Encode, EncodeLike, FullCodec, MaxEncodedLen};
+use frame_support::storage::StorageDecodeNonDedupLength;
 use sp_arithmetic::traits::SaturatedConversion;
 use sp_metadata_ir::{StorageEntryMetadataIR, StorageEntryTypeIR};
 use sp_std::prelude::*;
@@ -455,6 +456,31 @@ where
 		<Self as crate::storage::StorageDoubleMap<Key1, Key2, Value>>::decode_len(key1, key2)
 	}
 
+	/// Read the length of the storage value without decoding the entire value.
+	///
+	/// `Value` 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
+	///
+	///  - `None` does not mean that `get()` does not return a value. The default value is completly
+	/// ignored by this function.
+	///
+	/// - The value returned is the non-deduplicated length of the underlying Vector in storage.This
+	/// means that any duplicate items are included.
+	pub fn decode_non_dedup_len<KArg1, KArg2>(key1: KArg1, key2: KArg2) -> Option<usize>
+	where
+		KArg1: EncodeLike<Key1>,
+		KArg2: EncodeLike<Key2>,
+		Value: StorageDecodeNonDedupLength,
+	{
+		<Self as crate::storage::StorageDoubleMap<Key1, Key2, Value>>::decode_non_dedup_len(
+			key1, key2,
+		)
+	}
+
 	/// Migrate an item with the given `key1` and `key2` from defunct `OldHasher1` and
 	/// `OldHasher2` to the current hashers.
 	///

substrate/frame/support/src/storage/types/map.rs

@@ -27,6 +27,7 @@ use crate::{
 	StorageHasher, Twox128,
 };
 use codec::{Decode, Encode, EncodeLike, FullCodec, MaxEncodedLen};
+use frame_support::storage::StorageDecodeNonDedupLength;
 use sp_arithmetic::traits::SaturatedConversion;
 use sp_metadata_ir::{StorageEntryMetadataIR, StorageEntryTypeIR};
 use sp_std::prelude::*;
@@ -285,6 +286,27 @@ where
 		<Self as crate::storage::StorageMap<Key, Value>>::decode_len(key)
 	}
 
+	/// Read the length of the storage value without decoding the entire value.
+	///
+	/// `Value` 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
+	///
+	///  - `None` does not mean that `get()` does not return a value. The default value is completly
+	/// ignored by this function.
+	///
+	/// - The value returned is the non-deduplicated length of the underlying Vector in storage.This
+	/// means that any duplicate items are included.
+	pub fn decode_non_dedup_len<KeyArg: EncodeLike<Key>>(key: KeyArg) -> Option<usize>
+	where
+		Value: StorageDecodeNonDedupLength,
+	{
+		<Self as crate::storage::StorageMap<Key, Value>>::decode_non_dedup_len(key)
+	}
+
 	/// Migrate an item with the given `key` from a defunct `OldHasher` to the current hasher.
 	///
 	/// If the key doesn't exist, then it's a no-op. If it does, then it returns its value.

substrate/frame/support/src/storage/types/value.rs

@@ -26,6 +26,7 @@ use crate::{
 	traits::{GetDefault, StorageInfo, StorageInstance},
 };
 use codec::{Decode, Encode, EncodeLike, FullCodec, MaxEncodedLen};
+use frame_support::storage::StorageDecodeNonDedupLength;
 use sp_arithmetic::traits::SaturatedConversion;
 use sp_metadata_ir::{StorageEntryMetadataIR, StorageEntryTypeIR};
 use sp_std::prelude::*;
@@ -233,6 +234,27 @@ where
 		<Self as crate::storage::StorageValue<Value>>::decode_len()
 	}
 
+	/// Read the length of the storage value without decoding the entire value.
+	///
+	/// `Value` 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
+	///
+	///  - `None` does not mean that `get()` does not return a value. The default value is completly
+	/// ignored by this function.
+	///
+	/// - The value returned is the non-deduplicated length of the underlying Vector in storage.This
+	/// means that any duplicate items are included.
+	pub fn decode_non_dedup_len() -> Option<usize>
+	where
+		Value: StorageDecodeNonDedupLength,
+	{
+		<Self as crate::storage::StorageValue<Value>>::decode_non_dedup_len()
+	}
+
 	/// Try and append the given item to the value in the storage.
 	///
 	/// Is only available if `Value` of the storage implements [`StorageTryAppend`].