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/pallet/expand/storage.rs

@@ -90,6 +90,9 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 				Metadata::DoubleMap { .. } => quote::quote_spanned!(storage.attr_span =>
 					#frame_support::storage::types::StorageDoubleMapMetadata
 				),
+				Metadata::NMap { .. } => quote::quote_spanned!(storage.attr_span =>
+					#frame_support::storage::types::StorageNMapMetadata
+				),
 			};
 
 			let ty = match &storage.metadata {
@@ -126,6 +129,24 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 							value: #frame_support::metadata::DecodeDifferent::Encode(#value),
 						}
 					)
+				},
+				Metadata::NMap { keys, value, .. } => {
+					let keys = keys
+						.iter()
+						.map(|key| clean_type_string(&quote::quote!(#key).to_string()))
+						.collect::<Vec<_>>();
+					let value = clean_type_string(&quote::quote!(#value).to_string());
+					quote::quote_spanned!(storage.attr_span =>
+						#frame_support::metadata::StorageEntryType::NMap {
+							keys: #frame_support::metadata::DecodeDifferent::Encode(&[
+								#( #keys, )*
+							]),
+							hashers: #frame_support::metadata::DecodeDifferent::Encode(
+								<#full_ident as #metadata_trait>::HASHERS,
+							),
+							value: #frame_support::metadata::DecodeDifferent::Encode(#value),
+						}
+					)
 				}
 			};
 
@@ -227,6 +248,32 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						}
 					)
 				},
+				Metadata::NMap { keygen, value, .. } => {
+					let query = match storage.query_kind.as_ref().expect("Checked by def") {
+						QueryKind::OptionQuery => quote::quote_spanned!(storage.attr_span =>
+							Option<#value>
+						),
+						QueryKind::ValueQuery => quote::quote!(#value),
+					};
+					quote::quote_spanned!(storage.attr_span =>
+						#(#cfg_attrs)*
+						impl<#type_impl_gen> #pallet_ident<#type_use_gen> #completed_where_clause {
+							#( #docs )*
+							pub fn #getter<KArg>(key: KArg) -> #query
+							where
+								KArg: #frame_support::storage::types::EncodeLikeTuple<
+									<#keygen as #frame_support::storage::types::KeyGenerator>::KArg
+								>
+									+ #frame_support::storage::types::TupleToEncodedIter,
+							{
+								<
+									#full_ident as
+									#frame_support::storage::StorageNMap<#keygen, #value>
+								>::get(key)
+							}
+						}
+					)
+				}
 			}
 		} else {
 			Default::default()

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

@@ -50,7 +50,7 @@ impl syn::parse::Parse for PalletStorageAttr {
 }
 
 /// The value and key types used by storages. Needed to expand metadata.
-pub enum Metadata{
+pub enum Metadata {
 	Value { value: syn::GenericArgument },
 	Map { value: syn::GenericArgument, key: syn::GenericArgument },
 	DoubleMap {
@@ -58,6 +58,11 @@ pub enum Metadata{
 		key1: syn::GenericArgument,
 		key2: syn::GenericArgument
 	},
+	NMap {
+		keys: Vec<syn::Type>,
+		keygen: syn::GenericArgument,
+		value: syn::GenericArgument,
+	},
 }
 
 pub enum QueryKind {
@@ -115,6 +120,64 @@ fn retrieve_arg(
 	}
 }
 
+/// Parse the 2nd type argument to `StorageNMap` and return its keys.
+fn collect_keys(keygen: &syn::GenericArgument) -> syn::Result<Vec<syn::Type>> {
+	if let syn::GenericArgument::Type(syn::Type::Tuple(tup)) = keygen {
+		tup
+			.elems
+			.iter()
+			.map(extract_key)
+			.collect::<syn::Result<Vec<_>>>()
+	} else if let syn::GenericArgument::Type(ty) = keygen {
+		Ok(vec![extract_key(ty)?])
+	} else {
+		let msg = format!("Invalid pallet::storage, expected tuple of Key structs or Key struct");
+		Err(syn::Error::new(keygen.span(), msg))
+	}
+}
+
+/// In `Key<H, K>`, extract K and return it.
+fn extract_key(ty: &syn::Type) -> syn::Result<syn::Type> {
+	let typ = if let syn::Type::Path(typ) = ty {
+		typ
+	} else {
+		let msg = "Invalid pallet::storage, expected type path";
+		return Err(syn::Error::new(ty.span(), msg));
+	};
+
+	let key_struct = typ.path.segments.last().ok_or_else(|| {
+		let msg = "Invalid pallet::storage, expected type path with at least one segment";
+		syn::Error::new(typ.path.span(), msg)
+	})?;
+	if key_struct.ident != "Key" && key_struct.ident != "NMapKey" {
+		let msg = "Invalid pallet::storage, expected Key or NMapKey struct";
+		return Err(syn::Error::new(key_struct.ident.span(), msg));
+	}
+
+	let ty_params = if let syn::PathArguments::AngleBracketed(args) = &key_struct.arguments {
+		args
+	} else {
+		let msg = "Invalid pallet::storage, expected angle bracketed arguments";
+		return Err(syn::Error::new(key_struct.arguments.span(), msg));
+	};
+
+	if ty_params.args.len() != 2 {
+		let msg = format!("Invalid pallet::storage, unexpected number of generic arguments \
+			for Key struct, expected 2 args, found {}", ty_params.args.len());
+		return Err(syn::Error::new(ty_params.span(), msg));
+	}
+
+	let key = match &ty_params.args[1] {
+		syn::GenericArgument::Type(key_ty) => key_ty.clone(),
+		_ => {
+			let msg = "Invalid pallet::storage, expected type";
+			return Err(syn::Error::new(ty_params.args[1].span(), msg));
+		}
+	};
+
+	Ok(key)
+}
+
 impl StorageDef {
 	pub fn try_from(
 		attr_span: proc_macro2::Span,
@@ -177,11 +240,21 @@ impl StorageDef {
 					value: retrieve_arg(&typ.path.segments[0], 5)?,
 				}
 			}
+			"StorageNMap" => {
+				query_kind = retrieve_arg(&typ.path.segments[0], 3);
+				let keygen = retrieve_arg(&typ.path.segments[0], 1)?;
+				let keys = collect_keys(&keygen)?;
+				Metadata::NMap {
+					keys,
+					keygen,
+					value: retrieve_arg(&typ.path.segments[0], 2)?,
+				}
+			}
 			found => {
 				let msg = format!(
 					"Invalid pallet::storage, expected ident: `StorageValue` or \
-					`StorageMap` or `StorageDoubleMap` in order to expand metadata, found \
-					`{}`",
+					`StorageMap` or `StorageDoubleMap` or `StorageNMap` in order \
+					to expand metadata, found `{}`",
 					found,
 				);
 				return Err(syn::Error::new(item.ty.span(), msg));