Add iteration with key information when using $hash_concat

frame/support/procedural/src/storage/storage_struct.rs

@@ -121,10 +121,25 @@ pub fn decl_and_impl(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStre
 			},
 			StorageLineTypeDef::Map(map) => {
 				let hasher = map.hasher.to_storage_hasher_struct();
+				let key_type = &map.key;
+				let hasher_info = quote!(
+					<
+						<Self as #scrate::#storage_generator_trait>::Hasher
+						as #scrate::StorageHasherInfo<#key_type>
+					>
+				);
 				quote!(
 					impl<#impl_trait> #scrate::storage::StoragePrefixedMap<#value_type>
 						for #storage_struct #optional_storage_where_clause
 					{
+						type KeyInfo = #hasher_info::Info;
+
+						fn decode_key_info_from_encoded_key_without_prefix(encoded_key: &[u8])
+							-> Result<Self::KeyInfo, codec::Error>
+						{
+							#hasher_info::decode_hash_and_then_info(&mut &encoded_key[..])
+						}
+
 						fn module_prefix() -> &'static [u8] {
 							#instance_or_inherent::PREFIX.as_bytes()
 						}
@@ -201,11 +216,42 @@ pub fn decl_and_impl(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStre
 			},
 			StorageLineTypeDef::DoubleMap(map) => {
 				let hasher1 = map.hasher1.to_storage_hasher_struct();
+				let key1_type = &map.key1;
+				let hasher1_info = quote!(
+					<
+						<Self as #scrate::#storage_generator_trait>::Hasher1
+						as #scrate::StorageHasherInfo<#key1_type>
+					>
+				);
+
 				let hasher2 = map.hasher2.to_storage_hasher_struct();
+				let key2_type = &map.key2;
+				let hasher2_info = quote!(
+					<
+						<Self as #scrate::#storage_generator_trait>::Hasher2
+						as #scrate::StorageHasherInfo<#key2_type>
+					>
+				);
+
 				quote!(
 					impl<#impl_trait> #scrate::storage::StoragePrefixedMap<#value_type>
 						for #storage_struct #optional_storage_where_clause
 					{
+						type KeyInfo = (
+							#hasher1_info::Info,
+							#hasher2_info::Info
+						);
+
+						fn decode_key_info_from_encoded_key_without_prefix(encoded_key: &[u8])
+							-> Result<Self::KeyInfo, codec::Error>
+						{
+							let mut input = encoded_key;
+							Ok((
+								#hasher1_info::decode_hash_and_then_info(&mut input)?,
+								#hasher2_info::decode_hash_and_then_info(&mut input)?
+							))
+						}
+
 						fn module_prefix() -> &'static [u8] {
 							#instance_or_inherent::PREFIX.as_bytes()
 						}

frame/support/src/hash.rs

@@ -16,7 +16,7 @@
 
 //! Hash utilities.
 
-use codec::Codec;
+use codec::{Codec, Decode};
 use sp_std::prelude::Vec;
 use sp_io::hashing::{blake2_128, blake2_256, twox_64, twox_128, twox_256};
 
@@ -57,6 +57,20 @@ pub trait StorageHasher: 'static {
 	fn hash(x: &[u8]) -> Self::Output;
 }
 
+/// Trait to retrieve some info from hash of type `Key` encoded.
+pub trait StorageHasherInfo<Key> {
+	/// Some info contained in the hash of type `Key` encoded.
+	type Info;
+
+	/// Decode the hash and then decode the info from the decoded hash.
+	///
+	/// # WARNING
+	///
+	/// Even if info is (), input must be modified to have read the entire encoded hash.
+	fn decode_hash_and_then_info<I: codec::Input>(input: &mut I)
+		-> Result<Self::Info, codec::Error>;
+}
+
 /// Hash storage keys with `concat(twox64(key), key)`
 pub struct Twox64Concat;
 impl StorageHasher for Twox64Concat {
@@ -70,6 +84,16 @@ impl StorageHasher for Twox64Concat {
 	}
 }
 
+impl<Key: Decode> StorageHasherInfo<Key> for Twox64Concat {
+	type Info = Key;
+	fn decode_hash_and_then_info<I: codec::Input>(input: &mut I)
+		-> Result<Self::Info, codec::Error>
+	{
+		input.read(&mut [0u8; 8])?;
+		Key::decode(input)
+	}
+}
+
 /// Hash storage keys with `concat(blake2_128(key), key)`
 pub struct Blake2_128Concat;
 impl StorageHasher for Blake2_128Concat {
@@ -83,6 +107,16 @@ impl StorageHasher for Blake2_128Concat {
 	}
 }
 
+impl<Key: Decode> StorageHasherInfo<Key> for Blake2_128Concat {
+	type Info = Key;
+	fn decode_hash_and_then_info<I: codec::Input>(input: &mut I)
+		-> Result<Self::Info, codec::Error>
+	{
+		input.read(&mut [0u8; 16])?;
+		Key::decode(input)
+	}
+}
+
 /// Hash storage keys with blake2 128
 pub struct Blake2_128;
 impl StorageHasher for Blake2_128 {
@@ -92,6 +126,16 @@ impl StorageHasher for Blake2_128 {
 	}
 }
 
+impl<Key> StorageHasherInfo<Key> for Blake2_128 {
+	type Info = ();
+	fn decode_hash_and_then_info<I: codec::Input>(input: &mut I)
+		-> Result<Self::Info, codec::Error>
+	{
+		input.read(&mut [0u8; 16])?;
+		Ok(())
+	}
+}
+
 /// Hash storage keys with blake2 256
 pub struct Blake2_256;
 impl StorageHasher for Blake2_256 {
@@ -101,6 +145,16 @@ impl StorageHasher for Blake2_256 {
 	}
 }
 
+impl<Key> StorageHasherInfo<Key> for Blake2_256 {
+	type Info = ();
+	fn decode_hash_and_then_info<I: codec::Input>(input: &mut I)
+		-> Result<Self::Info, codec::Error>
+	{
+		input.read(&mut [0u8; 32])?;
+		Ok(())
+	}
+}
+
 /// Hash storage keys with twox 128
 pub struct Twox128;
 impl StorageHasher for Twox128 {
@@ -110,6 +164,16 @@ impl StorageHasher for Twox128 {
 	}
 }
 
+impl<Key> StorageHasherInfo<Key> for Twox128 {
+	type Info = ();
+	fn decode_hash_and_then_info<I: codec::Input>(input: &mut I)
+		-> Result<Self::Info, codec::Error>
+	{
+		input.read(&mut [0u8; 16])?;
+		Ok(())
+	}
+}
+
 /// Hash storage keys with twox 256
 pub struct Twox256;
 impl StorageHasher for Twox256 {
@@ -119,9 +183,20 @@ impl StorageHasher for Twox256 {
 	}
 }
 
+impl<Key> StorageHasherInfo<Key> for Twox256 {
+	type Info = ();
+	fn decode_hash_and_then_info<I: codec::Input>(input: &mut I)
+		-> Result<Self::Info, codec::Error>
+	{
+		input.read(&mut [0u8; 32])?;
+		Ok(())
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;
+	use codec::Encode;
 
 	#[test]
 	fn test_twox_64_concat() {
@@ -134,4 +209,34 @@ mod tests {
 		let r = Blake2_128Concat::hash(b"foo");
 		assert_eq!(r.split_at(16), (&blake2_128(b"foo")[..], &b"foo"[..]))
 	}
+
+	#[test]
+	fn test_storage_hasher_info() {
+		type KeyType = [u8; 15];
+		let key: KeyType = [3u8; 15];
+
+		let mut r = &Twox64Concat::hash(&(&key).encode()[..])[..];
+		assert_eq!(<Twox64Concat as StorageHasherInfo<KeyType>>::decode_hash_and_then_info(&mut r), Ok(key));
+		assert_eq!(r.len(), 0); // Assert input has indeed decoded the hash.
+
+		let mut r = &Twox128::hash(&(&key).encode()[..])[..];
+		assert_eq!(<Twox128 as StorageHasherInfo<KeyType>>::decode_hash_and_then_info(&mut r), Ok(()));
+		assert_eq!(r.len(), 0); // Assert input has indeed decoded the hash.
+
+		let mut r = &Twox256::hash(&(&key).encode()[..])[..];
+		assert_eq!(<Twox256 as StorageHasherInfo<KeyType>>::decode_hash_and_then_info(&mut r), Ok(()));
+		assert_eq!(r.len(), 0); // Assert input has indeed decoded the hash.
+
+		let mut r = &Blake2_128Concat::hash(&(&key).encode()[..])[..];
+		assert_eq!(<Blake2_128Concat as StorageHasherInfo<KeyType>>::decode_hash_and_then_info(&mut r), Ok(key));
+		assert_eq!(r.len(), 0); // Assert input has indeed decoded the hash.
+
+		let mut r = &Blake2_128::hash(&(&key).encode()[..])[..];
+		assert_eq!(<Blake2_128 as StorageHasherInfo<KeyType>>::decode_hash_and_then_info(&mut r), Ok(()));
+		assert_eq!(r.len(), 0); // Assert input has indeed decoded the hash.
+
+		let mut r = &Blake2_256::hash(&(&key).encode()[..])[..];
+		assert_eq!(<Blake2_256 as StorageHasherInfo<KeyType>>::decode_hash_and_then_info(&mut r), Ok(()));
+		assert_eq!(r.len(), 0); // Assert input has indeed decoded the hash.
+	}
 }

frame/support/src/lib.rs

@@ -68,7 +68,7 @@ pub mod weights;
 
 pub use self::hash::{
 	Twox256, Twox128, Blake2_256, Blake2_128, Twox64Concat, Blake2_128Concat, Hashable,
-	StorageHasher
+	StorageHasher, StorageHasherInfo
 };
 pub use self::storage::{
 	StorageValue, StorageMap, StorageLinkedMap, StorageDoubleMap, StoragePrefixedMap