From 80cc9dae631683fb0e353b1402eb9757fcd5973f Mon Sep 17 00:00:00 2001
From: kianenigma <kian@parity.io>
Date: Fri, 23 Dec 2022 14:49:52 +0100
Subject: [PATCH 01/10] demonstration

---
 .../procedural/src/pallet/expand/storage.rs   |  27 +++
 .../procedural/src/pallet/parse/storage.rs    |   2 +-
 frame/support/src/storage/generator/mod.rs    |   6 +-
 frame/support/src/storage/types/value.rs      |   4 +-
 frame/support/src/traits.rs                   |   2 +-
 frame/support/src/traits/try_runtime.rs       | 207 ++++++++++++++++++
 6 files changed, 241 insertions(+), 7 deletions(-)

diff --git a/frame/support/procedural/src/pallet/expand/storage.rs b/frame/support/procedural/src/pallet/expand/storage.rs
index 181f35b545496..68612081aab4a 100644
--- a/frame/support/procedural/src/pallet/expand/storage.rs
+++ b/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
 	)
 }
diff --git a/frame/support/procedural/src/pallet/parse/storage.rs b/frame/support/procedural/src/pallet/parse/storage.rs
index 8b551ab31d6c3..15d165e23df41 100644
--- a/frame/support/procedural/src/pallet/parse/storage.rs
+++ b/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,
diff --git a/frame/support/src/storage/generator/mod.rs b/frame/support/src/storage/generator/mod.rs
index 334e9b9e24e86..90af430ae5b83 100644
--- a/frame/support/src/storage/generator/mod.rs
+++ b/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;
diff --git a/frame/support/src/storage/types/value.rs b/frame/support/src/storage/types/value.rs
index f145e9fb30414..3875ccbaa4fb6 100644
--- a/frame/support/src/storage/types/value.rs
+++ b/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] {
diff --git a/frame/support/src/traits.rs b/frame/support/src/traits.rs
index 63c86c1f68459..550d50cdb8fc4 100644
--- a/frame/support/src/traits.rs
+++ b/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};
diff --git a/frame/support/src/traits/try_runtime.rs b/frame/support/src/traits/try_runtime.rs
index f741ca56a56fc..aab0a26618023 100644
--- a/frame/support/src/traits/try_runtime.rs
+++ b/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());
+	}
+}

From 2456895383c004dbe11b9b50dd3bfad63a74a6c4 Mon Sep 17 00:00:00 2001
From: kianenigma <kian@parity.io>
Date: Mon, 26 Dec 2022 09:47:56 +0100
Subject: [PATCH 02/10] remove empty

---
 frame/support/src/traits/try_runtime.rs | 1 -
 1 file changed, 1 deletion(-)

diff --git a/frame/support/src/traits/try_runtime.rs b/frame/support/src/traits/try_runtime.rs
index aab0a26618023..8bb5e979c5cfe 100644
--- a/frame/support/src/traits/try_runtime.rs
+++ b/frame/support/src/traits/try_runtime.rs
@@ -58,7 +58,6 @@ use crate::storage::generator::map::StorageMap;
 ///
 /// 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>;
 }

From db4b23422de7a40da346000c87c5f91e6e494fac Mon Sep 17 00:00:00 2001
From: kianenigma <kian@parity.io>
Date: Fri, 17 Mar 2023 10:06:45 -0600
Subject: [PATCH 03/10] make it all work again

---
 .../support/src/storage/types/counted_map.rs  |  14 +-
 frame/support/src/storage/types/mod.rs        |   2 +-
 frame/support/src/storage/types/value.rs      |   2 +-
 frame/support/src/traits.rs                   |   4 +-
 frame/support/src/traits/try_runtime.rs       | 240 +++++++++++++-----
 5 files changed, 191 insertions(+), 71 deletions(-)

diff --git a/frame/support/src/storage/types/counted_map.rs b/frame/support/src/storage/types/counted_map.rs
index 24b00be485e49..fca5cd37cec9d 100644
--- a/frame/support/src/storage/types/counted_map.rs
+++ b/frame/support/src/storage/types/counted_map.rs
@@ -74,7 +74,11 @@ impl<P: CountedStorageMapInstance, H, K, V, Q, O, M> MapWrapper
 	type Map = StorageMap<P, H, K, V, Q, O, M>;
 }
 
-type CounterFor<P> = StorageValue<<P as CountedStorageMapInstance>::CounterPrefix, u32, ValueQuery>;
+/// The numeric counter type.
+pub type Counter = u32;
+
+type CounterFor<P> =
+	StorageValue<<P as CountedStorageMapInstance>::CounterPrefix, Counter, ValueQuery>;
 
 /// On removal logic for updating counter while draining upon some prefix with
 /// [`crate::storage::PrefixIterator`].
@@ -378,14 +382,14 @@ where
 	/// can be very heavy, so use with caution.
 	///
 	/// Returns the number of items in the map which is used to set the counter.
