Implement ResultQuery

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

@@ -19,7 +19,9 @@ use crate::pallet::{
 	parse::storage::{Metadata, QueryKind, StorageDef, StorageGenerics},
 	Def,
 };
-use std::collections::HashMap;
+use quote::ToTokens;
+use std::{collections::HashMap, ops::IndexMut};
+use syn::spanned::Spanned;
 
 /// Generate the prefix_ident related to the storage.
 /// prefix_ident is used for the prefix struct to be given to storage as first generic param.
@@ -84,12 +86,28 @@ fn check_prefix_duplicates(
 	Ok(())
 }
 
+pub struct ResultOnEmptyStructMetadata {
+	/// The Rust ident that is going to be used as the name of the OnEmpty struct.
+	pub name: syn::Ident,
+	/// The path to the error type being returned by the ResultQuery.
+	pub error_path: syn::Path,
+	/// The visibility of the OnEmpty struct.
+	pub visibility: syn::Visibility,
+	/// The type of the storage item.
+	pub value_ty: syn::Type,
+	/// The name of the pallet error enum variant that is going to be returned.
+	pub variant_name: syn::Ident,
+	/// The span used to report compilation errors about the OnEmpty struct.
+	pub span: proc_macro2::Span,
+}
+
 ///
 /// * if generics are unnamed: replace the first generic `_` by the generated prefix structure
 /// * if generics are named: reorder the generic, remove their name, and add the missing ones.
 /// * Add `#[allow(type_alias_bounds)]`
