Implement StorageNMap (#8635)

* Implement StorageNMap

* Change copyright date to 2021

* Rewrite keys to use impl_for_tuples instead of recursion

* Implement prefix iteration on StorageNMap

* Implement EncodeLike for key arguments

* Rename KeyGenerator::Arg to KeyGenerator::KArg

* Support StorageNMap in decl_storage and #[pallet::storage] macros

* Use StorageNMap in assets pallet

* Support migrate_keys in StorageNMap

* Reduce line characters on select files

* Refactor crate imports in decl_storage macros

* Some more line char reductions and doc comment update

* Update UI test expectations

* Revert whitespace changes to untouched files

* Generate Key struct instead of a 1-tuple when only 1 pair of key and hasher is provided

* Revert formatting changes to unrelated files

* Introduce KeyGeneratorInner

* Add tests for StorageNMap in FRAMEv2 pallet macro

* Small fixes to unit tests for StorageNMap

* Bump runtime metadata version

* Remove unused import

* Update tests to use runtime metadata v13

* Introduce and use EncodeLikeTuple as a trait bound for KArg

* Add some rustdocs

* Revert usage of StorageNMap in assets pallet

* Make use of ext::PunctuatedTrailing

* Add rustdoc for final_hash

* Fix StorageNMap proc macro expansions for single key cases

* Create associated const in KeyGenerator for hasher metadata

* Refactor code according to comments from Basti

* Add module docs for generator/nmap.rs

* Re-export storage::Key as NMapKey in pallet prelude

* Seal the EncodeLikeTuple trait

* Extract sealing code out of key.rs

Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>

substrate/frame/support/procedural/src/storage/genesis_config/builder_def.rs

@@ -120,6 +120,21 @@ impl BuilderDef {
 							});
 						}}
 					},
+					StorageLineTypeDef::NMap(map) => {
+						let key_tuple = map.to_key_tuple();
+						let key_arg = if map.keys.len() == 1 {
+							quote!((k,))
+						} else {
+							quote!(k)
+						};
+						quote!{{
+							#data
+							let data: &#scrate::sp_std::vec::Vec<(#key_tuple, #value_type)> = data;
+							data.iter().for_each(|(k, v)| {
+								<#storage_struct as #scrate::#storage_trait>::insert(#key_arg, v);
+							});
+						}}
+					},
 				});
 			}
 		}

substrate/frame/support/procedural/src/storage/genesis_config/genesis_config_def.rs