-	pub fn initialize_counter() -> u32 {
-		let count = Self::iter_values().count() as u32;
+	pub fn initialize_counter() -> Counter {
+		let count = Self::iter_values().count() as Counter;
 		CounterFor::<Prefix>::set(count);
 		count
 	}
 
 	/// Return the count.
-	pub fn count() -> u32 {
+	pub fn count() -> Counter {
 		CounterFor::<Prefix>::get()
 	}
 }
@@ -1161,7 +1165,7 @@ mod test {
 				StorageEntryMetadata {
 					name: "counter_for_foo",
 					modifier: StorageEntryModifier::Default,
-					ty: StorageEntryType::Plain(scale_info::meta_type::<u32>()),
+					ty: StorageEntryType::Plain(scale_info::meta_type::<Counter>()),
 					default: vec![0, 0, 0, 0],
 					docs: if cfg!(feature = "no-metadata-docs") {
 						vec![]
diff --git a/frame/support/src/storage/types/mod.rs b/frame/support/src/storage/types/mod.rs
index 9a6f15d1ee820..b527ef9a1b194 100644
--- a/frame/support/src/storage/types/mod.rs
+++ b/frame/support/src/storage/types/mod.rs
@@ -29,7 +29,7 @@ mod map;
 mod nmap;
 mod value;
 
-pub use counted_map::{CountedStorageMap, CountedStorageMapInstance};
+pub use counted_map::{CountedStorageMap, CountedStorageMapInstance, Counter};
 pub use double_map::StorageDoubleMap;
 pub use key::{
 	EncodeLikeTuple, HasKeyPrefix, HasReversibleKeyPrefix, Key, KeyGenerator,
diff --git a/frame/support/src/storage/types/value.rs b/frame/support/src/storage/types/value.rs
index 2789b90dd2f51..3f23d564f11d4 100644
--- a/frame/support/src/storage/types/value.rs
+++ b/frame/support/src/storage/types/value.rs
@@ -24,7 +24,7 @@ use crate::{
 		types::{OptionQuery, QueryKindTrait, StorageEntryMetadataBuilder},
 		StorageAppend, StorageDecodeLength, StorageTryAppend,
 	},
-	traits::{GetDefault, StorageInfo, StorageInstance, Get},
+	traits::{Get, GetDefault, StorageInfo, StorageInstance},
 };
 use codec::{Decode, Encode, EncodeLike, FullCodec, MaxEncodedLen};
 use sp_arithmetic::traits::SaturatedConversion;
diff --git a/frame/support/src/traits.rs b/frame/support/src/traits.rs
index b827794d55cd0..101dd40b2a2da 100644
--- a/frame/support/src/traits.rs
+++ b/frame/support/src/traits.rs
@@ -121,4 +121,6 @@ pub use messages::{
 #[cfg(feature = "try-runtime")]
 mod try_runtime;
 #[cfg(feature = "try-runtime")]
-pub use try_runtime::{Select as TryStateSelect, TryState, UpgradeCheckSelect, TryDecodeEntireStorage};
+pub use try_runtime::{
+	Select as TryStateSelect, TryDecodeEntireStorage, TryState, UpgradeCheckSelect,
+};
diff --git a/frame/support/src/traits/try_runtime.rs b/frame/support/src/traits/try_runtime.rs
index 1b7f4ff45c95f..2d6bb21facb23 100644
--- a/frame/support/src/traits/try_runtime.rs
+++ b/frame/support/src/traits/try_runtime.rs
@@ -19,10 +19,8 @@
 
 use super::StorageInstance;
 use crate::{
-	storage::{
-		generator::value::StorageValue,
-		types::{CountedStorageMapInstance, KeyGenerator, QueryKindTrait},
-	},
+	storage::types::{CountedStorageMapInstance, KeyGenerator, QueryKindTrait},
+	traits::{PartialStorageInfoTrait, StorageInfo},
 	StorageHasher,
 };
 use codec::{Decode, FullCodec};
@@ -52,13 +50,15 @@ impl Default for Select {
 	}
 }
 
-use crate::storage::generator::map::StorageMap;
-
-/// Decode the entire values under the given storage type.
+/// Decode the entire data 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.
+///
+/// Tuple implementations are provided and simply decode each type in the tuple, summing up the
+/// decoded bytes if `Ok(_)`.
 pub trait TryDecodeEntireStorage {
+	/// Decode the entire data under the given storage, returning `Ok(bytes_decoded)` if success.
 	fn try_decode_entire_state() -> Result<usize, &'static str>;
 }
 
@@ -73,15 +73,43 @@ impl TryDecodeEntireStorage for Tuple {
 	}
 }
 
-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)
-		},
+/// Decode all the values based on the prefix of `info` to `V`.
+///
+/// Basically, it decodes and sums up all the values who's key start with `info.prefix`. For values,
+/// this would be the value itself. For all sorts of maps, this should be all map items in the
+/// absence of key collision.
+fn decode_storage_info<V: Decode>(info: StorageInfo) -> Result<usize, &'static str> {
+	let mut next_key = info.prefix.clone();
+	let mut decoded = 0;
+
+	let decode_key = |key: &[u8]| {
+		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(|_| {
+					log::error!(target: crate::LOG_TARGET, "failed to decoded {:?}", info,);
+					"failed to decode value under existing key"
+				})?;
+				Ok::<usize, &'static str>(len)
+			},
+		}
+	};
+
+	decoded += decode_key(&next_key)?;
+	loop {
+		match sp_io::storage::next_key(&next_key) {
+			Some(key) if key.starts_with(&info.prefix) => {
+				decoded += decode_key(&key)?;
+				next_key = key;
+			},
+			_ => break
+		}
 	}
+
+	Ok(decoded)
 }
 
 impl<Prefix, Value, QueryKind, OnEmpty> TryDecodeEntireStorage
@@ -93,17 +121,11 @@ where
 	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)
+		let info = Self::partial_storage_info()
+			.first()
+			.cloned()
+			.expect("Value has only one storage info; qed");
+		decode_storage_info::<Value>(info)
 	}
 }
 
@@ -118,14 +140,11 @@ where
 	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)
+		let info = Self::partial_storage_info()
+			.first()
+			.cloned()
+			.expect("Map has only one storage info; qed");
+		decode_storage_info::<Value>(info)
 	}
 }
 
@@ -140,15 +159,13 @@ where
 	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)
+		let (map_info, counter_info) = match &Self::partial_storage_info()[..] {
+			[a, b] => (a.clone(), b.clone()), // TODO better way?
+			_ => panic!("Counted map has two storage info items; qed"),
+		};
+		let mut decoded = decode_storage_info::<Value>(counter_info)?;
+		decoded += decode_storage_info::<Value>(map_info)?;
+		Ok(decoded)
 	}
 }
 
@@ -173,7 +190,11 @@ impl<Prefix, Hasher1, Key1, Hasher2, Key2, Value, QueryKind, OnEmpty> TryDecodeE
 	OnEmpty: Get<QueryKind::Query> + 'static,
 {
 	fn try_decode_entire_state() -> Result<usize, &'static str> {
-		todo!();
+		let info = Self::partial_storage_info()
+			.first()
+			.cloned()
+			.expect("Double-map has only one storage info; qed");
+		decode_storage_info::<Value>(info)
 	}
 }
 
@@ -187,7 +208,11 @@ where
 	OnEmpty: Get<QueryKind::Query> + 'static,
 {
 	fn try_decode_entire_state() -> Result<usize, &'static str> {
-		todo!();
+		let info = Self::partial_storage_info()
+			.first()
+			.cloned() // TODO better way?
+			.expect("N-map has only one storage info; qed");
+		decode_storage_info::<Value>(info)
 	}
 }
 
@@ -332,24 +357,46 @@ impl<BlockNumber: Clone + sp_std::fmt::Debug + AtLeast32BitUnsigned> TryState<Bl
 mod tests {
 	use super::*;
 	use crate::Blake2_128Concat;
-
-	struct Prefix;
-	impl StorageInstance for Prefix {
-		fn pallet_prefix() -> &'static str {
-			"test_pallet"
+	use crate::storage::types;
+
+	// TODO reuse?
+	macro_rules! build_prefix {
+		($name:ident) => {
+			struct $name;
+			impl StorageInstance for $name {
+				fn pallet_prefix() -> &'static str {
+					"test_pallet"
+				}
+				const STORAGE_PREFIX: &'static str = stringify!($name);
+			}
 		}
-		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>;
+	build_prefix!(ValuePrefix);
+	type Value = types::StorageValue<ValuePrefix, u32>;
+
+	build_prefix!(MapPrefix);
+	type Map = types::StorageMap<MapPrefix, Blake2_128Concat, u32, u32>;
+
+	build_prefix!(CMapCounterPrefix);
+	build_prefix!(CMapPrefix);
+	impl CountedStorageMapInstance for CMapPrefix {
+		type CounterPrefix = CMapCounterPrefix;
+	}
+	type CMap = types::CountedStorageMap<CMapPrefix, Blake2_128Concat, u32, u32>;
+
+	build_prefix!(DMapPrefix);
+	type DMap =
+		types::StorageDoubleMap<DMapPrefix, Blake2_128Concat, u32, 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));
 
@@ -361,26 +408,93 @@ mod tests {
 
 	#[test]
 	fn try_decode_entire_state_map_works() {
-		todo!()
-	}
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			assert_eq!(Map::try_decode_entire_state(), Ok(0));
 
-	#[test]
-	fn try_decode_entire_state_double_works() {
-		todo!()
+			Map::insert(0, 42);
+			assert_eq!(Map::try_decode_entire_state(), Ok(4));
+
+			Map::insert(0, 42);
+			assert_eq!(Map::try_decode_entire_state(), Ok(4));
+
+			Map::insert(1, 42);
+			assert_eq!(Map::try_decode_entire_state(), Ok(8));
+
+			Map::remove(0);
+			assert_eq!(Map::try_decode_entire_state(), Ok(4));
+
+			// two bytes, cannot be decoded into u32.
+			sp_io::storage::set(&Map::hashed_key_for(2), &[0u8, 1]);
+			assert!(Map::try_decode_entire_state().is_err());
+		})
 	}
 
 	#[test]