-pub fn process_generics(def: &mut Def) -> syn::Result<()> {
+pub fn process_generics(def: &mut Def) -> syn::Result<Vec<ResultOnEmptyStructMetadata>> {
 	let frame_support = &def.frame_support;
+	let mut on_empty_struct_metadata = Vec::new();
 
 	for storage_def in def.storages.iter_mut() {
 		let item = &mut def.item.content.as_mut().expect("Checked by def").1[storage_def.index];
@@ -120,27 +138,72 @@ pub fn process_generics(def: &mut Def) -> syn::Result<()> {
 
 		let default_query_kind: syn::Type =
 			syn::parse_quote!(#frame_support::storage::types::OptionQuery);
-		let default_on_empty: syn::Type = syn::parse_quote!(#frame_support::traits::GetDefault);
+		let mut default_on_empty = |value_ty: syn::Type| -> syn::Type {
+			if let Some(QueryKind::ResultQuery(error_path, variant_name)) =
+				storage_def.query_kind.as_ref()
+			{
+				let on_empty_ident =
+					quote::format_ident!("__Frame_Internal_Get{}Result", storage_def.ident);
+				on_empty_struct_metadata.push(ResultOnEmptyStructMetadata {
+					name: on_empty_ident.clone(),
+					visibility: storage_def.vis.clone(),
+					value_ty,
+					error_path: error_path.clone(),
+					variant_name: variant_name.clone(),
+					span: storage_def.attr_span,
+				});
+				return syn::parse_quote!(#on_empty_ident)
+			}
+			syn::parse_quote!(#frame_support::traits::GetDefault)
+		};
 		let default_max_values: syn::Type = syn::parse_quote!(#frame_support::traits::GetDefault);
 
+		let set_result_query_type_parameter = |query_type: &mut syn::Type| -> syn::Result<()> {
+			if let Some(QueryKind::ResultQuery(error_path, _)) = storage_def.query_kind.as_ref() {
+				if let syn::Type::Path(syn::TypePath { path: syn::Path { segments, .. }, .. }) =
+					query_type
+				{
+					if let Some(seg) = segments.last_mut() {
+						if let syn::PathArguments::AngleBracketed(
+							syn::AngleBracketedGenericArguments { args, .. },
+						) = &mut seg.arguments
+						{
+							args.clear();
+							args.push(syn::GenericArgument::Type(syn::parse_quote!(#error_path)));
+						}
+					}
+				} else {
+					let msg = format!(
+						"Invalid pallet::storage, unexpected type for query, expected ResultQuery \
+						with 1 type parameter, found `{}`",
+						query_type.to_token_stream().to_string()
+					);
+					return Err(syn::Error::new(query_type.span(), msg))
+				}
+			}
+			Ok(())
+		};
+
 		if let Some(named_generics) = storage_def.named_generics.clone() {
 			args.args.clear();
 			args.args.push(syn::parse_quote!( #prefix_ident<#type_use_gen> ));
 			match named_generics {
 				StorageGenerics::Value { value, query_kind, on_empty } => {
-					args.args.push(syn::GenericArgument::Type(value));
-					let query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
+					args.args.push(syn::GenericArgument::Type(value.clone()));
+					let mut query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
+					set_result_query_type_parameter(&mut query_kind)?;
 					args.args.push(syn::GenericArgument::Type(query_kind));
-					let on_empty = on_empty.unwrap_or_else(|| default_on_empty.clone());
+					let on_empty = on_empty.unwrap_or_else(|| default_on_empty(value));
 					args.args.push(syn::GenericArgument::Type(on_empty));
 				},
 				StorageGenerics::Map { hasher, key, value, query_kind, on_empty, max_values } => {
 					args.args.push(syn::GenericArgument::Type(hasher));
 					args.args.push(syn::GenericArgument::Type(key));
-					args.args.push(syn::GenericArgument::Type(value));
-					let query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
+					args.args.push(syn::GenericArgument::Type(value.clone()));
+					let mut query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
+					set_result_query_type_parameter(&mut query_kind)?;
 					args.args.push(syn::GenericArgument::Type(query_kind));
-					let on_empty = on_empty.unwrap_or_else(|| default_on_empty.clone());
+					let on_empty = on_empty.unwrap_or_else(|| default_on_empty(value));
 					args.args.push(syn::GenericArgument::Type(on_empty));
 					let max_values = max_values.unwrap_or_else(|| default_max_values.clone());
 					args.args.push(syn::GenericArgument::Type(max_values));
@@ -155,10 +218,11 @@ pub fn process_generics(def: &mut Def) -> syn::Result<()> {
 				} => {
 					args.args.push(syn::GenericArgument::Type(hasher));
 					args.args.push(syn::GenericArgument::Type(key));
-					args.args.push(syn::GenericArgument::Type(value));
-					let query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
+					args.args.push(syn::GenericArgument::Type(value.clone()));
+					let mut query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
+					set_result_query_type_parameter(&mut query_kind)?;
 					args.args.push(syn::GenericArgument::Type(query_kind));
-					let on_empty = on_empty.unwrap_or_else(|| default_on_empty.clone());
+					let on_empty = on_empty.unwrap_or_else(|| default_on_empty(value));
 					args.args.push(syn::GenericArgument::Type(on_empty));
 					let max_values = max_values.unwrap_or_else(|| default_max_values.clone());
 					args.args.push(syn::GenericArgument::Type(max_values));
@@ -177,31 +241,65 @@ pub fn process_generics(def: &mut Def) -> syn::Result<()> {
 					args.args.push(syn::GenericArgument::Type(key1));
 					args.args.push(syn::GenericArgument::Type(hasher2));
 					args.args.push(syn::GenericArgument::Type(key2));
-					args.args.push(syn::GenericArgument::Type(value));
-					let query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
+					args.args.push(syn::GenericArgument::Type(value.clone()));
+					let mut query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
+					set_result_query_type_parameter(&mut query_kind)?;
 					args.args.push(syn::GenericArgument::Type(query_kind));
-					let on_empty = on_empty.unwrap_or_else(|| default_on_empty.clone());
+					let on_empty = on_empty.unwrap_or_else(|| default_on_empty(value));
 					args.args.push(syn::GenericArgument::Type(on_empty));
 					let max_values = max_values.unwrap_or_else(|| default_max_values.clone());
 					args.args.push(syn::GenericArgument::Type(max_values));
 				},
 				StorageGenerics::NMap { keygen, value, query_kind, on_empty, max_values } => {
 					args.args.push(syn::GenericArgument::Type(keygen));
-					args.args.push(syn::GenericArgument::Type(value));
-					let query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
+					args.args.push(syn::GenericArgument::Type(value.clone()));
+					let mut query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
+					set_result_query_type_parameter(&mut query_kind)?;
 					args.args.push(syn::GenericArgument::Type(query_kind));
-					let on_empty = on_empty.unwrap_or_else(|| default_on_empty.clone());
+					let on_empty = on_empty.unwrap_or_else(|| default_on_empty(value));
 					args.args.push(syn::GenericArgument::Type(on_empty));
 					let max_values = max_values.unwrap_or_else(|| default_max_values.clone());
 					args.args.push(syn::GenericArgument::Type(max_values));
 				},
 			}
 		} else {
 			args.args[0] = syn::parse_quote!( #prefix_ident<#type_use_gen> );
+
+			let (value_idx, query_idx, on_empty_idx) = match storage_def.metadata {
+				Metadata::Value { .. } => (1, 2, 3),
+				Metadata::NMap { .. } => (2, 3, 4),
+				Metadata::Map { .. } | Metadata::CountedMap { .. } => (3, 4, 5),
+				Metadata::DoubleMap { .. } => (5, 6, 7),
+			};
+
+			// #FIXME: Really stupid, the index and index_mut methods on syn::Punctuated don't
+			// return Options so we have to manually check whether they're over the length limit.
+			if query_idx < args.args.len() {
+				if let syn::GenericArgument::Type(query_kind) = args.args.index_mut(query_idx) {
+					set_result_query_type_parameter(query_kind)?;
+				}
+			} else if let Some(QueryKind::ResultQuery(error_path, _)) =
+				storage_def.query_kind.as_ref()
+			{
+				args.args.push(syn::GenericArgument::Type(syn::parse_quote!(#error_path)))
+			}
+
+			// Here, we only need to check if OnEmpty is *not* specified, and if so, then we have to
+			// generate a default OnEmpty struct for it.
+			if on_empty_idx >= args.args.len() &&
+				matches!(storage_def.query_kind.as_ref(), Some(QueryKind::ResultQuery(_, _)))
+			{
+				let value_ty = match args.args[value_idx].clone() {
+					syn::GenericArgument::Type(ty) => ty,
+					_ => unreachable!(),
+				};
+				let on_empty = default_on_empty(value_ty);
+				args.args.push(syn::GenericArgument::Type(on_empty));
+			}
 		}
 	}
 
-	Ok(())
+	Ok(on_empty_struct_metadata)
 }
 
 ///
@@ -212,9 +310,10 @@ pub fn process_generics(def: &mut Def) -> syn::Result<()> {
 /// * Add `#[allow(type_alias_bounds)]` on storages type alias
 /// * generate metadatas
 pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
-	if let Err(e) = process_generics(def) {
-		return e.into_compile_error()
-	}
+	let on_empty_struct_metadata = match process_generics(def) {
+		Ok(idents) => idents,
+		Err(e) => return e.into_compile_error(),
+	};
 
 	// Check for duplicate prefixes
 	let mut prefix_set = HashMap::new();
@@ -277,6 +376,10 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						QueryKind::OptionQuery => quote::quote_spanned!(storage.attr_span =>
 							Option<#value>
 						),
+						QueryKind::ResultQuery(error_path, _) =>
+							quote::quote_spanned!(storage.attr_span =>
+								Result<#value, #error_path>
+							),
 						QueryKind::ValueQuery => quote::quote!(#value),
 					};
 					quote::quote_spanned!(storage.attr_span =>
@@ -296,6 +399,10 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						QueryKind::OptionQuery => quote::quote_spanned!(storage.attr_span =>
 							Option<#value>
 						),
+						QueryKind::ResultQuery(error_path, _) =>
+							quote::quote_spanned!(storage.attr_span =>
+								Result<#value, #error_path>
+							),
 						QueryKind::ValueQuery => quote::quote!(#value),
 					};
 					quote::quote_spanned!(storage.attr_span =>
@@ -317,6 +424,10 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						QueryKind::OptionQuery => quote::quote_spanned!(storage.attr_span =>
 							Option<#value>
 						),
+						QueryKind::ResultQuery(error_path, _) =>
+							quote::quote_spanned!(storage.attr_span =>
+								Result<#value, #error_path>
+							),
 						QueryKind::ValueQuery => quote::quote!(#value),
 					};
 					quote::quote_spanned!(storage.attr_span =>
@@ -338,6 +449,10 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						QueryKind::OptionQuery => quote::quote_spanned!(storage.attr_span =>
 							Option<#value>
 						),
+						QueryKind::ResultQuery(error_path, _) =>
+							quote::quote_spanned!(storage.attr_span =>
+								Result<#value, #error_path>
+							),
 						QueryKind::ValueQuery => quote::quote!(#value),
 					};
 					quote::quote_spanned!(storage.attr_span =>
@@ -361,6 +476,10 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						QueryKind::OptionQuery => quote::quote_spanned!(storage.attr_span =>
 							Option<#value>
 						),
+						QueryKind::ResultQuery(error_path, _) =>
+							quote::quote_spanned!(storage.attr_span =>
+								Result<#value, #error_path>
+							),
 						QueryKind::ValueQuery => quote::quote!(#value),
 					};
 					quote::quote_spanned!(storage.attr_span =>
@@ -459,6 +578,61 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 		)
 	});
 
+	let on_empty_structs = on_empty_struct_metadata.into_iter().map(|metadata| {
+		use crate::pallet::parse::GenericKind;
+		use syn::{GenericArgument, Path, PathArguments, PathSegment, Type, TypePath};
+
+		let ResultOnEmptyStructMetadata {
+			name,
+			visibility,
+			value_ty,
+			error_path,
+			variant_name,
+			span,
+		} = metadata;
+
+		let generic_kind = match error_path.segments.last() {
+			Some(PathSegment { arguments: PathArguments::AngleBracketed(args), .. }) => {
+				let (has_config, has_instance) =
+					args.args.iter().fold((false, false), |(has_config, has_instance), arg| {
+						match arg {
+							GenericArgument::Type(Type::Path(TypePath {
+								path: Path { segments, .. },
+								..
+							})) => {
+								let maybe_config =
+									segments.first().map_or(false, |seg| seg.ident == "T");
+								let maybe_instance =
+									segments.first().map_or(false, |seg| seg.ident == "I");
+
+								(has_config || maybe_config, has_instance || maybe_instance)
+							},
+							_ => (has_config, has_instance),
+						}
+					});
+				GenericKind::from_gens(has_config, has_instance).unwrap_or(GenericKind::None)
+			},
+			_ => GenericKind::None,
+		};
+		let type_impl_gen = generic_kind.type_impl_gen(proc_macro2::Span::call_site());
+		let config_where_clause = &def.config.where_clause;
+
+		quote::quote_spanned!(span =>
+			#[doc(hidden)]
+			#[allow(non_camel_case_types)]
+			#visibility struct #name;
+
+			impl<#type_impl_gen> #frame_support::traits::Get<Result<#value_ty, #error_path>>
+				for #name
+				#config_where_clause
+			{
+				fn get() -> Result<#value_ty, #error_path> {
+					Err(<#error_path>::#variant_name)
+				}
+			}
+		)
+	});
+
 	let mut where_clauses = vec![&def.config.where_clause];
 	where_clauses.extend(def.storages.iter().map(|storage| &storage.where_clause));
 	let completed_where_clause = super::merge_where_clauses(&where_clauses);
@@ -489,5 +663,6 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 
 		#( #getters )*
 		#( #prefix_structs )*
+		#( #on_empty_structs )*
 	)
 }