@@ -104,6 +104,10 @@ impl GenesisConfigDef {
 
 					parse_quote!( Vec<(#key1, #key2, #value_type)> )
 				},
+				StorageLineTypeDef::NMap(map) => {
+					let key_tuple = map.to_key_tuple();
+					parse_quote!( Vec<(#key_tuple, #value_type)> )
+				}
 			};
 
 			let default = line.default_value.as_ref()

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

@@ -177,10 +177,8 @@ fn impl_build_storage(
 	}
 }
 
-pub fn genesis_config_and_build_storage(
-	scrate: &TokenStream,
-	def: &DeclStorageDefExt,
-) -> TokenStream {
+pub fn genesis_config_and_build_storage(def: &DeclStorageDefExt) -> TokenStream {
+	let scrate = &def.hidden_crate;
 	let builders = BuilderDef::from_def(scrate, def);
 	if !builders.blocks.is_empty() {
 		let genesis_config = match GenesisConfigDef::from_def(def) {

substrate/frame/support/procedural/src/storage/getters.rs

@@ -21,7 +21,8 @@ use proc_macro2::TokenStream;
 use quote::quote;
 use super::{DeclStorageDefExt, StorageLineTypeDef};
 
-pub fn impl_getters(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStream {
+pub fn impl_getters(def: &DeclStorageDefExt) -> TokenStream {
+	let scrate = &def.hidden_crate;
 	let mut getters = TokenStream::new();
 
 	for (get_fn, line) in def.storage_lines.iter()
@@ -65,6 +66,21 @@ pub fn impl_getters(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStrea
 					}
 				}
 			},
+			StorageLineTypeDef::NMap(map) => {
+				let keygen = map.to_keygen_struct(&def.hidden_crate);
+				let value = &map.value;
+				quote!{
+					pub fn #get_fn<KArg>(key: KArg) -> #value
+					where
+						KArg: #scrate::storage::types::EncodeLikeTuple<
+							<#keygen as #scrate::storage::types::KeyGenerator>::KArg
+						>
+							+ #scrate::storage::types::TupleToEncodedIter,
+					{
+						<#storage_struct as #scrate::#storage_trait>::get(key)
+					}
+				}
+			}
 		};
 		getters.extend(getter);
 	}

substrate/frame/support/procedural/src/storage/instance_trait.rs

@@ -34,7 +34,8 @@ struct InstanceDef {
 	index: u8,
 }
 
-pub fn decl_and_impl(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStream {
+pub fn decl_and_impl(def: &DeclStorageDefExt) -> TokenStream {
+	let scrate = &def.hidden_crate;
 	let mut impls = TokenStream::new();
 
 	impls.extend(reexport_instance_trait(scrate, def));

substrate/frame/support/procedural/src/storage/metadata.rs

@@ -63,6 +63,27 @@ fn storage_line_metadata_type(scrate: &TokenStream, line: &StorageLineDefExt) ->
 				}
 			}
 		},
+		StorageLineTypeDef::NMap(map) => {
+			let keys = map.keys
+				.iter()
+				.map(|key| clean_type_string(&quote!(#key).to_string()))
+				.collect::<Vec<_>>();
+			let hashers = map.hashers
+				.iter()
+				.map(|hasher| hasher.to_storage_hasher_struct())
+				.collect::<Vec<_>>();
+			quote!{
+				#scrate::metadata::StorageEntryType::NMap {
+					keys: #scrate::metadata::DecodeDifferent::Encode(&[
+						#( #keys, )*
+					]),
+					hashers: #scrate::metadata::DecodeDifferent::Encode(&[
+						#( #scrate::metadata::StorageHasher::#hashers, )*
+					]),
+					value: #scrate::metadata::DecodeDifferent::Encode(#value_type),
+				}
+			}
+		}
 	}
 }
 
@@ -140,7 +161,8 @@ fn default_byte_getter(
 	(struct_def, struct_instance)
 }
 
-pub fn impl_metadata(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStream {
+pub fn impl_metadata(def: &DeclStorageDefExt) -> TokenStream {
+	let scrate = &def.hidden_crate;
 	let mut entries = TokenStream::new();
 	let mut default_byte_getter_struct_defs = TokenStream::new();
 

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

@@ -70,7 +70,9 @@ impl syn::parse::Parse for DeclStorageDef {
 /// Extended version of `DeclStorageDef` with useful precomputed value.
 pub struct DeclStorageDefExt {
 	/// Name of the module used to import hidden imports.
-	hidden_crate: Option<syn::Ident>,
+	hidden_crate: proc_macro2::TokenStream,
+	/// Hidden imports used by the module.
+	hidden_imports: proc_macro2::TokenStream,
 	/// Visibility of store trait.
 	visibility: syn::Visibility,
 	/// Name of store trait: usually `Store`.
@@ -108,9 +110,15 @@ pub struct DeclStorageDefExt {
 
 impl From<DeclStorageDef> for DeclStorageDefExt {
 	fn from(mut def: DeclStorageDef) -> Self {
+		let hidden_crate_name = def.hidden_crate.as_ref().map(|i| i.to_string())
+			.unwrap_or_else(|| "decl_storage".to_string());
+	
+		let hidden_crate = generate_crate_access(&hidden_crate_name, "frame-support");
+		let hidden_imports = generate_hidden_includes(&hidden_crate_name, "frame-support");
+
 		let storage_lines = def.storage_lines.drain(..).collect::<Vec<_>>();
 		let storage_lines = storage_lines.into_iter()
-			.map(|line| StorageLineDefExt::from_def(line, &def))
+			.map(|line| StorageLineDefExt::from_def(line, &def, &hidden_crate))
 			.collect();
 
 		let (
@@ -144,7 +152,8 @@ impl From<DeclStorageDef> for DeclStorageDefExt {
 		);
 
 		Self {
-			hidden_crate: def.hidden_crate,
+			hidden_crate,
+			hidden_imports,
 			visibility: def.visibility,
 			store_trait: def.store_trait,
 			module_name: def.module_name,
@@ -230,7 +239,11 @@ pub struct StorageLineDefExt {
 }
 
 impl StorageLineDefExt {
-	fn from_def(storage_def: StorageLineDef, def: &DeclStorageDef) -> Self {
+	fn from_def(
+		storage_def: StorageLineDef,
+		def: &DeclStorageDef,
+		hidden_crate: &proc_macro2::TokenStream,
+	) -> Self {
 		let is_generic = match &storage_def.storage_type {
 			StorageLineTypeDef::Simple(value) => {
 				ext::type_contains_ident(&value, &def.module_runtime_generic)
@@ -244,12 +257,17 @@ impl StorageLineDefExt {
 					|| ext::type_contains_ident(&map.key2, &def.module_runtime_generic)
 					|| ext::type_contains_ident(&map.value, &def.module_runtime_generic)
 			}
+			StorageLineTypeDef::NMap(map) => {
+				map.keys.iter().any(|key| ext::type_contains_ident(key, &def.module_runtime_generic))
+					|| ext::type_contains_ident(&map.value, &def.module_runtime_generic)
+			}
 		};
 
 		let query_type = match &storage_def.storage_type {
 			StorageLineTypeDef::Simple(value) => value.clone(),
 			StorageLineTypeDef::Map(map) => map.value.clone(),
 			StorageLineTypeDef::DoubleMap(map) => map.value.clone(),
+			StorageLineTypeDef::NMap(map) => map.value.clone(),
 		};
 		let is_option = ext::extract_type_option(&query_type).is_some();
 		let value_type = ext::extract_type_option(&query_type).unwrap_or_else(|| query_type.clone());
@@ -295,6 +313,10 @@ impl StorageLineDefExt {
 				let key2 = &map.key2;
 				quote!( StorageDoubleMap<#key1, #key2, #value_type> )
 			},
+			StorageLineTypeDef::NMap(map) => {
+				let keygen = map.to_keygen_struct(hidden_crate);
+				quote!( StorageNMap<#keygen, #value_type> )
+			}
 		};
 
 		let storage_trait = quote!( storage::#storage_trait_truncated );
@@ -332,6 +354,7 @@ impl StorageLineDefExt {
 pub enum StorageLineTypeDef {
 	Map(MapDef),
 	DoubleMap(Box<DoubleMapDef>),
+	NMap(NMapDef),
 	Simple(syn::Type),
 }
 
@@ -351,6 +374,42 @@ pub struct DoubleMapDef {
 	pub value: syn::Type,
 }
 
+pub struct NMapDef {
+	pub hashers: Vec<HasherKind>,
+	pub keys: Vec<syn::Type>,
+	pub value: syn::Type,
+}
+
+impl NMapDef {
+	fn to_keygen_struct(&self, scrate: &proc_macro2::TokenStream) -> proc_macro2::TokenStream {
+		if self.keys.len() == 1 {
+			let hasher = &self.hashers[0].to_storage_hasher_struct();
+			let key = &self.keys[0];
+			return quote!( #scrate::storage::types::Key<#scrate::#hasher, #key> );
+		}
+
+		let key_hasher = self.keys.iter().zip(&self.hashers).map(|(key, hasher)| {
+			let hasher = hasher.to_storage_hasher_struct();
+			quote!( #scrate::storage::types::Key<#scrate::#hasher, #key> )
+		})
+		.collect::<Vec<_>>();
+		quote!(( #(#key_hasher,)* ))
+	}
+
+	fn to_key_tuple(&self) -> proc_macro2::TokenStream {
+		if self.keys.len() == 1 {
+			let key = &self.keys[0];
+			return quote!(#key);
+		}
+
+		let tuple = self.keys.iter().map(|key| {
+			quote!(#key)
+		})
+		.collect::<Vec<_>>();
+		quote!(( #(#tuple,)* ))
+	}
+}
+
 pub struct ExtraGenesisLineDef {
 	attrs: Vec<syn::Attribute>,
 	name: syn::Ident,
@@ -402,26 +461,24 @@ pub fn decl_storage_impl(input: proc_macro::TokenStream) -> proc_macro::TokenStr
 
 	print_pallet_upgrade::maybe_print_pallet_upgrade(&def_ext);
 
-	let hidden_crate_name = def_ext.hidden_crate.as_ref().map(|i| i.to_string())
-		.unwrap_or_else(|| "decl_storage".to_string());
-
-	let scrate = generate_crate_access(&hidden_crate_name, "frame-support");
-	let scrate_decl = generate_hidden_includes(&hidden_crate_name, "frame-support");
-
+	let scrate = &def_ext.hidden_crate;
+	let scrate_decl = &def_ext.hidden_imports;
 	let store_trait = store_trait::decl_and_impl(&def_ext);
-	let getters = getters::impl_getters(&scrate, &def_ext);
-	let metadata = metadata::impl_metadata(&scrate, &def_ext);
-	let instance_trait = instance_trait::decl_and_impl(&scrate, &def_ext);
-	let genesis_config = genesis_config::genesis_config_and_build_storage(&scrate, &def_ext);
-	let storage_struct = storage_struct::decl_and_impl(&scrate, &def_ext);
+	let getters = getters::impl_getters(&def_ext);
+	let metadata = metadata::impl_metadata(&def_ext);
+	let instance_trait = instance_trait::decl_and_impl(&def_ext);
+	let genesis_config = genesis_config::genesis_config_and_build_storage(&def_ext);
+	let storage_struct = storage_struct::decl_and_impl(&def_ext);
 
 	quote!(
 		use #scrate::{
 			StorageValue as _,
 			StorageMap as _,
 			StorageDoubleMap as _,
+			StorageNMap as _,
 			StoragePrefixedMap as _,
 			IterableStorageMap as _,
+			IterableStorageNMap as _,
 			IterableStorageDoubleMap as _,
 		};
 

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

@@ -29,6 +29,7 @@ mod keyword {
 	syn::custom_keyword!(get);
 	syn::custom_keyword!(map);
 	syn::custom_keyword!(double_map);
+	syn::custom_keyword!(nmap);
 	syn::custom_keyword!(opaque_blake2_256);
 	syn::custom_keyword!(opaque_blake2_128);
 	syn::custom_keyword!(blake2_128_concat);
@@ -199,6 +200,7 @@ impl_parse_for_opt!(DeclStorageBuild => keyword::build);
 enum DeclStorageType {
 	Map(DeclStorageMap),
 	DoubleMap(Box<DeclStorageDoubleMap>),
+	NMap(DeclStorageNMap),
 	Simple(syn::Type),
 }
 
@@ -208,6 +210,8 @@ impl syn::parse::Parse for DeclStorageType {
 			Ok(Self::Map(input.parse()?))
 		} else if input.peek(keyword::double_map) {
 			Ok(Self::DoubleMap(input.parse()?))
+		} else if input.peek(keyword::nmap) {
+			Ok(Self::NMap(input.parse()?))
 		} else {
 			Ok(Self::Simple(input.parse()?))
 		}
@@ -235,7 +239,21 @@ struct DeclStorageDoubleMap {
 	pub value: syn::Type,
 }
 
-#[derive(ToTokens, Debug)]
+#[derive(Parse, ToTokens, Debug)]
+struct DeclStorageKey {
+	pub hasher: Opt<SetHasher>,
+	pub key: syn::Type,
+}
+
+#[derive(Parse, ToTokens, Debug)]
+struct DeclStorageNMap {
+	pub map_keyword: keyword::nmap,
+	pub storage_keys: ext::PunctuatedTrailing<DeclStorageKey, Token![,]>,
+	pub ass_keyword: Token![=>],
+	pub value: syn::Type,
+}
+
+#[derive(Clone, ToTokens, Debug)]
 enum Hasher {
 	Blake2_256(keyword::opaque_blake2_256),
 	Blake2_128(keyword::opaque_blake2_128),
@@ -291,7 +309,7 @@ impl syn::parse::Parse for Opt<DeclStorageDefault> {
 	}
 }
 
-#[derive(Parse, ToTokens, Debug)]
+#[derive(Clone, Parse, ToTokens, Debug)]
 struct SetHasher {
 	pub hasher_keyword: keyword::hasher,
 	pub inner: ext::Parens<Hasher>,
@@ -495,6 +513,18 @@ fn parse_storage_line_defs(
 					value: map.value,
 				})
 			),
+			DeclStorageType::NMap(map) => super::StorageLineTypeDef::NMap(
+				super::NMapDef {
+					hashers: map
+						.storage_keys
+						.inner
+						.iter()
+						.map(|pair| Ok(pair.hasher.inner.clone().ok_or_else(no_hasher_error)?.into()))
+						.collect::<Result<Vec<_>, syn::Error>>()?,
+					keys: map.storage_keys.inner.iter().map(|pair| pair.key.clone()).collect(),
+					value: map.value,
+				}
+			),
 			DeclStorageType::Simple(expr) => super::StorageLineTypeDef::Simple(expr),
 		};
 

substrate/frame/support/procedural/src/storage/print_pallet_upgrade.rs

@@ -239,6 +239,15 @@ pub fn maybe_print_pallet_upgrade(def: &super::DeclStorageDefExt) {
 					comma_default_value_getter_name = comma_default_value_getter_name,
 				)
 			},
+			StorageLineTypeDef::NMap(map) => {
+				format!("StorageNMap<_, {keygen}, {value_type}{comma_query_kind}\
+					{comma_default_value_getter_name}>",
+					keygen = map.to_keygen_struct(&def.hidden_crate),
+					value_type = to_cleaned_string(&value_type),
+					comma_query_kind = comma_query_kind,
+					comma_default_value_getter_name = comma_default_value_getter_name,
+				)
+			}
 			StorageLineTypeDef::Simple(_) => {
 				format!("StorageValue<_, {value_type}{comma_query_kind}\
 					{comma_default_value_getter_name}>",

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

@@ -47,7 +47,8 @@ fn from_query_to_optional_value(is_option: bool) -> TokenStream {
 	}
 }
 
-pub fn decl_and_impl(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStream {
+pub fn decl_and_impl(def: &DeclStorageDefExt) -> TokenStream {
+	let scrate = &def.hidden_crate;
 	let mut impls = TokenStream::new();
 
 	for line in &def.storage_lines {
@@ -199,6 +200,43 @@ pub fn decl_and_impl(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStre
 							#from_optional_value_to_query
 						}
 
+						fn from_query_to_optional_value(v: Self::Query) -> Option<#value_type> {
+							#from_query_to_optional_value
+						}
+					}
+				)
+			},
+			StorageLineTypeDef::NMap(_) => {
+				quote!(
+					impl<#impl_trait> #scrate::storage::StoragePrefixedMap<#value_type>
+						for #storage_struct #optional_storage_where_clause
+					{
+						fn module_prefix() -> &'static [u8] {
+							<#instance_or_inherent as #scrate::traits::Instance>::PREFIX.as_bytes()
+						}
+
+						fn storage_prefix() -> &'static [u8] {
+							#storage_name_bstr
+						}
+					}
+
+					impl<#impl_trait> #scrate::#storage_generator_trait for #storage_struct
+					#optional_storage_where_clause
+					{
+						type Query = #query_type;
+
+						fn module_prefix() -> &'static [u8] {
+							<#instance_or_inherent as #scrate::traits::Instance>::PREFIX.as_bytes()
+						}
+
+						fn storage_prefix() -> &'static [u8] {
+							#storage_name_bstr
+						}
+
+						fn from_optional_value_to_query(v: Option<#value_type>) -> Self::Query {
+							#from_optional_value_to_query
+						}
+
 						fn from_query_to_optional_value(v: Self::Query) -> Option<#value_type> {
 							#from_query_to_optional_value
 						}