-	fn try_decode_entire_state_nmap_works() {
-		todo!()
+	fn try_decode_entire_state_counted_map_works() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			// counter is not even initialized;
+			assert_eq!(CMap::try_decode_entire_state(), Ok(0 + 0));
+
+			let counter = 4;
+
+			CMap::insert(0, 42);
+			assert_eq!(CMap::try_decode_entire_state(), Ok(4 + counter));
+
+			CMap::insert(0, 42);
+			assert_eq!(CMap::try_decode_entire_state(), Ok(4 + counter));
+
+			CMap::insert(1, 42);
+			assert_eq!(CMap::try_decode_entire_state(), Ok(8 + counter));
+
+			CMap::remove(0);
+			assert_eq!(CMap::try_decode_entire_state(), Ok(4 + counter));
+
+			// counter is cleared again.
+			CMap::clear(u32::MAX, None).unwrap();
+			assert_eq!(CMap::try_decode_entire_state(), Ok(0 + 0));
+
+			// two bytes, cannot be decoded into u32.
+			sp_io::storage::set(&CMap::hashed_key_for(2), &[0u8, 1]);
+			assert!(CMap::try_decode_entire_state().is_err());
+		})
 	}
 
 	#[test]
-	fn try_decode_entire_state_counted_map_works() {
-		todo!()
+	fn try_decode_entire_state_double_works() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			assert_eq!(DMap::try_decode_entire_state(), Ok(0));
+
+			DMap::insert(0, 0, 42);
+			assert_eq!(DMap::try_decode_entire_state(), Ok(4));
+
+			DMap::insert(0, 0, 42);
+			assert_eq!(DMap::try_decode_entire_state(), Ok(4));
+
+			DMap::insert(0, 1, 42);
+			assert_eq!(DMap::try_decode_entire_state(), Ok(8));
+
+			DMap::insert(1, 0, 42);
+			assert_eq!(DMap::try_decode_entire_state(), Ok(12));
+
+			DMap::remove(0, 0);
+			assert_eq!(DMap::try_decode_entire_state(), Ok(8));
+
+			// two bytes, cannot be decoded into u32.
+			sp_io::storage::set(&DMap::hashed_key_for(1, 1), &[0u8, 1]);
+			assert!(DMap::try_decode_entire_state().is_err());
+		})
 	}
 
 	#[test]
 	fn try_decode_entire_state_tuple_of_storage_works() {
-		assert!(<(Value, Map) as TryDecodeEntireStorage>::try_decode_entire_state().is_ok());
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			assert_eq!(<(Value, Map) as TryDecodeEntireStorage>::try_decode_entire_state(), Ok(0));
+
+			Value::put(42);
+			assert_eq!(<(Value, Map) as TryDecodeEntireStorage>::try_decode_entire_state(), Ok(4));
+
+			Map::insert(0, 42);
+			assert_eq!(<(Value, Map) as TryDecodeEntireStorage>::try_decode_entire_state(), Ok(8));
+		});
 	}
 }

From 4cd6761854e0c51077c8e57c88ccd92b241b0257 Mon Sep 17 00:00:00 2001
From: kianenigma <kian@parity.io>
Date: Fri, 17 Mar 2023 13:08:08 -0600
Subject: [PATCH 04/10] all tests work

---
 bin/node-template/runtime/src/lib.rs          |  9 +++
 bin/node/runtime/src/lib.rs                   |  2 +
 frame/glutton/src/lib.rs                      |  5 +-
 .../procedural/src/pallet/expand/storage.rs   | 27 ++++----
 frame/support/src/storage/migration.rs        | 33 ++++++++++
 frame/support/src/traits/try_runtime.rs       | 62 +++++++++++--------
 6 files changed, 94 insertions(+), 44 deletions(-)

diff --git a/bin/node-template/runtime/src/lib.rs b/bin/node-template/runtime/src/lib.rs
index 57e4470c762ac..bbd4f2d9f4f63 100644
--- a/bin/node-template/runtime/src/lib.rs
+++ b/bin/node-template/runtime/src/lib.rs
@@ -318,6 +318,14 @@ pub type SignedExtra = (
 	pallet_transaction_payment::ChargeTransactionPayment<Runtime>,
 );
 
+/// All migrations of the runtime, aside from the ones declared in the pallets.
+///
+/// This can be a tuple of types, each implementing `OnRuntimeUpgrade`. Add other migration types
+/// before `EnsureStateDecodes` as needed -- this is only for testing, and
+// should come last.
+#[allow(unused_parens)]
+type Migrations = (frame_support::migration::EnsureStateDecodes<AllPalletsWithSystem>);
+
 /// Unchecked extrinsic type as expected by this runtime.
 pub type UncheckedExtrinsic =
 	generic::UncheckedExtrinsic<Address, RuntimeCall, Signature, SignedExtra>;
@@ -330,6 +338,7 @@ pub type Executive = frame_executive::Executive<
 	frame_system::ChainContext<Runtime>,
 	Runtime,
 	AllPalletsWithSystem,
+	Migrations,
 >;
 
 #[cfg(feature = "runtime-benchmarks")]
diff --git a/bin/node/runtime/src/lib.rs b/bin/node/runtime/src/lib.rs
index 0c21fcc904fa1..e49358b277681 100644
--- a/bin/node/runtime/src/lib.rs
+++ b/bin/node/runtime/src/lib.rs
@@ -1860,6 +1860,8 @@ type Migrations = (
 	pallet_nomination_pools::migration::v2::MigrateToV2<Runtime>,
 	pallet_alliance::migration::Migration<Runtime>,
 	pallet_contracts::Migration<Runtime>,
+	// This should always be the last migration item.
+	frame_support::storage::migration::EnsureStateDecodes<AllPalletsWithSystem>,
 );
 
 /// MMR helper types.
diff --git a/frame/glutton/src/lib.rs b/frame/glutton/src/lib.rs
index e9a46374a5ade..3d0dacfe5f17f 100644
--- a/frame/glutton/src/lib.rs
+++ b/frame/glutton/src/lib.rs
@@ -17,9 +17,8 @@
 
 //! # Glutton Pallet
 //!
-//! Pallet that consumes `ref_time` and `proof_size` of a block. Based on the
-//! `Compute` and `Storage` parameters the pallet consumes the adequate amount
-//! of weight.
+//! Pallet that consumes `ref_time` and `proof_size` of a block. Based on the `Compute` and
+//! `Storage` parameters the pallet consumes the adequate amount of weight.
 
 #![cfg_attr(not(feature = "std"), no_std)]
 
diff --git a/frame/support/procedural/src/pallet/expand/storage.rs b/frame/support/procedural/src/pallet/expand/storage.rs
index 6a45c0b8159e2..0a8a8e15512b9 100644
--- a/frame/support/procedural/src/pallet/expand/storage.rs
+++ b/frame/support/procedural/src/pallet/expand/storage.rs
@@ -127,7 +127,6 @@ 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 {
@@ -637,23 +636,23 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 	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()?;
-			)
-		});
+		let storage_names = def
+			.storages
+			.iter()
+			.map(|storage| {
+				let ident = &storage.ident;
+				let gen = &def.type_use_generics(storage.attr_span);
+				quote::quote_spanned!(storage.attr_span => #ident<#gen> )
+			})
+			.collect::<Vec<_>>();
 		quote::quote!(
-			#[cfg(any(feature = "try-runtime", test))]
+			#[cfg(feature = "try-runtime")]
 			impl<#type_impl_gen> #frame_support::traits::TryDecodeEntireStorage
-				for #pallet_ident<#type_use_gen> #completed_where_clause
+			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)
+					// simply delegate impl to a tuple of all storage items we have.
+					<( #( #storage_names ),*) as frame_support::traits::TryDecodeEntireStorage>::try_decode_entire_state()
 				}
 			}
 		)
