demonstration

frame/support/procedural/src/pallet/expand/storage.rs

@@ -130,6 +130,7 @@ pub fn process_generics(def: &mut Def) -> syn::Result<Vec<ResultOnEmptyStructMet
 		};
 
 		let prefix_ident = prefix_ident(storage_def);
+		// TODO: maybe use a helper function
 		let type_use_gen = if def.config.has_instance {
 			quote::quote_spanned!(storage_def.attr_span => T, I)
 		} else {
@@ -631,12 +632,36 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 		)
 	});
 
+	// aggregated where clause of all storage types and the whole pallet.
 	let mut where_clauses = vec![&def.config.where_clause];
 	where_clauses.extend(def.storages.iter().map(|storage| &storage.where_clause));
 	let completed_where_clause = super::merge_where_clauses(&where_clauses);
 	let type_impl_gen = &def.type_impl_generics(proc_macro2::Span::call_site());
 	let type_use_gen = &def.type_use_generics(proc_macro2::Span::call_site());
 
+	let try_decode_entire_state = {
+		let decode_each = def.storages.iter().map(|storage_def| {
+			let ident = &storage_def.ident;
+			let gen = &def.type_use_generics(storage_def.attr_span);
+			let full_ident = quote::quote_spanned!(storage_def.attr_span => #ident<#gen> );
+			quote::quote!(
+				total_size += <#full_ident as #frame_support::traits::DecodeAll>::decode_all()?;
+			)
+		});
+		quote::quote!(
+			#[cfg(any(feature = "try-runtime", test))]
+			impl<#type_impl_gen> #frame_support::traits::TryDecodeEntireStorage
+				for #pallet_ident<#type_use_gen> #completed_where_clause
+			{
+				fn try_decode_entire_state() -> Result<usize, &'static str> {
+					let mut total_size = 0usize;
+					#(#decode_each)*
+					Ok(total_size)
+				}
+			}
+		)
+	};
+
 	quote::quote!(
 		impl<#type_impl_gen> #pallet_ident<#type_use_gen>
 			#completed_where_clause
@@ -662,5 +687,7 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 		#( #getters )*
 		#( #prefix_structs )*
 		#( #on_empty_structs )*
+
+		#try_decode_entire_state
 	)
 }

frame/support/procedural/src/pallet/parse/storage.rs

