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>

frame/metadata/src/lib.rs

@@ -300,6 +300,11 @@ pub enum StorageEntryType {
 		value: DecodeDifferentStr,
 		key2_hasher: StorageHasher,
 	},
+	NMap {
+		keys: DecodeDifferentArray<&'static str, StringBuf>,
+		hashers: DecodeDifferentArray<StorageHasher>,
+		value: DecodeDifferentStr,
+	},
 }
 
 /// A storage entry modifier.
@@ -364,8 +369,10 @@ pub enum RuntimeMetadata {
 	V10(RuntimeMetadataDeprecated),
 	/// Version 11 for runtime metadata. No longer used.
 	V11(RuntimeMetadataDeprecated),
-	/// Version 12 for runtime metadata.
-	V12(RuntimeMetadataV12),
+	/// Version 12 for runtime metadata. No longer used.
+	V12(RuntimeMetadataDeprecated),
+	/// Version 13 for runtime metadata.
+	V13(RuntimeMetadataV13),
 }
 
 /// Enum that should fail.
@@ -389,15 +396,15 @@ impl Decode for RuntimeMetadataDeprecated {
 /// The metadata of a runtime.
 #[derive(Eq, Encode, PartialEq, RuntimeDebug)]
 #[cfg_attr(feature = "std", derive(Decode, Serialize))]
-pub struct RuntimeMetadataV12 {
+pub struct RuntimeMetadataV13 {
 	/// Metadata of all the modules.
 	pub modules: DecodeDifferentArray<ModuleMetadata>,
 	/// Metadata of the extrinsic.
 	pub extrinsic: ExtrinsicMetadata,
 }
 
 /// The latest version of the metadata.
-pub type RuntimeMetadataLastVersion = RuntimeMetadataV12;
+pub type RuntimeMetadataLastVersion = RuntimeMetadataV13;
 
 /// All metadata about an runtime module.
 #[derive(Clone, PartialEq, Eq, Encode, RuntimeDebug)]
@@ -425,6 +432,6 @@ impl Into<sp_core::OpaqueMetadata> for RuntimeMetadataPrefixed {
 
 impl Into<RuntimeMetadataPrefixed> for RuntimeMetadataLastVersion {
 	fn into(self) -> RuntimeMetadataPrefixed {
-		RuntimeMetadataPrefixed(META_RESERVED, RuntimeMetadata::V12(self))
+		RuntimeMetadataPrefixed(META_RESERVED, RuntimeMetadata::V13(self))
 	}
 }

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()

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

@@ -50,14 +50,19 @@ 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 {
 		value: syn::GenericArgument,
 		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));

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);
+							});
+						}}
+					},
 				});
 			}
 		}

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()

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) {

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);
 	}

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));

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();