Improve mock_impl_runtime_apis! (paritytech#7370)

* Improve `mock_impl_runtime_apis!`

This adds a new attribute for functions being implemented in the
`mock_impl_runtime_apis!` macro, the `advanced` attribute. When this
attribute is given the user gets access to the `at` parameter and is
able to return a `Result`, instead of letting the macro generate this
stuff.

* Use the `at_param_name` directly

* Prevent clashing of `params`

primitives/api/proc-macro/src/mock_impl_runtime_apis.rs

@@ -24,16 +24,22 @@ use crate::utils::{
 
 use proc_macro2::{Span, TokenStream};
 
-use quote::quote;
+use quote::{quote, quote_spanned};
 
 use syn::{
 	spanned::Spanned, parse_macro_input, Ident, Type, ItemImpl, ImplItem, TypePath, parse_quote,
-	parse::{Parse, ParseStream, Result, Error}, fold::{self, Fold},
+	parse::{Parse, ParseStream, Result, Error}, fold::{self, Fold}, Attribute, Pat,
 };
 
 /// Unique identifier used to make the hidden includes unique for this macro.
 const HIDDEN_INCLUDES_ID: &str = "MOCK_IMPL_RUNTIME_APIS";
 
+/// The `advanced` attribute.
+///
+/// If this attribute is given to a function, the function gets access to the `BlockId` as first
+/// parameter and needs to return a `Result` with the appropiate error type.
+const ADVANCED_ATTRIBUTE: &str = "advanced";
+
 /// The structure used for parsing the runtime api implementations.
 struct RuntimeApiImpls {
 	impls: Vec<ItemImpl>,
@@ -65,10 +71,16 @@ fn implement_common_api_traits(
 
 	let error_type = error_type.map(|e| quote!(#e)).unwrap_or_else(|| quote!(String));
 
-	Ok(quote!(
+	// Quote using the span from `error_type` to generate nice error messages when the type is
+	// not implementing a trait or similar.
+	let api_error_ext = quote_spanned! { error_type.span() =>
 		impl #crate_::ApiErrorExt for #self_ty {
 			type Error = #error_type;
 		}
+	};
+
+	Ok(quote!(
+		#api_error_ext
 
 		impl #crate_::ApiExt<#block_type> for #self_ty {
 			type StateBackend = #crate_::InMemoryBackend<#crate_::HashFor<#block_type>>;
@@ -89,7 +101,7 @@ fn implement_common_api_traits(
 
 			fn has_api_with<A: #crate_::RuntimeApiInfo + ?Sized, P: Fn(u32) -> bool>(
 				&self,
-				at: &#crate_::BlockId<#block_type>,
+				_: &#crate_::BlockId<#block_type>,
 				pred: P,
 			) -> std::result::Result<bool, #error_type> where Self: Sized {
 				Ok(pred(A::VERSION))
@@ -153,6 +165,61 @@ fn implement_common_api_traits(
 	))
 }
 
+/// Returns if the advanced attribute is present in the given `attributes`.
+///
+/// If the attribute was found, it will be automatically removed from the vec.
+fn has_advanced_attribute(attributes: &mut Vec<Attribute>) -> bool {
+	let mut found = false;
+	attributes.retain(|attr| if attr.path.is_ident(ADVANCED_ATTRIBUTE) {
+		found = true;
+		false
+	} else {
+		true
+	});
+
+	found
+}
+
+/// Get the name and type of the `at` parameter that is passed to a runtime api function.
+///
+/// If `is_advanced` is `false`, the name is `_`.
+fn get_at_param_name(
+	is_advanced: bool,
+	param_names: &mut Vec<Pat>,
+	param_types_and_borrows: &mut Vec<(TokenStream, bool)>,
+	function_span: Span,
+	default_block_id_type: &TokenStream,
+) -> Result<(TokenStream, TokenStream)> {
+	if is_advanced {
+		if param_names.is_empty() {
+			return Err(Error::new(
+				function_span,
+				format!(
+					"If using the `{}` attribute, it is required that the function \
+					 takes at least one argument, the `BlockId`.",
+					ADVANCED_ATTRIBUTE,
+				),
+			))
+		}
+
+		// `param_names` and `param_types` have the same length, so if `param_names` is not empty
+		// `param_types` can not be empty as well.
+		let ptype_and_borrows = param_types_and_borrows.remove(0);
+		let span = ptype_and_borrows.1.span();
+		if !ptype_and_borrows.1 {
+			return Err(Error::new(
+				span,
+				"`BlockId` needs to be taken by reference and not by value!",
+			))
+		}
+
+		let name = param_names.remove(0);
+		Ok((quote!( #name ), ptype_and_borrows.0))
+	} else {
+		Ok((quote!( _ ), default_block_id_type.clone()))
+	}
+}
+
 /// Auxialiry structure to fold a runtime api trait implementation into the expected format.
 ///
 /// This renames the methods, changes the method parameters and extracts the error type.
@@ -170,8 +237,10 @@ impl<'a> Fold for FoldRuntimeApiImpl<'a> {
 	fn fold_impl_item_method(&mut self, mut input: syn::ImplItemMethod) -> syn::ImplItemMethod {
 		let block = {
 			let crate_ = generate_crate_access(HIDDEN_INCLUDES_ID);
+			let is_advanced = has_advanced_attribute(&mut input.attrs);
+			let mut errors = Vec::new();
 
-			let (param_names, param_types, error) = match extract_parameter_names_types_and_borrows(
+			let (mut param_names, mut param_types_and_borrows) = match extract_parameter_names_types_and_borrows(
 				&input.sig,
 				AllowSelfRefInParameters::YesButIgnore,
 			) {
@@ -180,49 +249,81 @@ impl<'a> Fold for FoldRuntimeApiImpl<'a> {
 					res.iter().map(|v| {
 						let ty = &v.1;
 						let borrow = &v.2;
-						quote!( #borrow #ty )
+						(quote_spanned!(ty.span() => #borrow #ty ), v.2.is_some())
 					}).collect::<Vec<_>>(),
-					None
 				),
-				Err(e) => (Vec::new(), Vec::new(), Some(e.to_compile_error())),
+				Err(e) => {
+					errors.push(e.to_compile_error());
+
+					(Default::default(), Default::default())
+				}
 			};
 
 			let block_type = &self.block_type;
+			let block_id_type = quote!( &#crate_::BlockId<#block_type> );
+
+			let (at_param_name, block_id_type) = match get_at_param_name(
+				is_advanced,
+				&mut param_names,
+				&mut param_types_and_borrows,
+				input.span(),
+				&block_id_type,
+			) {
+				Ok(res) => res,
+				Err(e) => {
+					errors.push(e.to_compile_error());
+					(quote!( _ ), block_id_type)
+				}
+			};
 
+			let param_types = param_types_and_borrows.iter().map(|v| &v.0);
 			// Rewrite the input parameters.
 			input.sig.inputs = parse_quote! {
 				&self,
-				_: &#crate_::BlockId<#block_type>,
+				#at_param_name: #block_id_type,
 				_: #crate_::ExecutionContext,
-				params: Option<( #( #param_types ),* )>,
+				___params___sp___api___: Option<( #( #param_types ),* )>,
 				_: Vec<u8>,
 			};
 
 			input.sig.ident = generate_method_runtime_api_impl_name(
 				&self.impl_trait,
 				&input.sig.ident,
 			);
-			let ret_type = return_type_extract_type(&input.sig.output);
 
-			// Generate the correct return type.
-			input.sig.output = parse_quote!(
-				-> std::result::Result<#crate_::NativeOrEncoded<#ret_type>, Self::Error>
-			);
+			// When using advanced, the user needs to declare the correct return type on its own,
+			// otherwise do it for the user.
+			if !is_advanced {
+				let ret_type = return_type_extract_type(&input.sig.output);
+
+				// Generate the correct return type.
+				input.sig.output = parse_quote!(
+					-> std::result::Result<#crate_::NativeOrEncoded<#ret_type>, Self::Error>
+				);
+			}
 
 			let orig_block = input.block.clone();
 
+			let construct_return_value = if is_advanced {
+				quote!( (move || #orig_block)() )
+			} else {
+				quote! {
+					let __fn_implementation__ = move || #orig_block;
+
+					Ok(#crate_::NativeOrEncoded::Native(__fn_implementation__()))
+				}
+			};
+
 			// Generate the new method implementation that calls into the runtime.
 			parse_quote!(
 				{
 					// Get the error to the user (if we have one).
-					#error
+					#( #errors )*
 
-					let (#( #param_names ),*) = params
+					let (#( #param_names ),*) = ___params___sp___api___
 						.expect("Mocked runtime apis don't support calling deprecated api versions");
 
-					let __fn_implementation__ = move || #orig_block;
-
-					Ok(#crate_::NativeOrEncoded::Native(__fn_implementation__()))
+					#construct_return_value
 				}
 			)
 		};
@@ -240,10 +341,17 @@ impl<'a> Fold for FoldRuntimeApiImpl<'a> {
 				if ty.ident == "Error" {
 					if let Some(error_type) = self.error_type {
 						if *error_type != ty.ty {
-							let error = Error::new(
+							let mut error = Error::new(
 								ty.span(),
 								"Error type can not change between runtime apis",
 							);
+							let error_first = Error::new(
+								error_type.span(),
+								"First error type was declared here."
+							);
+
+							error.combine(error_first);
+
 							ImplItem::Verbatim(error.to_compile_error())
 						} else {
 							ImplItem::Verbatim(Default::default())

primitives/api/src/lib.rs

@@ -241,20 +241,18 @@ pub use sp_api_proc_macro::impl_runtime_apis;
 
 /// Mocks given trait implementations as runtime apis.
 ///
-/// Accepts similar syntax as [`impl_runtime_apis!`](macro.impl_runtime_apis.html) and generates
+/// Accepts similar syntax as [`impl_runtime_apis!`] and generates
 /// simplified mock implementations of the given runtime apis. The difference in syntax is that the
 /// trait does not need to be referenced by a qualified path, methods accept the `&self` parameter
-/// and the error type can be specified as associated type. If no error type is specified `String`
+/// and the error type can be specified as associated type. If no error type is specified [`String`]
 /// is used as error type.
 ///
-/// Besides implementing the given traits, the [`Core`], [`ApiExt`] and [`ApiErrorExt`] are
-/// implemented automatically.
+/// Besides implementing the given traits, the [`Core`](sp_api::Core), [`ApiExt`](sp_api::ApiExt)
+/// and [`ApiErrorExt`](sp_api::ApiErrorExt) are implemented automatically.
 ///
 /// # Example
 ///
 /// ```rust
-/// use sp_version::create_runtime_str;
-/// #
 /// # use sp_runtime::traits::Block as BlockT;
 /// # use sp_test_primitives::Block;
 /// #
@@ -270,7 +268,6 @@ pub use sp_api_proc_macro::impl_runtime_apis;
 /// #        fn build_block() -> Block;
 /// #     }
 /// # }
-///
 /// struct MockApi {
 ///     balance: u64,
 /// }
@@ -301,6 +298,59 @@ pub use sp_api_proc_macro::impl_runtime_apis;
 ///
 /// # fn main() {}
 /// ```
+///
+/// # `advanced` attribute
+///
+/// This attribute can be placed above individual function in the mock implementation to request
+/// more control over the function declaration. From the client side each runtime api function is
+/// called with the `at` parameter that is a [`BlockId`](sp_api::BlockId). When using the `advanced`
+/// attribute, the macro expects that the first parameter of the function is this `at` parameter.
+/// Besides that the macro also doesn't do the automatic return value rewrite, which means that full
+/// return value must be specified. The full return value is constructed like
+/// [`Result`]`<`[`NativeOrEncoded`](sp_api::NativeOrEncoded)`<ReturnValue>, Error>` while
+/// `ReturnValue` being the return value that is specified in the trait declaration.
+///
+/// ## Example
+/// ```rust
+/// # use sp_runtime::{traits::Block as BlockT, generic::BlockId};
+/// # use sp_test_primitives::Block;
+/// # use sp_core::NativeOrEncoded;
+/// #
+/// # sp_api::decl_runtime_apis! {
+/// #     /// Declare the api trait.
+/// #     pub trait Balance {
+/// #         /// Get the balance.
+/// #         fn get_balance() -> u64;
+/// #         /// Set the balance.
+/// #         fn set_balance(val: u64);
+/// #     }
+/// # }
+/// struct MockApi {
+///     balance: u64,
+/// }
+///
+/// sp_api::mock_impl_runtime_apis! {
+///     impl Balance<Block> for MockApi {
+///         type Error = String;
+///         #[advanced]
+///         fn get_balance(&self, at: &BlockId<Block>) -> Result<NativeOrEncoded<u64>, String> {
+///             println!("Being called at: {}", at);
+///
+///             Ok(self.balance.into())
+///         }
+///         #[advanced]
+///         fn set_balance(at: &BlockId<Block>, val: u64) -> Result<NativeOrEncoded<()>, String> {
+///             if let BlockId::Number(1) = at {
+///                 println!("Being called to set balance to: {}", val);
+///             }
+///
+///             Ok(().into())
+///         }
+///     }
+/// }
+///
+/// # fn main() {}
+/// ```
 pub use sp_api_proc_macro::mock_impl_runtime_apis;
 
 /// A type that records all accessed trie nodes and generates a proof out of it.

primitives/api/test/tests/decl_and_impl.rs

@@ -19,9 +19,8 @@ use sp_api::{
 	RuntimeApiInfo, decl_runtime_apis, impl_runtime_apis, mock_impl_runtime_apis,
 	ApiExt,
 };
-
 use sp_runtime::{traits::{GetNodeBlockType, Block as BlockT}, generic::BlockId};
-
+use sp_core::NativeOrEncoded;
 use substrate_test_runtime_client::runtime::Block;
 use sp_blockchain::Result;
 
@@ -103,9 +102,20 @@ mock_impl_runtime_apis! {
 			unimplemented!()
 		}
 
-		fn same_name() {}
+		#[advanced]
+		fn same_name(_: &BlockId<Block>) -> std::result::Result<NativeOrEncoded<()>, String> {
+			Ok(().into())
+		}
 
-		fn wild_card(_: u32) {}
+		#[advanced]
+		fn wild_card(at: &BlockId<Block>, _: u32) -> std::result::Result<NativeOrEncoded<()>, String> {
+			if let BlockId::Number(1337) = at {
+				// yeah
+				Ok(().into())
+			} else {
+				Err("Ohh noooo".into())
+			}
+		}
 	}
 
 	impl ApiWithCustomVersion<Block> for MockApi {
@@ -180,3 +190,12 @@ fn mock_runtime_api_panics_on_calling_old_version() {
 	#[allow(deprecated)]
 	let _ = mock.same_name_before_version_2(&BlockId::Number(0));
 }
+
+#[test]
+fn mock_runtime_api_works_with_advanced() {
+	let mock = MockApi { block: None };
+
+	Api::<Block>::same_name(&mock, &BlockId::Number(0)).unwrap();
+	mock.wild_card(&BlockId::Number(1337), 1).unwrap();
+	assert_eq!(String::from("Ohh noooo"), mock.wild_card(&BlockId::Number(1336), 1).unwrap_err());
+}