diff --git a/frame/support/src/storage/migration.rs b/frame/support/src/storage/migration.rs
index 8e945afdb6441..13816c157619e 100644
--- a/frame/support/src/storage/migration.rs
+++ b/frame/support/src/storage/migration.rs
@@ -389,6 +389,39 @@ pub fn move_prefix(from_prefix: &[u8], to_prefix: &[u8]) {
 	}
 }
 
+/// A phony migration that does nothing, except executing [`TryDecodeEntireStorage`] on
+/// `post_upgrade`, which implies it is only available if `try-state` feature is used.
+///
+/// This can be used typically in the top level runtime, where `AllPallets` typically comes from
+/// `construct_runtime!`.
+pub struct EnsureStateDecodes<AllPallets>(sp_std::marker::PhantomData<AllPallets>);
+
+#[cfg(not(feature = "try-runtime"))]
+impl<AllPallets> crate::traits::OnRuntimeUpgrade for EnsureStateDecodes<AllPallets> {
+	fn on_runtime_upgrade() -> crate::weights::Weight {
+		Default::default()
+	}
+}
+
+#[cfg(feature = "try-runtime")]
+impl<AllPallets: crate::traits::TryDecodeEntireStorage> crate::traits::OnRuntimeUpgrade
+	for EnsureStateDecodes<AllPallets>
+{
+	fn on_runtime_upgrade() -> sp_weights::Weight {
+		Default::default()
+	}
+
+	fn post_upgrade(_: Vec<u8>) -> Result<(), &'static str> {
+		let decoded = AllPallets::try_decode_entire_state()?;
+		crate::log::info!(
+			target: crate::LOG_TARGET,
+			"decoded the entire state, total size = {} bytes",
+			decoded
+		);
+		Ok(())
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::{
diff --git a/frame/support/src/traits/try_runtime.rs b/frame/support/src/traits/try_runtime.rs
index 2d6bb21facb23..15dd7d97492d2 100644
--- a/frame/support/src/traits/try_runtime.rs
+++ b/frame/support/src/traits/try_runtime.rs
@@ -82,20 +82,16 @@ fn decode_storage_info<V: Decode>(info: StorageInfo) -> Result<usize, &'static s
 	let mut next_key = info.prefix.clone();
 	let mut decoded = 0;
 
-	let decode_key = |key: &[u8]| {
-		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(|_| {
-					log::error!(target: crate::LOG_TARGET, "failed to decoded {:?}", info,);
-					"failed to decode value under existing key"
-				})?;
-				Ok::<usize, &'static str>(len)
-			},
-		}
+	let decode_key = |key: &[u8]| 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(|_| {
+				log::error!(target: crate::LOG_TARGET, "failed to decoded {:?}", info,);
+				"failed to decode value under existing key"
+			})?;
+			Ok::<usize, &'static str>(len)
+		},
 	};
 
 	decoded += decode_key(&next_key)?;
@@ -105,7 +101,7 @@ fn decode_storage_info<V: Decode>(info: StorageInfo) -> Result<usize, &'static s
 				decoded += decode_key(&key)?;
 				next_key = key;
 			},
-			_ => break
+			_ => break,
 		}
 	}
 
@@ -129,8 +125,8 @@ where
 	}
 }
 
-impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty> TryDecodeEntireStorage
-	for crate::storage::types::StorageMap<Prefix, Hasher, Key, Value, QueryKind, OnEmpty>
+impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireStorage
+	for crate::storage::types::StorageMap<Prefix, Hasher, Key, Value, QueryKind, OnEmpty, MaxValues>
 where
 	Prefix: StorageInstance,
 	Hasher: StorageHasher,
@@ -138,6 +134,7 @@ where
 	Value: FullCodec,
 	QueryKind: QueryKindTrait<Value, OnEmpty>,
 	OnEmpty: Get<QueryKind::Query> + 'static,
+	MaxValues: Get<Option<u32>>,
 {
 	fn try_decode_entire_state() -> Result<usize, &'static str> {
 		let info = Self::partial_storage_info()
@@ -148,15 +145,23 @@ where
 	}
 }
 
-impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty> TryDecodeEntireStorage
-	for crate::storage::types::CountedStorageMap<Prefix, Hasher, Key, Value, QueryKind, OnEmpty>
-where
+impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireStorage
+	for crate::storage::types::CountedStorageMap<
+		Prefix,
+		Hasher,
+		Key,
+		Value,
+		QueryKind,
+		OnEmpty,
+		MaxValues,
+	> where
 	Prefix: CountedStorageMapInstance,
 	Hasher: StorageHasher,
 	Key: FullCodec,
 	Value: FullCodec,
 	QueryKind: QueryKindTrait<Value, OnEmpty>,
 	OnEmpty: Get<QueryKind::Query> + 'static,
+	MaxValues: Get<Option<u32>>,
 {
 	fn try_decode_entire_state() -> Result<usize, &'static str> {
 		let (map_info, counter_info) = match &Self::partial_storage_info()[..] {
@@ -169,7 +174,8 @@ where
 	}
 }
 
-impl<Prefix, Hasher1, Key1, Hasher2, Key2, Value, QueryKind, OnEmpty> TryDecodeEntireStorage
+impl<Prefix, Hasher1, Key1, Hasher2, Key2, Value, QueryKind, OnEmpty, MaxValues>
+	TryDecodeEntireStorage
 	for crate::storage::types::StorageDoubleMap<
 		Prefix,
 		Hasher1,
@@ -179,6 +185,7 @@ impl<Prefix, Hasher1, Key1, Hasher2, Key2, Value, QueryKind, OnEmpty> TryDecodeE
 		Value,
 		QueryKind,
 		OnEmpty,
+		MaxValues,
 	> where
 	Prefix: StorageInstance,
 	Hasher1: StorageHasher,
@@ -188,6 +195,7 @@ impl<Prefix, Hasher1, Key1, Hasher2, Key2, Value, QueryKind, OnEmpty> TryDecodeE
 	Value: FullCodec,
 	QueryKind: QueryKindTrait<Value, OnEmpty>,
 	OnEmpty: Get<QueryKind::Query> + 'static,
+	MaxValues: Get<Option<u32>>,
 {
 	fn try_decode_entire_state() -> Result<usize, &'static str> {
 		let info = Self::partial_storage_info()
@@ -198,14 +206,15 @@ impl<Prefix, Hasher1, Key1, Hasher2, Key2, Value, QueryKind, OnEmpty> TryDecodeE
 	}
 }
 
-impl<Prefix, Key, Value, QueryKind, OnEmpty> TryDecodeEntireStorage
-	for crate::storage::types::StorageNMap<Prefix, Key, Value, QueryKind, OnEmpty>
+impl<Prefix, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireStorage
+	for crate::storage::types::StorageNMap<Prefix, Key, Value, QueryKind, OnEmpty, MaxValues>
 where
 	Prefix: StorageInstance,
 	Key: KeyGenerator,
 	Value: FullCodec,
 	QueryKind: QueryKindTrait<Value, OnEmpty>,
 	OnEmpty: Get<QueryKind::Query> + 'static,