@@ -150,7 +150,7 @@ pub enum QueryKind {
 /// `type MyStorage = StorageValue<MyStorageP, u32>`
 /// The keys and values types are parsed in order to get metadata
 pub struct StorageDef {
-	/// The index of error item in pallet module.
+	/// The index of storage item in pallet module.
 	pub index: usize,
 	/// Visibility of the storage type.
 	pub vis: syn::Visibility,

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

@@ -24,10 +24,10 @@
 //!
 //! This is internal api and is subject to change.
 
-mod double_map;
+pub(crate) mod double_map;
 pub(crate) mod map;
-mod nmap;
-mod value;
+pub(crate) mod nmap;
+pub(crate) mod value;
 
 pub use double_map::StorageDoubleMap;
 pub use map::StorageMap;

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

@@ -24,7 +24,7 @@ use crate::{
 		types::{OptionQuery, QueryKindTrait, StorageEntryMetadataBuilder},
 		StorageAppend, StorageDecodeLength, StorageTryAppend,
 	},
-	traits::{GetDefault, StorageInfo, StorageInstance},
+	traits::{GetDefault, StorageInfo, StorageInstance, Get},
 };
 use codec::{Decode, Encode, EncodeLike, FullCodec, MaxEncodedLen};
 use sp_arithmetic::traits::SaturatedConversion;
@@ -46,7 +46,7 @@ where
 	Prefix: StorageInstance,
 	Value: FullCodec,
 	QueryKind: QueryKindTrait<Value, OnEmpty>,
-	OnEmpty: crate::traits::Get<QueryKind::Query> + 'static,
+	OnEmpty: Get<QueryKind::Query> + 'static,
 {
 	type Query = QueryKind::Query;
 	fn module_prefix() -> &'static [u8] {

frame/support/src/traits.rs

@@ -121,4 +121,4 @@ pub use messages::{
 #[cfg(feature = "try-runtime")]
 mod try_runtime;
 #[cfg(feature = "try-runtime")]
-pub use try_runtime::{Select as TryStateSelect, TryState};
+pub use try_runtime::{Select as TryStateSelect, TryState, TryDecodeEntireStorage};

frame/support/src/traits/try_runtime.rs

@@ -17,8 +17,18 @@
 
 //! Try-runtime specific traits and types.
 
+use super::StorageInstance;
+use crate::{
+	storage::{
+		generator::value::StorageValue,
+		types::{CountedStorageMapInstance, KeyGenerator, QueryKindTrait},
+	},
+	StorageHasher,
+};
+use codec::{Decode, FullCodec};
 use impl_trait_for_tuples::impl_for_tuples;
 use sp_arithmetic::traits::AtLeast32BitUnsigned;
+use sp_core::Get;
 use sp_std::prelude::*;
 
 // Which state tests to execute.
@@ -42,6 +52,146 @@ impl Default for Select {
 	}
 }
 
+use crate::storage::generator::map::StorageMap;
+
+/// Decode the entire values under the given storage type.
+///
+/// For values, this is trivial. For all kinds of maps, it should decode all the values associated
+/// with all keys existing in the map.
+
+pub trait TryDecodeEntireStorage {
+	fn try_decode_entire_state() -> Result<usize, &'static str>;
+}
+
+#[cfg_attr(all(not(feature = "tuples-96"), not(feature = "tuples-128")), impl_for_tuples(64))]
+#[cfg_attr(all(feature = "tuples-96", not(feature = "tuples-128")), impl_for_tuples(96))]
+#[cfg_attr(feature = "tuples-128", impl_for_tuples(128))]
+impl TryDecodeEntireStorage for Tuple {
+	fn try_decode_entire_state() -> Result<usize, &'static str> {
+		let mut len = 0usize;
+		for_tuples!( #( len = len.saturating_add(Tuple::try_decode_entire_state()?); )* );
+		Ok(len)
+	}
+}
+
+fn decode_key<V: Decode>(key: &[u8]) -> Result<usize, &'static str> {
+	match sp_io::storage::get(key) {
+		None => Ok(0),
+		Some(bytes) => {
+			let len = bytes.len();
+			let _ = <V as Decode>::decode(&mut bytes.as_ref()).map_err(|_| "TODO")?;
+			Ok(len)
+		},
+	}
+}
+
+impl<Prefix, Value, QueryKind, OnEmpty> TryDecodeEntireStorage
+	for crate::storage::types::StorageValue<Prefix, Value, QueryKind, OnEmpty>
+where
+	Prefix: StorageInstance,
+	Value: FullCodec,
+	QueryKind: QueryKindTrait<Value, OnEmpty>,
+	OnEmpty: Get<QueryKind::Query> + 'static,
+{
+	fn try_decode_entire_state() -> Result<usize, &'static str> {
+		let len = decode_key::<Value>(&Self::hashed_key())?;
+		// TODO: it will be easier to long such things if we implement these one by one in the
+		// macro, which is more or less then same thing as the blanket implementation.
+		log::debug!(
+			target: crate::LOG_TARGET,
+			"decoded {:?}::{:?}, total of {} bytes",
+			Self::module_prefix(),
+			Self::storage_prefix(),
+			len
+		);
+		Ok(len)
+	}
+}
+
+impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty> TryDecodeEntireStorage
+	for crate::storage::types::StorageMap<Prefix, Hasher, Key, Value, QueryKind, OnEmpty>
+where
+	Prefix: StorageInstance,
+	Hasher: StorageHasher,
+	Key: FullCodec,
+	Value: FullCodec,
+	QueryKind: QueryKindTrait<Value, OnEmpty>,
+	OnEmpty: Get<QueryKind::Query> + 'static,
+{
+	fn try_decode_entire_state() -> Result<usize, &'static str> {
+		let mut next_key = Self::prefix_hash();
+		let mut len = 0usize;
+		while let Some(key) = sp_io::storage::next_key(&next_key) {
+			len += decode_key::<Value>(&key)?;
+			next_key = key;
+		}
+
+		Ok(len)
+	}
+}
+
+impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty> TryDecodeEntireStorage
+	for crate::storage::types::CountedStorageMap<Prefix, Hasher, Key, Value, QueryKind, OnEmpty>
+where
+	Prefix: CountedStorageMapInstance,
+	Hasher: StorageHasher,
+	Key: FullCodec,
+	Value: FullCodec,
+	QueryKind: QueryKindTrait<Value, OnEmpty>,
+	OnEmpty: Get<QueryKind::Query> + 'static,
+{
+	fn try_decode_entire_state() -> Result<usize, &'static str> {
+		// TODO: rather not re-hardcode the u32 here.
+		let mut len = decode_key::<u32>(&Self::counter_storage_final_key())?;
+		let mut next_key = Self::map_storage_final_prefix();
+		while let Some(key) = sp_io::storage::next_key(&next_key) {
+			len += decode_key::<Value>(&key)?;
+			next_key = key;
+		}
+
+		Ok(len)
+	}
+}
+
+impl<Prefix, Hasher1, Key1, Hasher2, Key2, Value, QueryKind, OnEmpty> TryDecodeEntireStorage
+	for crate::storage::types::StorageDoubleMap<
+		Prefix,
+		Hasher1,
+		Key1,
+		Hasher2,
+		Key2,
+		Value,
+		QueryKind,
+		OnEmpty,
+	> where
+	Prefix: StorageInstance,
+	Hasher1: StorageHasher,
+	Key1: FullCodec,
+	Hasher2: StorageHasher,
+	Key2: FullCodec,
+	Value: FullCodec,
+	QueryKind: QueryKindTrait<Value, OnEmpty>,
+	OnEmpty: Get<QueryKind::Query> + 'static,
+{
+	fn try_decode_entire_state() -> Result<usize, &'static str> {
+		todo!();
+	}
+}
+
+impl<Prefix, Key, Value, QueryKind, OnEmpty> TryDecodeEntireStorage
+	for crate::storage::types::StorageNMap<Prefix, Key, Value, QueryKind, OnEmpty>
+where
+	Prefix: StorageInstance,
+	Key: KeyGenerator,
+	Value: FullCodec,
+	QueryKind: QueryKindTrait<Value, OnEmpty>,
+	OnEmpty: Get<QueryKind::Query> + 'static,
+{
+	fn try_decode_entire_state() -> Result<usize, &'static str> {
+		todo!();
+	}
+}
+
 impl sp_std::fmt::Debug for Select {
 	fn fmt(&self, f: &mut sp_std::fmt::Formatter<'_>) -> sp_std::fmt::Result {
 		match self {
@@ -138,3 +288,60 @@ impl<BlockNumber: Clone + sp_std::fmt::Debug + AtLeast32BitUnsigned> TryState<Bl
 		}
 	}
 }
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use crate::Blake2_128Concat;
+
+	struct Prefix;
+	impl StorageInstance for Prefix {
+		fn pallet_prefix() -> &'static str {
+			"test_pallet"
+		}
+		const STORAGE_PREFIX: &'static str = "val_prefix";
+	}
+
+	type Value = crate::storage::types::StorageValue<Prefix, u32>;
+	type Map = crate::storage::types::StorageMap<Prefix, Blake2_128Concat, u32, u32>;
+
+	#[test]
+	fn try_decode_entire_state_value_works() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			assert_eq!(Value::try_decode_entire_state(), Ok(0));
+			Value::put(42);
+			assert_eq!(Value::try_decode_entire_state(), Ok(4));
+			Value::kill();
+			assert_eq!(Value::try_decode_entire_state(), Ok(0));
+
+			// two bytes, cannot be decoded into u32.
+			sp_io::storage::set(&Value::hashed_key(), &[0u8, 1]);
+			assert!(Value::try_decode_entire_state().is_err());
+		})
+	}
+
+	#[test]
+	fn try_decode_entire_state_map_works() {
+		todo!()
+	}
+
+	#[test]
+	fn try_decode_entire_state_double_works() {
+		todo!()
+	}
+
+	#[test]
+	fn try_decode_entire_state_nmap_works() {
+		todo!()
+	}
+
+	#[test]
+	fn try_decode_entire_state_counted_map_works() {
+		todo!()
+	}
+
+	#[test]
+	fn try_decode_entire_state_tuple_of_storage_works() {
+		assert!(<(Value, Map) as TryDecodeEntireStorage>::try_decode_entire_state().is_ok());
+	}
+}