Allow to name the generic for storages in #[pallet::storage] (#8751)

* implement named generic for storages

* fix error message on unexpected name for generic

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

@@ -18,6 +18,7 @@
 use super::helper;
 use syn::spanned::Spanned;
 use quote::ToTokens;
+use std::collections::HashMap;
 
 /// List of additional token to be used for parsing.
 mod keyword {
@@ -51,17 +52,17 @@ impl syn::parse::Parse for PalletStorageAttr {
 
 /// The value and key types used by storages. Needed to expand metadata.
 pub enum Metadata {
-	Value { value: syn::GenericArgument },
-	Map { value: syn::GenericArgument, key: syn::GenericArgument },
+	Value { value: syn::Type },
+	Map { value: syn::Type, key: syn::Type },
 	DoubleMap {
-		value: syn::GenericArgument,
-		key1: syn::GenericArgument,
-		key2: syn::GenericArgument
+		value: syn::Type,
+		key1: syn::Type,
+		key2: syn::Type
 	},
 	NMap {
 		keys: Vec<syn::Type>,
-		keygen: syn::GenericArgument,
-		value: syn::GenericArgument,
+		keygen: syn::Type,
+		value: syn::Type,
 	},
 }
 
@@ -98,41 +99,402 @@ pub struct StorageDef {
 	pub attr_span: proc_macro2::Span,
 	/// The `cfg` attributes.
 	pub cfg_attrs: Vec<syn::Attribute>,
+	/// If generics are named (e.g. `StorageValue<Value = u32, ..>`) then this contains all the
+	/// generics of the storage.
+	/// If generics are not named, this is none.
+	pub named_generics: Option<StorageGenerics>,
 }
 
-/// In `Foo<A, B, C>` retrieve the argument at given position, i.e. A is argument at position 0.
-fn retrieve_arg(
-	segment: &syn::PathSegment,
-	arg_pos: usize,
-) -> syn::Result<syn::GenericArgument> {
-	if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
-		if arg_pos < args.args.len() {
-			Ok(args.args[arg_pos].clone())
-		} else {
-			let msg = format!("pallet::storage unexpected number of generic argument, expected at \
-				least {} args, found {}", arg_pos + 1, args.args.len());
-			Err(syn::Error::new(args.span(), msg))
+
+/// The parsed generic from the
+#[derive(Clone)]
+pub enum StorageGenerics {
+	DoubleMap {
+		hasher1: syn::Type,
+		key1: syn::Type,
+		hasher2: syn::Type,
+		key2: syn::Type,
+		value: syn::Type,
+		query_kind: Option<syn::Type>,
+		on_empty: Option<syn::Type>,
+		max_values: Option<syn::Type>,
+	},
+	Map {
+		hasher: syn::Type,
+		key: syn::Type,
+		value: syn::Type,
+		query_kind: Option<syn::Type>,
+		on_empty: Option<syn::Type>,
+		max_values: Option<syn::Type>,
+	},
+	Value {
+		value: syn::Type,
+		query_kind: Option<syn::Type>,
+		on_empty: Option<syn::Type>,
+	},
+	NMap {
+		keygen: syn::Type,
+		value: syn::Type,
+		query_kind: Option<syn::Type>,
+		on_empty: Option<syn::Type>,
+		max_values: Option<syn::Type>,
+	},
+}
+
+impl StorageGenerics {
+	/// Return the metadata from the defined generics
+	fn metadata(&self) -> syn::Result<Metadata> {
+		let res = match self.clone() {
+			Self::DoubleMap { value, key1, key2, .. } => Metadata::DoubleMap { value, key1, key2 },
+			Self::Map { value, key, .. } => Metadata::Map { value, key },
+			Self::Value { value, .. } => Metadata::Value { value },
+			Self::NMap { keygen, value, .. } => Metadata::NMap {
+				keys: collect_keys(&keygen)?,
+				keygen,
+				value,
+			},
+		};
+
+		Ok(res)
+	}
+
+	/// Return the query kind from the defined generics
+	fn query_kind(&self) -> Option<syn::Type> {
+		match &self {
+			Self::DoubleMap { query_kind, .. }
+			| Self::Map { query_kind, .. }
+			| Self::Value { query_kind, .. }
+			| Self::NMap { query_kind, .. }
+				=> query_kind.clone(),
 		}
+	}
+}
+
+enum StorageKind {
+	Value,
+	Map,
+	DoubleMap,
+	NMap,
+}
+
+/// Check the generics in the `map` contains the generics in `gen` may contains generics in
+/// `optional_gen`, and doesn't contains any other.
+fn check_generics(
+	map: &HashMap<String, syn::Binding>,
+	mandatory_generics: &[&str],
+	optional_generics: &[&str],
+	storage_type_name: &str,
+	args_span: proc_macro2::Span,
+) -> syn::Result<()> {
+	let mut errors = vec![];
+
+	let expectation = {
+		let mut e = format!(
+			"`{}` expect generics {}and optional generics {}",
+			storage_type_name,
+			mandatory_generics.iter().map(|name| format!("`{}`, ", name)).collect::<String>(),
+			&optional_generics.iter().map(|name| format!("`{}`, ", name)).collect::<String>(),
+		);
+		e.pop();
+		e.pop();
+		e.push_str(".");
+		e
+	};
+
+	for (gen_name, gen_binding) in map {
+		if !mandatory_generics.contains(&gen_name.as_str())
+			&& !optional_generics.contains(&gen_name.as_str())
+		{
+			let msg = format!(
+				"Invalid pallet::storage, Unexpected generic `{}` for `{}`. {}",
+				gen_name,
+				storage_type_name,
+				expectation,
+			);
+			errors.push(syn::Error::new(gen_binding.span(), msg));
+		}
+	}
+
+	for mandatory_generic in mandatory_generics {
+		if !map.contains_key(&mandatory_generic.to_string()) {
+			let msg = format!(
+				"Invalid pallet::storage, cannot find `{}` generic, required for `{}`.",
+				mandatory_generic,
+				storage_type_name
+			);
+			errors.push(syn::Error::new(args_span, msg));
+		}
+	}
+
+	let mut errors = errors.drain(..);
+	if let Some(mut error) = errors.next() {
+		for other_error in errors {
+			error.combine(other_error);
+		}
+		Err(error)
 	} else {
-		let msg = format!("pallet::storage unexpected number of generic argument, expected at \
-			least {} args, found none", arg_pos + 1);
+		Ok(())
+	}
+}
+
+/// Returns `(named generics, metadata, query kind)`
+fn process_named_generics(
+	storage: &StorageKind,
+	args_span: proc_macro2::Span,
+	args: &[syn::Binding],
+) -> syn::Result<(Option<StorageGenerics>, Metadata, Option<syn::Type>)> {
+	let mut parsed = HashMap::<String, syn::Binding>::new();
+
+	// Ensure no duplicate.
+	for arg in args {
+		if let Some(other) = parsed.get(&arg.ident.to_string()) {
+			let msg = "Invalid pallet::storage, Duplicated named generic";
+			let mut err = syn::Error::new(arg.ident.span(), msg);
+			err.combine(syn::Error::new(other.ident.span(), msg));
+			return Err(err);
+		}
+		parsed.insert(arg.ident.to_string(), arg.clone());
+	}
+
+	let generics = match storage {
+		StorageKind::Value => {
+			check_generics(
+				&parsed,
+				&["Value"],
+				&["QueryKind", "OnEmpty"],
+				"StorageValue",
+				args_span,
+			)?;
+
+			StorageGenerics::Value {
+				value: parsed.remove("Value")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				query_kind: parsed.remove("QueryKind")
+					.map(|binding| binding.ty),
+				on_empty: parsed.remove("OnEmpty")
+					.map(|binding| binding.ty),
+			}
+		}
+		StorageKind::Map => {
+			check_generics(
+				&parsed,
+				&["Hasher", "Key", "Value"],
+				&["QueryKind", "OnEmpty", "MaxValues"],
+				"StorageMap",
+				args_span,
+			)?;
+
+			StorageGenerics::Map {
+				hasher: parsed.remove("Hasher")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				key: parsed.remove("Key")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				value: parsed.remove("Value")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				query_kind: parsed.remove("QueryKind").map(|binding| binding.ty),
+				on_empty: parsed.remove("OnEmpty").map(|binding| binding.ty),
+				max_values: parsed.remove("MaxValues").map(|binding| binding.ty),
+			}
+		}
+		StorageKind::DoubleMap => {
+			check_generics(
+				&parsed,
+				&["Hasher1", "Key1", "Hasher2", "Key2", "Value"],
+				&["QueryKind", "OnEmpty", "MaxValues"],
+				"StorageDoubleMap",
+				args_span,
+			)?;
+
+			StorageGenerics::DoubleMap {
+				hasher1: parsed.remove("Hasher1")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				key1: parsed.remove("Key1")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				hasher2: parsed.remove("Hasher2")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				key2: parsed.remove("Key2")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				value: parsed.remove("Value")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				query_kind: parsed.remove("QueryKind").map(|binding| binding.ty),
+				on_empty: parsed.remove("OnEmpty").map(|binding| binding.ty),
+				max_values: parsed.remove("MaxValues").map(|binding| binding.ty),
+			}
+		}
+		StorageKind::NMap => {
+			check_generics(
+				&parsed,
+				&["Key", "Value"],
+				&["QueryKind", "OnEmpty", "MaxValues"],
+				"StorageNMap",
+				args_span,
+			)?;
+
+			StorageGenerics::NMap {
+				keygen: parsed.remove("Key")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				value: parsed.remove("Value")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				query_kind: parsed.remove("QueryKind").map(|binding| binding.ty),
+				on_empty: parsed.remove("OnEmpty").map(|binding| binding.ty),
+				max_values: parsed.remove("MaxValues").map(|binding| binding.ty),
+			}
+		}
+	};
+
+	let metadata = generics.metadata()?;
+	let query_kind = generics.query_kind();
+
+	Ok((Some(generics), metadata, query_kind))
+}
+
+/// Returns `(named generics, metadata, query kind)`
+fn process_unnamed_generics(
+	storage: &StorageKind,
+	args_span: proc_macro2::Span,
+	args: &[syn::Type],
+) -> syn::Result<(Option<StorageGenerics>, Metadata, Option<syn::Type>)> {
+	let retrieve_arg = |arg_pos| {
+		args.get(arg_pos)
+			.cloned()
+			.ok_or_else(|| {
+				let msg = format!(
+					"Invalid pallet::storage, unexpected number of generic argument, \
+						expect at least {} args, found {}.",
+					arg_pos + 1,
+					args.len(),
+				);
+				syn::Error::new(args_span, msg)
+			})
+	};
+
+	let prefix_arg = retrieve_arg(0)?;
+	syn::parse2::<syn::Token![_]>(prefix_arg.to_token_stream())
+		.map_err(|e| {
+			let msg = "Invalid pallet::storage, for unnamed generic arguments the type \
+				first generic argument must be `_`, the argument is then replaced by macro.";
+			let mut err = syn::Error::new(prefix_arg.span(), msg);
+			err.combine(e);
+			err
+		})?;
+
+	let res = match storage {
+		StorageKind::Value => (
+			None,
+			Metadata::Value { value: retrieve_arg(1)? },
+			retrieve_arg(2).ok(),
+		),
+		StorageKind::Map => (
+			None,
+			Metadata::Map {
+				key: retrieve_arg(2)?,
+				value: retrieve_arg(3)?,
+			},
+			retrieve_arg(4).ok(),
+		),
+		StorageKind::DoubleMap => (
+			None,
+			Metadata::DoubleMap {
+				key1: retrieve_arg(2)?,
+				key2: retrieve_arg(4)?,
+				value: retrieve_arg(5)?,
+			},
+			retrieve_arg(6).ok(),
+		),
+		StorageKind::NMap => {
+			let keygen = retrieve_arg(1)?;
+			let keys = collect_keys(&keygen)?;
+			(
+				None,
+				Metadata::NMap {
+					keys,
+					keygen,
+					value: retrieve_arg(2)?,
+				},
+				retrieve_arg(3).ok(),
+			)
+		},
+	};
+
+	Ok(res)
+}
+
+/// Returns `(named generics, metadata, query kind)`
+fn process_generics(
+	segment: &syn::PathSegment,
+) -> syn::Result<(Option<StorageGenerics>, Metadata, Option<syn::Type>)> {
+	let storage_kind = match &*segment.ident.to_string() {
+		"StorageValue" => StorageKind::Value,
+		"StorageMap" => StorageKind::Map,
+		"StorageDoubleMap" => StorageKind::DoubleMap,
+		"StorageNMap" => StorageKind::NMap,
+		found => {
+			let msg = format!(
+				"Invalid pallet::storage, expected ident: `StorageValue` or \
+				`StorageMap` or `StorageDoubleMap` or `StorageNMap` in order to expand metadata, \
+				found `{}`.",
+				found,
+			);
+			return Err(syn::Error::new(segment.ident.span(), msg));
+		}
+	};
+
+	let args_span = segment.arguments.span();
+
+	let args = match &segment.arguments {
+		syn::PathArguments::AngleBracketed(args) if args.args.len() != 0 => args,
+		_ => {
+			let msg = "Invalid pallet::storage, invalid number of generic generic arguments, \
+				expect more that 0 generic arguments.";
+			return Err(syn::Error::new(segment.span(), msg));
+		}
+	};
+
+	if args.args.iter().all(|gen| matches!(gen, syn::GenericArgument::Type(_))) {
+		let args = args.args.iter()
+			.map(|gen| match gen {
+				syn::GenericArgument::Type(gen) => gen.clone(),
+				_ => unreachable!("It is asserted above that all generics are types"),
+			})
+			.collect::<Vec<_>>();
+		process_unnamed_generics(&storage_kind, args_span, &args)
+	} else if args.args.iter().all(|gen| matches!(gen, syn::GenericArgument::Binding(_))) {
+		let args = args.args.iter()
+			.map(|gen| match gen {
+				syn::GenericArgument::Binding(gen) => gen.clone(),
+				_ => unreachable!("It is asserted above that all generics are bindings"),
+			})
+			.collect::<Vec<_>>();
+		process_named_generics(&storage_kind, args_span, &args)
+	} else {
+		let msg = "Invalid pallet::storage, invalid generic declaration for storage. Expect only \
+			type generics or binding generics, e.g. `<Name1 = Gen1, Name2 = Gen2, ..>` or \
+			`<Gen1, Gen2, ..>`.";
 		Err(syn::Error::new(segment.span(), msg))
 	}
 }
 
 /// 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 {
+fn collect_keys(keygen: &syn::Type) -> syn::Result<Vec<syn::Type>> {
+	if let 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))
+		Ok(vec![extract_key(keygen)?])
 	}
 }
 
@@ -187,7 +549,7 @@ impl StorageDef {
 		let item = if let syn::Item::Type(item) = item {
 			item
 		} else {
-			return Err(syn::Error::new(item.span(), "Invalid pallet::storage, expected item type"));
+			return Err(syn::Error::new(item.span(), "Invalid pallet::storage, expect item type."));
 		};
 
 		let mut attrs: Vec<PalletStorageAttr> = helper::take_item_pallet_attrs(&mut item.attrs)?;
@@ -217,55 +579,14 @@ impl StorageDef {
 			return Err(syn::Error::new(item.ty.span(), msg));
 		}
 
-		let query_kind;
-		let metadata = match &*typ.path.segments[0].ident.to_string() {
-			"StorageValue" => {
-				query_kind = retrieve_arg(&typ.path.segments[0], 2);
-				Metadata::Value {
-					value: retrieve_arg(&typ.path.segments[0], 1)?,
-				}
-			}
-			"StorageMap" => {
-				query_kind = retrieve_arg(&typ.path.segments[0], 4);
-				Metadata::Map {
-					key: retrieve_arg(&typ.path.segments[0], 2)?,
-					value: retrieve_arg(&typ.path.segments[0], 3)?,
-				}
-			}
-			"StorageDoubleMap" => {
-				query_kind = retrieve_arg(&typ.path.segments[0], 6);
-				Metadata::DoubleMap {
-					key1: retrieve_arg(&typ.path.segments[0], 2)?,
-					key2: retrieve_arg(&typ.path.segments[0], 4)?,
-					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` or `StorageNMap` in order \
-					to expand metadata, found `{}`",
-					found,
-				);
-				return Err(syn::Error::new(item.ty.span(), msg));
-			}
-		};
+		let (named_generics, metadata, query_kind) = process_generics(&typ.path.segments[0])?;
+
 		let query_kind = query_kind
 			.map(|query_kind| match query_kind {
-				syn::GenericArgument::Type(syn::Type::Path(path))
+				syn::Type::Path(path)
 					if path.path.segments.last().map_or(false, |s| s.ident == "OptionQuery")
 				=> Some(QueryKind::OptionQuery),
-				syn::GenericArgument::Type(syn::Type::Path(path))
+				syn::Type::Path(path)
 					if path.path.segments.last().map_or(false, |s| s.ident == "ValueQuery")
 				=> Some(QueryKind::ValueQuery),
 				_ => None,
@@ -279,16 +600,6 @@ impl StorageDef {
 			return Err(syn::Error::new(getter.unwrap().span(), msg));
 		}
 
-		let prefix_arg = retrieve_arg(&typ.path.segments[0], 0)?;
-		syn::parse2::<syn::Token![_]>(prefix_arg.to_token_stream())
-			.map_err(|e| {
-				let msg = "Invalid use of `#[pallet::storage]`, the type first generic argument \
-					must be `_`, the final argument is automatically set by macro.";
-				let mut err = syn::Error::new(prefix_arg.span(), msg);
-				err.combine(e);
-				err
-			})?;
-
 		Ok(StorageDef {
 			attr_span,
 			index,
@@ -301,6 +612,7 @@ impl StorageDef {
 			query_kind,
 			where_clause,
 			cfg_attrs,
+			named_generics,
 		})
 	}
 }