+	MaxValues: Get<Option<u32>>,
 {
 	fn try_decode_entire_state() -> Result<usize, &'static str> {
 		let info = Self::partial_storage_info()
@@ -356,8 +365,7 @@ impl<BlockNumber: Clone + sp_std::fmt::Debug + AtLeast32BitUnsigned> TryState<Bl
 #[cfg(test)]
 mod tests {
 	use super::*;
-	use crate::Blake2_128Concat;
-	use crate::storage::types;
+	use crate::{storage::types, Blake2_128Concat};
 
 	// TODO reuse?
 	macro_rules! build_prefix {
@@ -369,7 +377,7 @@ mod tests {
 				}
 				const STORAGE_PREFIX: &'static str = stringify!($name);
 			}
-		}
+		};
 	}
 
 	build_prefix!(ValuePrefix);
@@ -450,7 +458,7 @@ mod tests {
 			assert_eq!(CMap::try_decode_entire_state(), Ok(4 + counter));
 
 			// counter is cleared again.
-			CMap::clear(u32::MAX, None).unwrap();
+			let _ = CMap::clear(u32::MAX, None);
 			assert_eq!(CMap::try_decode_entire_state(), Ok(0 + 0));
 
 			// two bytes, cannot be decoded into u32.

From d4d880bd85328c4dbdebe4f07d3c296b0e594a4f Mon Sep 17 00:00:00 2001
From: kianenigma <kian@parity.io>
Date: Fri, 17 Mar 2023 13:17:43 -0600
Subject: [PATCH 05/10] doc

---
 frame/support/src/storage/migration.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/frame/support/src/storage/migration.rs b/frame/support/src/storage/migration.rs
index 13816c157619e..59e47baa5fd22 100644
--- a/frame/support/src/storage/migration.rs
+++ b/frame/support/src/storage/migration.rs
@@ -389,7 +389,7 @@ pub fn move_prefix(from_prefix: &[u8], to_prefix: &[u8]) {
 	}
 }
 
-/// A phony migration that does nothing, except executing [`TryDecodeEntireStorage`] on
+/// A phony migration that does nothing, except executing `TryDecodeEntireStorage` on
 /// `post_upgrade`, which implies it is only available if `try-state` feature is used.
 ///
 /// This can be used typically in the top level runtime, where `AllPallets` typically comes from

From 15c5dc5b429bfdc6aa2da73d5ba4dfbb6d75baf0 Mon Sep 17 00:00:00 2001
From: Kian Paimani <5588131+kianenigma@users.noreply.github.com>
Date: Sun, 13 Aug 2023 12:29:08 -0700
Subject: [PATCH 06/10] update

---
 frame/support/src/traits/try_runtime.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/frame/support/src/traits/try_runtime.rs b/frame/support/src/traits/try_runtime.rs
index c361de5aebd68..390c782d61cdd 100644
--- a/frame/support/src/traits/try_runtime.rs
+++ b/frame/support/src/traits/try_runtime.rs
@@ -19,7 +19,7 @@
 
 use super::StorageInstance;
 use crate::{
-	storage::types::{CountedStorageMapInstance, KeyGenerator, QueryKindTrait},
+	storage::types::{CountedStorageMapInstance, KeyGenerator, QueryKindTrait, Counter},
 	traits::{PartialStorageInfoTrait, StorageInfo},
 	StorageHasher,
 };
@@ -169,7 +169,7 @@ impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireS
 			[a, b] => (a.clone(), b.clone()), // TODO better way?
 			_ => panic!("Counted map has two storage info items; qed"),
 		};
-		let mut decoded = decode_storage_info::<Value>(counter_info)?;
+		let mut decoded = decode_storage_info::<Counter>(counter_info)?;
 		decoded += decode_storage_info::<Value>(map_info)?;
 		Ok(decoded)
 	}

From 734453a1bb8bb0adc3a92f641fef5f89f1c0ef63 Mon Sep 17 00:00:00 2001
From: kianenigma <kian@parity.io>
Date: Sun, 13 Aug 2023 12:55:54 -0700
Subject: [PATCH 07/10] fix

---
 frame/support/src/storage/migration.rs        |   3 +-
 .../support/src/storage/types/counted_nmap.rs |   6 +-
 frame/support/src/storage/types/mod.rs        |   2 +-
 frame/support/src/traits/try_runtime.rs       | 143 +++++++++++++++---
 4 files changed, 133 insertions(+), 21 deletions(-)

diff --git a/frame/support/src/storage/migration.rs b/frame/support/src/storage/migration.rs
index f54dd21a4ea0a..1c921172b2086 100644
--- a/frame/support/src/storage/migration.rs
+++ b/frame/support/src/storage/migration.rs
@@ -23,6 +23,7 @@ use crate::{
 	StorageHasher, Twox128,
 };
 use codec::{Decode, Encode};
+use sp_runtime::TryRuntimeError;
 use sp_std::prelude::*;
 
 use super::PrefixIterator;
@@ -407,7 +408,7 @@ impl<AllPallets: crate::traits::TryDecodeEntireStorage> crate::traits::OnRuntime
 		Default::default()
 	}
 
-	fn post_upgrade(_: Vec<u8>) -> Result<(), &'static str> {
+	fn post_upgrade(_: Vec<u8>) -> Result<(), TryRuntimeError> {
 		let decoded = AllPallets::try_decode_entire_state()?;
 		crate::log::info!(
 			target: crate::LOG_TARGET,
diff --git a/frame/support/src/storage/types/counted_nmap.rs b/frame/support/src/storage/types/counted_nmap.rs
index 43a243cc5d681..03d162a637259 100644
--- a/frame/support/src/storage/types/counted_nmap.rs
+++ b/frame/support/src/storage/types/counted_nmap.rs
@@ -71,8 +71,10 @@ impl<P: CountedStorageNMapInstance, K, V, Q, O, M> MapWrapper
 	type Map = StorageNMap<P, K, V, Q, O, M>;
 }
 
+type Counter = super::counted_map::Counter;
+
 type CounterFor<P> =
-	StorageValue<<P as CountedStorageNMapInstance>::CounterPrefix, u32, ValueQuery>;
+	StorageValue<<P as CountedStorageNMapInstance>::CounterPrefix, Counter, ValueQuery>;
 
 /// On removal logic for updating counter while draining upon some prefix with
 /// [`crate::storage::PrefixIterator`].
@@ -429,7 +431,7 @@ where
 	}
 
 	/// Return the count.
-	pub fn count() -> u32 {
+	pub fn count() -> Counter {
 		CounterFor::<Prefix>::get()
 	}
 }
diff --git a/frame/support/src/storage/types/mod.rs b/frame/support/src/storage/types/mod.rs
index 99b0455fc6748..44dd403c1e6f9 100644
--- a/frame/support/src/storage/types/mod.rs
+++ b/frame/support/src/storage/types/mod.rs
@@ -30,7 +30,7 @@ mod map;
 mod nmap;
 mod value;
 
-pub use counted_map::{CountedStorageMap, CountedStorageMapInstance};
+pub use counted_map::{CountedStorageMap, CountedStorageMapInstance, Counter};
 pub use counted_nmap::{CountedStorageNMap, CountedStorageNMapInstance};
 pub use double_map::StorageDoubleMap;
 pub use key::{
diff --git a/frame/support/src/traits/try_runtime.rs b/frame/support/src/traits/try_runtime.rs
index 390c782d61cdd..19d09ce344456 100644
--- a/frame/support/src/traits/try_runtime.rs
+++ b/frame/support/src/traits/try_runtime.rs
@@ -19,11 +19,14 @@
 
 use super::StorageInstance;
 use crate::{
-	storage::types::{CountedStorageMapInstance, KeyGenerator, QueryKindTrait, Counter},
+	storage::types::{
+		CountedStorageMapInstance, CountedStorageNMapInstance, Counter, KeyGenerator,
+		QueryKindTrait,
+	},
 	traits::{PartialStorageInfoTrait, StorageInfo},
 	StorageHasher,
 };
-use codec::{Decode, FullCodec};
+use codec::{Decode, DecodeAll, FullCodec};
 use impl_trait_for_tuples::impl_for_tuples;
 use sp_arithmetic::traits::AtLeast32BitUnsigned;
 use sp_core::Get;
@@ -87,7 +90,7 @@ fn decode_storage_info<V: Decode>(info: StorageInfo) -> Result<usize, &'static s
 		None => Ok(0),
 		Some(bytes) => {
 			let len = bytes.len();
-			let _ = <V as Decode>::decode(&mut bytes.as_ref()).map_err(|_| {
+			let _ = <V as DecodeAll>::decode_all(&mut bytes.as_ref()).map_err(|_| {
 				log::error!(target: crate::LOG_TARGET, "failed to decoded {:?}", info,);
 				"failed to decode value under existing key"
 			})?;
@@ -166,7 +169,7 @@ impl<Prefix, Hasher, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireS
 {
 	fn try_decode_entire_state() -> Result<usize, &'static str> {
 		let (map_info, counter_info) = match &Self::partial_storage_info()[..] {
-			[a, b] => (a.clone(), b.clone()), // TODO better way?
+			[a, b] => (a.clone(), b.clone()),
 			_ => panic!("Counted map has two storage info items; qed"),
 		};
 		let mut decoded = decode_storage_info::<Counter>(counter_info)?;
@@ -220,12 +223,34 @@ where
 	fn try_decode_entire_state() -> Result<usize, &'static str> {
 		let info = Self::partial_storage_info()
 			.first()
-			.cloned() // TODO better way?
+			.cloned()
 			.expect("N-map has only one storage info; qed");
 		decode_storage_info::<Value>(info)
 	}
 }
 
+impl<Prefix, Key, Value, QueryKind, OnEmpty, MaxValues> TryDecodeEntireStorage
+	for crate::storage::types::CountedStorageNMap<Prefix, Key, Value, QueryKind, OnEmpty, MaxValues>
+where
+	Prefix: CountedStorageNMapInstance,
+	Key: KeyGenerator,
+	Value: FullCodec,
+	QueryKind: QueryKindTrait<Value, OnEmpty>,
+	OnEmpty: Get<QueryKind::Query> + 'static,
+	MaxValues: Get<Option<u32>>,
+{
+	fn try_decode_entire_state() -> Result<usize, &'static str> {
+		let (map_info, counter_info) = match &Self::partial_storage_info()[..] {
+			[a, b] => (a.clone(), b.clone()),
+			_ => panic!("Counted NMap has two storage info items; qed"),
+		};
+
+		let mut decoded = decode_storage_info::<Counter>(counter_info)?;
+		decoded += decode_storage_info::<Value>(map_info)?;
+		Ok(decoded)
+	}
+}
+
 impl sp_std::fmt::Debug for Select {
 	fn fmt(&self, f: &mut sp_std::fmt::Formatter<'_>) -> sp_std::fmt::Result {
 		match self {
@@ -374,9 +399,13 @@ impl<BlockNumber: Clone + sp_std::fmt::Debug + AtLeast32BitUnsigned> TryState<Bl
 #[cfg(test)]
 mod tests {
 	use super::*;
-	use crate::{storage::types, Blake2_128Concat};
+	use crate::{
+		storage::types::{self, Key},
+		Blake2_128Concat,
+	};
+
+	type H = Blake2_128Concat;
 
-	// TODO reuse?
 	macro_rules! build_prefix {
 		($name:ident) => {
 			struct $name;
@@ -393,18 +422,27 @@ mod tests {
 	type Value = types::StorageValue<ValuePrefix, u32>;
 
 	build_prefix!(MapPrefix);
-	type Map = types::StorageMap<MapPrefix, Blake2_128Concat, u32, u32>;
+	type Map = types::StorageMap<MapPrefix, H, u32, u32>;
 
 	build_prefix!(CMapCounterPrefix);
 	build_prefix!(CMapPrefix);
 	impl CountedStorageMapInstance for CMapPrefix {
 		type CounterPrefix = CMapCounterPrefix;
 	}
-	type CMap = types::CountedStorageMap<CMapPrefix, Blake2_128Concat, u32, u32>;
+	type CMap = types::CountedStorageMap<CMapPrefix, H, u8, u16>;
 
 	build_prefix!(DMapPrefix);
-	type DMap =
-		types::StorageDoubleMap<DMapPrefix, Blake2_128Concat, u32, Blake2_128Concat, u32, u32>;
+	type DMap = types::StorageDoubleMap<DMapPrefix, H, u32, H, u32, u32>;
+
+	build_prefix!(NMapPrefix);
+	type NMap = types::StorageNMap<NMapPrefix, (Key<H, u8>, Key<H, u8>), u128>;
+
+	build_prefix!(CountedNMapCounterPrefix);
+	build_prefix!(CountedNMapPrefix);
+	impl CountedStorageNMapInstance for CountedNMapPrefix {
+		type CounterPrefix = CountedNMapCounterPrefix;
+	}
+	type CNMap = types::CountedStorageNMap<CountedNMapPrefix, (Key<H, u8>, Key<H, u8>), u128>;
 
 	#[test]
 	fn try_decode_entire_state_value_works() {
@@ -453,25 +491,26 @@ mod tests {
 			assert_eq!(CMap::try_decode_entire_state(), Ok(0 + 0));
 
 			let counter = 4;
+			let value_size = std::mem::size_of::<u16>();
 
 			CMap::insert(0, 42);
-			assert_eq!(CMap::try_decode_entire_state(), Ok(4 + counter));
+			assert_eq!(CMap::try_decode_entire_state(), Ok(value_size + counter));
 
 			CMap::insert(0, 42);
-			assert_eq!(CMap::try_decode_entire_state(), Ok(4 + counter));
+			assert_eq!(CMap::try_decode_entire_state(), Ok(value_size + counter));
 
 			CMap::insert(1, 42);
-			assert_eq!(CMap::try_decode_entire_state(), Ok(8 + counter));
+			assert_eq!(CMap::try_decode_entire_state(), Ok(value_size * 2 + counter));
 
 			CMap::remove(0);
-			assert_eq!(CMap::try_decode_entire_state(), Ok(4 + counter));
+			assert_eq!(CMap::try_decode_entire_state(), Ok(value_size + counter));
 
 			// counter is cleared again.
 			let _ = CMap::clear(u32::MAX, None);
 			assert_eq!(CMap::try_decode_entire_state(), Ok(0 + 0));
 
-			// two bytes, cannot be decoded into u32.
-			sp_io::storage::set(&CMap::hashed_key_for(2), &[0u8, 1]);
+			// 1 bytes, cannot be decoded into u16.
+			sp_io::storage::set(&CMap::hashed_key_for(2), &[0u8]);
 			assert!(CMap::try_decode_entire_state().is_err());
 		})
 	}
@@ -502,6 +541,76 @@ mod tests {
 		})
 	}
 
+	#[test]
+	fn try_decode_entire_state_n_map_works() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			assert_eq!(NMap::try_decode_entire_state(), Ok(0));
+
+			let value_size = std::mem::size_of::<u128>();
+
+			NMap::insert((0u8, 0), 42);
+			assert_eq!(NMap::try_decode_entire_state(), Ok(value_size));
+
+			NMap::insert((0, 0), 42);
+			assert_eq!(NMap::try_decode_entire_state(), Ok(value_size));
+
+			NMap::insert((0, 1), 42);
+			assert_eq!(NMap::try_decode_entire_state(), Ok(value_size * 2));
+
+			NMap::insert((1, 0), 42);
+			assert_eq!(NMap::try_decode_entire_state(), Ok(value_size * 3));
+
+			NMap::remove((0, 0));
+			assert_eq!(NMap::try_decode_entire_state(), Ok(value_size * 2));
+
+			// two bytes, cannot be decoded into u128.
+			sp_io::storage::set(&NMap::hashed_key_for((1, 1)), &[0u8, 1]);
+			assert!(NMap::try_decode_entire_state().is_err());
+		})
+	}
+
+	#[test]
+	fn try_decode_entire_state_counted_n_map_works() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			sp_io::TestExternalities::new_empty().execute_with(|| {
+				assert_eq!(NMap::try_decode_entire_state(), Ok(0));
+
+				let value_size = std::mem::size_of::<u128>();
+				let counter = 4;
+
+				CNMap::insert((0u8, 0), 42);
+				assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size + counter));
+
+				CNMap::insert((0, 0), 42);
+				assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size + counter));
+
+				CNMap::insert((0, 1), 42);
+				assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size * 2 + counter));
+
+				CNMap::insert((1, 0), 42);
+				assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size * 3 + counter));
+
+				CNMap::remove((0, 0));
+				assert_eq!(CNMap::try_decode_entire_state(), Ok(value_size * 2 + counter));
+
+				// two bytes, cannot be decoded into u128.
+				sp_io::storage::set(&CNMap::hashed_key_for((1, 1)), &[0u8, 1]);
+				assert!(CNMap::try_decode_entire_state().is_err());
+			})
+		})
+	}
+
+	#[test]
+	fn extra_bytes_are_rejected() {
+		sp_io::TestExternalities::new_empty().execute_with(|| {
+			assert_eq!(Map::try_decode_entire_state(), Ok(0));
+
+			// 6bytes, too many to fit in u32, should be rejected.
+			sp_io::storage::set(&Map::hashed_key_for(2), &[0u8, 1, 3, 4, 5, 6]);
+			assert!(Map::try_decode_entire_state().is_err());
+		})
+	}
+
 	#[test]
 	fn try_decode_entire_state_tuple_of_storage_works() {
 		sp_io::TestExternalities::new_empty().execute_with(|| {

From 375b71bf32a44cb416b46689344326b61c785235 Mon Sep 17 00:00:00 2001
From: kianenigma <kian@parity.io>
Date: Sun, 13 Aug 2023 13:06:11 -0700
Subject: [PATCH 08/10] fux

---
 frame/support/src/storage/migration.rs | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/frame/support/src/storage/migration.rs b/frame/support/src/storage/migration.rs
index 1c921172b2086..e581c8a00a21e 100644
--- a/frame/support/src/storage/migration.rs
+++ b/frame/support/src/storage/migration.rs
@@ -23,7 +23,6 @@ use crate::{
 	StorageHasher, Twox128,
 };
 use codec::{Decode, Encode};
-use sp_runtime::TryRuntimeError;
 use sp_std::prelude::*;
 
 use super::PrefixIterator;
@@ -408,7 +407,7 @@ impl<AllPallets: crate::traits::TryDecodeEntireStorage> crate::traits::OnRuntime
 		Default::default()
 	}
 
-	fn post_upgrade(_: Vec<u8>) -> Result<(), TryRuntimeError> {
+	fn post_upgrade(_: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
 		let decoded = AllPallets::try_decode_entire_state()?;
 		crate::log::info!(
 			target: crate::LOG_TARGET,

From 07aac4ad66fdce9255f5d3fcb3a250d1f70e209b Mon Sep 17 00:00:00 2001
From: kianenigma <kian@parity.io>
Date: Sun, 13 Aug 2023 17:36:12 -0700
Subject: [PATCH 09/10] fix

---
 .../support/procedural/src/pallet/expand/event.rs |  3 ++-
 .../procedural/src/pallet/expand/storage.rs       | 15 +++++++++++----
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/frame/support/procedural/src/pallet/expand/event.rs b/frame/support/procedural/src/pallet/expand/event.rs
index f94bdef332d9d..49e9344c6a9a8 100644
--- a/frame/support/procedural/src/pallet/expand/event.rs
+++ b/frame/support/procedural/src/pallet/expand/event.rs
@@ -127,11 +127,12 @@ pub fn expand_event(def: &mut Def) -> proc_macro2::TokenStream {
 		let trait_use_gen = &def.trait_use_generics(event.attr_span);
 		let type_impl_gen = &def.type_impl_generics(event.attr_span);
 		let type_use_gen = &def.type_use_generics(event.attr_span);
+		let pallet_ident = &def.pallet_struct.pallet;
 
 		let PalletEventDepositAttr { fn_vis, fn_span, .. } = deposit_event;
 
 		quote::quote_spanned!(*fn_span =>
-			impl<#type_impl_gen> Pallet<#type_use_gen> #completed_where_clause {
+			impl<#type_impl_gen> #pallet_ident<#type_use_gen> #completed_where_clause {
 				#fn_vis fn deposit_event(event: Event<#event_use_gen>) {
 					let event = <
 						<T as Config #trait_use_gen>::RuntimeEvent as
diff --git a/frame/support/procedural/src/pallet/expand/storage.rs b/frame/support/procedural/src/pallet/expand/storage.rs
index 1508f5abc95ba..3c3c5b51b5272 100644
--- a/frame/support/procedural/src/pallet/expand/storage.rs
+++ b/frame/support/procedural/src/pallet/expand/storage.rs
@@ -791,12 +791,17 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 		let storage_names = def
 			.storages
 			.iter()
-			.map(|storage| {
-				let ident = &storage.ident;
-				let gen = &def.type_use_generics(storage.attr_span);
-				quote::quote_spanned!(storage.attr_span => #ident<#gen> )
+			.filter_map(|storage| {
+				if storage.cfg_attrs.is_empty() {
+					let ident = &storage.ident;
+					let gen = &def.type_use_generics(storage.attr_span);
+					Some(quote::quote_spanned!(storage.attr_span => #ident<#gen> ))
+				} else {
+					None
+				}
 			})
 			.collect::<Vec<_>>();
+
 		quote::quote!(
 			#[cfg(feature = "try-runtime")]
 			impl<#type_impl_gen> #frame_support::traits::TryDecodeEntireStorage
@@ -804,6 +809,8 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 			{
 				fn try_decode_entire_state() -> Result<usize, &'static str> {
 					// simply delegate impl to a tuple of all storage items we have.
+					//
+					// NOTE: for now, we have to exclude storage items that are feature gated.
 					<( #( #storage_names ),*) as frame_support::traits::TryDecodeEntireStorage>::try_decode_entire_state()
 				}
 			}

From ed7b61860bb0c20bd6c67e1c46ec5cb3bafe005c Mon Sep 17 00:00:00 2001
From: kianenigma <kian@parity.io>
Date: Wed, 23 Aug 2023 13:32:17 +0200
Subject: [PATCH 10/10] try and fix ui build

---
 .../test/tests/pallet_ui/call_argument_invalid_bound.stderr | 2 +-
 .../tests/pallet_ui/call_argument_invalid_bound_2.stderr    | 2 +-
 .../tests/pallet_ui/call_argument_invalid_bound_3.stderr    | 2 +-
 .../test/tests/pallet_ui/event_field_not_member.stderr      | 2 +-
 .../test/tests/pallet_ui/inherent_check_inner_span.stderr   | 6 +++---
 .../tests/pallet_ui/storage_info_unsatisfied_nmap.stderr    | 4 ++--
 6 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/frame/support/test/tests/pallet_ui/call_argument_invalid_bound.stderr b/frame/support/test/tests/pallet_ui/call_argument_invalid_bound.stderr
index d10bf1359019a..cc8c4fda76f2f 100644
--- a/frame/support/test/tests/pallet_ui/call_argument_invalid_bound.stderr
+++ b/frame/support/test/tests/pallet_ui/call_argument_invalid_bound.stderr
@@ -19,7 +19,7 @@ error[E0277]: `<T as pallet::Config>::Bar` doesn't implement `std::fmt::Debug`
    |
    = help: the trait `std::fmt::Debug` is not implemented for `<T as pallet::Config>::Bar`
    = note: required for `&<T as pallet::Config>::Bar` to implement `std::fmt::Debug`
-   = note: required for the cast from `&<T as pallet::Config>::Bar` to the object type `dyn std::fmt::Debug`
+   = note: required for the cast from `&&<T as pallet::Config>::Bar` to `&dyn std::fmt::Debug`
 
 error[E0277]: the trait bound `<T as pallet::Config>::Bar: Clone` is not satisfied
   --> tests/pallet_ui/call_argument_invalid_bound.rs:21:36
diff --git a/frame/support/test/tests/pallet_ui/call_argument_invalid_bound_2.stderr b/frame/support/test/tests/pallet_ui/call_argument_invalid_bound_2.stderr
index 7173cdcd47361..b8d0bbf34730e 100644
--- a/frame/support/test/tests/pallet_ui/call_argument_invalid_bound_2.stderr
+++ b/frame/support/test/tests/pallet_ui/call_argument_invalid_bound_2.stderr
@@ -19,7 +19,7 @@ error[E0277]: `<T as pallet::Config>::Bar` doesn't implement `std::fmt::Debug`
    |
    = help: the trait `std::fmt::Debug` is not implemented for `<T as pallet::Config>::Bar`
    = note: required for `&<T as pallet::Config>::Bar` to implement `std::fmt::Debug`
-   = note: required for the cast from `&<T as pallet::Config>::Bar` to the object type `dyn std::fmt::Debug`
+   = note: required for the cast from `&&<T as pallet::Config>::Bar` to `&dyn std::fmt::Debug`
 
 error[E0277]: the trait bound `<T as pallet::Config>::Bar: Clone` is not satisfied
   --> tests/pallet_ui/call_argument_invalid_bound_2.rs:21:36
diff --git a/frame/support/test/tests/pallet_ui/call_argument_invalid_bound_3.stderr b/frame/support/test/tests/pallet_ui/call_argument_invalid_bound_3.stderr
index 4cbed3709626c..158d0c7638e31 100644
--- a/frame/support/test/tests/pallet_ui/call_argument_invalid_bound_3.stderr
+++ b/frame/support/test/tests/pallet_ui/call_argument_invalid_bound_3.stderr
@@ -20,7 +20,7 @@ error[E0277]: `Bar` doesn't implement `std::fmt::Debug`
    = help: the trait `std::fmt::Debug` is not implemented for `Bar`
    = note: add `#[derive(Debug)]` to `Bar` or manually `impl std::fmt::Debug for Bar`
    = note: required for `&Bar` to implement `std::fmt::Debug`
-   = note: required for the cast from `&Bar` to the object type `dyn std::fmt::Debug`
+   = note: required for the cast from `&&Bar` to `&dyn std::fmt::Debug`
 help: consider annotating `Bar` with `#[derive(Debug)]`
    |
 17 +     #[derive(Debug)]
diff --git a/frame/support/test/tests/pallet_ui/event_field_not_member.stderr b/frame/support/test/tests/pallet_ui/event_field_not_member.stderr
index 1161f4a190231..d853406c95817 100644
--- a/frame/support/test/tests/pallet_ui/event_field_not_member.stderr
+++ b/frame/support/test/tests/pallet_ui/event_field_not_member.stderr
@@ -18,4 +18,4 @@ error[E0277]: `<T as pallet::Config>::Bar` doesn't implement `std::fmt::Debug`
    |
    = help: the trait `std::fmt::Debug` is not implemented for `<T as pallet::Config>::Bar`
    = note: required for `&<T as pallet::Config>::Bar` to implement `std::fmt::Debug`
-   = note: required for the cast from `&<T as pallet::Config>::Bar` to the object type `dyn std::fmt::Debug`
+   = note: required for the cast from `&&<T as pallet::Config>::Bar` to `&dyn std::fmt::Debug`
diff --git a/frame/support/test/tests/pallet_ui/inherent_check_inner_span.stderr b/frame/support/test/tests/pallet_ui/inherent_check_inner_span.stderr
index bc34c55241a76..4acf57fb3e3dc 100644
--- a/frame/support/test/tests/pallet_ui/inherent_check_inner_span.stderr
+++ b/frame/support/test/tests/pallet_ui/inherent_check_inner_span.stderr
@@ -1,11 +1,11 @@
 error[E0046]: not all trait items implemented, missing: `Call`, `Error`, `INHERENT_IDENTIFIER`, `create_inherent`, `is_inherent`
-  --> $DIR/inherent_check_inner_span.rs:19:2
+  --> tests/pallet_ui/inherent_check_inner_span.rs:19:2
    |
 19 |     impl<T: Config> ProvideInherent for Pallet<T> {}
    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ missing `Call`, `Error`, `INHERENT_IDENTIFIER`, `create_inherent`, `is_inherent` in implementation
    |
-   = help: implement the missing item: `type Call = Type;`
-   = help: implement the missing item: `type Error = Type;`
+   = help: implement the missing item: `type Call = /* Type */;`
+   = help: implement the missing item: `type Error = /* Type */;`
    = help: implement the missing item: `const INHERENT_IDENTIFIER: [u8; 8] = value;`
    = help: implement the missing item: `fn create_inherent(_: &InherentData) -> std::option::Option<<Self as ProvideInherent>::Call> { todo!() }`
    = help: implement the missing item: `fn is_inherent(_: &<Self as ProvideInherent>::Call) -> bool { todo!() }`
diff --git a/frame/support/test/tests/pallet_ui/storage_info_unsatisfied_nmap.stderr b/frame/support/test/tests/pallet_ui/storage_info_unsatisfied_nmap.stderr
index c34c796fe59c1..bbbb3ed1b2e47 100644
--- a/frame/support/test/tests/pallet_ui/storage_info_unsatisfied_nmap.stderr
+++ b/frame/support/test/tests/pallet_ui/storage_info_unsatisfied_nmap.stderr
@@ -14,5 +14,5 @@ error[E0277]: the trait bound `Bar: MaxEncodedLen` is not satisfied
              (TupleElement0, TupleElement1, TupleElement2, TupleElement3, TupleElement4, TupleElement5, TupleElement6)
              (TupleElement0, TupleElement1, TupleElement2, TupleElement3, TupleElement4, TupleElement5, TupleElement6, TupleElement7)
            and $N others
-   = note: required for `Key<frame_support::Twox64Concat, Bar>` to implement `KeyGeneratorMaxEncodedLen`
-   = note: required for `frame_support::pallet_prelude::StorageNMap<_GeneratedPrefixForStorageFoo<T>, Key<frame_support::Twox64Concat, Bar>, u32>` to implement `StorageInfoTrait`
+   = note: required for `NMapKey<frame_support::Twox64Concat, Bar>` to implement `KeyGeneratorMaxEncodedLen`
+   = note: required for `frame_support::pallet_prelude::StorageNMap<_GeneratedPrefixForStorageFoo<T>, NMapKey<frame_support::Twox64Concat, Bar>, u32>` to implement `StorageInfoTrait`