don't allow ? in non-result v2 benchmark function defs, require either Result or () return type

frame/balances/src/benchmarking.rs

@@ -225,7 +225,7 @@ mod benchmarks {
 	}
 
 	#[benchmark]
-	fn force_unreserve() {
+	fn force_unreserve() -> Result<(), BenchmarkError> {
 		let user: T::AccountId = account("user", 0, SEED);
 		let user_lookup = T::Lookup::unlookup(user.clone());
 

frame/benchmarking/src/lib.rs

@@ -238,9 +238,10 @@ pub mod v2 {
 	};
 
 	// Used in #[benchmark] implementation to ensure that benchmark function arguments
-	// implement [`ParamRange`].
+	// implement [`ParamRange`] and benchmark function return types are compatible with
+	// `BenchmarkResult`
 	#[doc(hidden)]
-	pub use static_assertions::assert_impl_all;
+	pub use static_assertions::{assert_impl_all, assert_type_eq_all};
 
 	/// Used by the new benchmarking code to specify that a benchmarking variable is linear
 	/// over some specified range, i.e. `Linear<0, 1_000>` means that the corresponding variable
@@ -272,4 +273,7 @@ pub mod v2 {
 			B
 		}
 	}
+
+	/// The return type for benchmark function definitions that make use of the `?` operator
+	pub type BenchmarkResult = Result<(), BenchmarkError>;
 }

frame/support/procedural/src/benchmark.rs

@@ -28,19 +28,23 @@ use syn::{
 	punctuated::Punctuated,
 	spanned::Spanned,
 	token::{Colon2, Comma, Gt, Lt, Paren},
-	Attribute, Error, Expr, ExprBlock, ExprCall, ExprPath, FnArg, Item, ItemFn, ItemMod, LitInt,
-	Pat, Path, PathArguments, PathSegment, Result, Stmt, Token, Type, WhereClause,
+	visit::{self, Visit},
+	Attribute, Error, Expr, ExprBlock, ExprCall, ExprPath, ExprTry, FnArg, Item, ItemFn, ItemMod,
+	LitInt, Pat, Path, PathArguments, PathSegment, Result, ReturnType, Stmt, Token, Type,
+	WhereClause,
 };
 
 mod keywords {
 	use syn::custom_keyword;
 
 	custom_keyword!(benchmark);
 	custom_keyword!(benchmarks);
+	custom_keyword!(BenchmarkError);
 	custom_keyword!(block);
 	custom_keyword!(extra);
 	custom_keyword!(extrinsic_call);
 	custom_keyword!(skip_meta);
+	custom_keyword!(Result);
 }
 
 /// This represents the raw parsed data for a param definition such as `x: Linear<10, 20>`.
@@ -62,6 +66,23 @@ struct RangeArgs {
 	_gt_token: Gt,
 }
 
+struct ExprTryVisitor<'ast> {
+	results: Vec<&'ast ExprTry>,
+}
+
+impl<'ast> ExprTryVisitor<'ast> {
+	pub fn new() -> ExprTryVisitor<'ast> {
+		ExprTryVisitor { results: Vec::new() }
+	}
+}
+
+impl<'ast> Visit<'ast> for ExprTryVisitor<'ast> {
+	fn visit_expr_try(&mut self, expr_try: &'ast ExprTry) {
+		self.results.push(expr_try);
+		visit::visit_expr_try(self, expr_try);
+	}
+}
+
 #[derive(Clone, Debug)]
 struct BenchmarkAttrs {
 	skip_meta: bool,
@@ -145,16 +166,38 @@ struct BenchmarkDef {
 	setup_stmts: Vec<Stmt>,
 	call_def: BenchmarkCallDef,
 	verify_stmts: Vec<Stmt>,
+	return_type: Option<Type>, // `None` if `()` return type
 	extra: bool,
 	skip_meta: bool,
 }
 
+/// Parses the return type of a benchmark function definition
+fn parse_return_type(item_fn: &ItemFn) -> Result<Option<Type>> {
+	match &item_fn.sig.output {
+		ReturnType::Default => {
+			// return type is `()` which is permitted
+			// check for question mark ops as these aren't allowed here
+			let mut visitor = ExprTryVisitor::new();
+			visitor.visit_item_fn(item_fn);
+			if let Some(expr) = visitor.results.first() {
+				return Err(Error::new(
+					expr.question_token.span,
+					"The `?` operator can only be used in benchmark functions that specify `Result<(), BenchmarkError>`as their return type"
+				))
+			}
+			Ok(None)
+		},
+		ReturnType::Type(_, typ) => Ok(Some(*typ.clone())),
+	}
+}
+
 impl BenchmarkDef {
 	/// Constructs a [`BenchmarkDef`] by traversing an existing [`ItemFn`] node.
 	pub fn from(item_fn: &ItemFn, extra: bool, skip_meta: bool) -> Result<BenchmarkDef> {
-		let mut params: Vec<ParamDef> = Vec::new();
+		let return_type = parse_return_type(item_fn)?;
 
 		// parse params such as "x: Linear<0, 1>"
+		let mut params: Vec<ParamDef> = Vec::new();
 		for arg in &item_fn.sig.inputs {
 			let invalid_param = |span| {
 				return Err(Error::new(span, "Invalid benchmark function param. A valid example would be `x: Linear<5, 10>`.", ))
@@ -253,6 +296,7 @@ impl BenchmarkDef {
 			verify_stmts: Vec::from(&item_fn.block.stmts[(i + 1)..item_fn.block.stmts.len()]),
 			extra,
 			skip_meta,
+			return_type,
 		})
 	}
 }
@@ -661,6 +705,12 @@ fn expand_benchmark(
 		true => quote!(T: Config<I>, I: 'static),
 	};
 
+	let return_type_assertion = match benchmark_def.return_type {
+		Some(return_type) =>
+			quote!(#home::assert_type_eq_all!(#return_type, #home::BenchmarkResult);),
+		_ => quote!(),
+	};
+
 	let (pre_call, post_call) = match benchmark_def.call_def {
 		BenchmarkCallDef::ExtrinsicCall { origin, expr_call, attr_span: _ } => {
 			let mut expr_call = expr_call.clone();
@@ -734,6 +784,8 @@ fn expand_benchmark(
 			#home::assert_impl_all!(#param_types: #home::ParamRange);
 		)*
 
+		#return_type_assertion
+
 		#[allow(non_camel_case_types)]
 		struct #name;
 

frame/support/test/tests/benchmark_ui/bad_return_non_result.rs

@@ -0,0 +1,16 @@
+use frame_benchmarking::v2::*;
+#[allow(unused_imports)]
+use frame_support_test::Config;
+
+#[benchmarks]
+mod benchmarks {
+	use super::*;
+
+	#[benchmark]
+	fn bench() -> Option<String> {
+		#[block]
+		{}
+	}
+}
+
+fn main() {}

frame/support/test/tests/benchmark_ui/bad_return_non_result.stderr

@@ -0,0 +1,22 @@
+error[E0271]: type mismatch resolving `<std::option::Option<std::string::String> as TypeEq>::This == Result<(), BenchmarkError>`
+  --> tests/benchmark_ui/bad_return_non_result.rs:10:16
+   |
+10 |     fn bench() -> Option<String> {
+   |                   ^^^^^^^^^^^^^^ type mismatch resolving `<std::option::Option<std::string::String> as TypeEq>::This == Result<(), BenchmarkError>`
+   |
+note: expected this to be `Result<(), BenchmarkError>`
+  --> tests/benchmark_ui/bad_return_non_result.rs:5:1
+   |
+5  | #[benchmarks]
+   | ^^^^^^^^^^^^^
+   = note: expected enum `Result<(), BenchmarkError>`
+              found enum `std::option::Option<std::string::String>`
+note: required by a bound in `assert_type_eq_all`
+  --> tests/benchmark_ui/bad_return_non_result.rs:5:1
+   |
+5  | #[benchmarks]
+   | ^^^^^^^^^^^^^
+   | |
+   | required by a bound in this
+   | required by this bound in `assert_type_eq_all`
+   = note: this error originates in the macro `frame_benchmarking::v2::assert_type_eq_all` which comes from the expansion of the attribute macro `benchmarks` (in Nightly builds, run with -Z macro-backtrace for more info)

frame/support/test/tests/benchmark_ui/bad_return_wrong_result.rs

@@ -0,0 +1,16 @@
+use frame_benchmarking::v2::*;
+#[allow(unused_imports)]
+use frame_support_test::Config;
+
+#[benchmarks]
+mod benchmarks {
+	use super::*;
+
+	#[benchmark]
+	fn bench() -> Result<String, BenchmarkError> {
+		#[block]
+		{}
+	}
+}
+
+fn main() {}

frame/support/test/tests/benchmark_ui/bad_return_wrong_result.stderr

@@ -0,0 +1,22 @@
+error[E0271]: type mismatch resolving `<Result<std::string::String, frame_benchmarking::BenchmarkError> as TypeEq>::This == Result<(), frame_benchmarking::BenchmarkError>`
+  --> tests/benchmark_ui/bad_return_wrong_result.rs:10:16
+   |
+10 |     fn bench() -> Result<String, BenchmarkError> {
+   |                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ type mismatch resolving `<Result<std::string::String, frame_benchmarking::BenchmarkError> as TypeEq>::This == Result<(), frame_benchmarking::BenchmarkError>`
+   |
+note: expected this to be `Result<(), frame_benchmarking::BenchmarkError>`
+  --> tests/benchmark_ui/bad_return_wrong_result.rs:5:1
+   |
+5  | #[benchmarks]
+   | ^^^^^^^^^^^^^
+   = note: expected enum `Result<(), _>`
+              found enum `Result<std::string::String, _>`
+note: required by a bound in `assert_type_eq_all`
+  --> tests/benchmark_ui/bad_return_wrong_result.rs:5:1
+   |
+5  | #[benchmarks]
+   | ^^^^^^^^^^^^^
+   | |
+   | required by a bound in this
+   | required by this bound in `assert_type_eq_all`
+   = note: this error originates in the macro `frame_benchmarking::v2::assert_type_eq_all` which comes from the expansion of the attribute macro `benchmarks` (in Nightly builds, run with -Z macro-backtrace for more info)

frame/support/test/tests/benchmark_ui/invalid_use_of_question_op.rs

@@ -0,0 +1,21 @@
+use frame_benchmarking::v2::*;
+#[allow(unused_imports)]
+use frame_support_test::Config;
+
+#[benchmarks]
+mod benchmarks {
+	use super::*;
+
+	fn something() -> BenchmarkResult {
+		Ok(())
+	}
+
+	#[benchmark]
+	fn bench() {
+		something()?;
+		#[block]
+		{}
+	}
+}
+
+fn main() {}

frame/support/test/tests/benchmark_ui/invalid_use_of_question_op.stderr

@@ -0,0 +1,5 @@
+error: The `?` operator can only be used in benchmark functions that specify `Result<(), BenchmarkError>`as their return type
+  --> tests/benchmark_ui/invalid_use_of_question_op.rs:15:14
+   |
+15 |         something()?;
+   |                    ^

frame/support/test/tests/benchmark_ui/pass/valid_with_question_compact.rs

@@ -0,0 +1,21 @@
+use frame_benchmarking::v2::*;
+#[allow(unused_imports)]
+use frame_support_test::Config;
+
+#[benchmarks]
+mod benchmarks {
+	use super::*;
+
+	fn something() -> Result<(), BenchmarkError> {
+		Ok(())
+	}
+
+	#[benchmark]
+	fn bench() -> BenchmarkResult {
+		something()?;
+		#[block]
+		{}
+	}
+}
+
+fn main() {}

frame/support/test/tests/benchmark_ui/pass/valid_with_question_verbose.rs

@@ -0,0 +1,21 @@
+use frame_benchmarking::v2::*;
+#[allow(unused_imports)]
+use frame_support_test::Config;
+
+#[benchmarks]
+mod benchmarks {
+	use super::*;
+
+	fn something() -> Result<(), BenchmarkError> {
+		Ok(())
+	}
+
+	#[benchmark]
+	fn bench() -> Result<(), BenchmarkError> {
+		something()?;
+		#[block]
+		{}
+	}
+}
+
+fn main() {}

frame/support/test/tests/benchmark_ui/question_op_in_non_result.rs

@@ -0,0 +1,15 @@
+use frame_benchmarking::v2::*;
+#[allow(unused_imports)]
+use frame_support_test::Config;
+
+#[benchmarks]
+mod benchmarks {
+	#[benchmark]
+	fn bench() {
+		something()?;
+		#[block]
+		{}
+	}
+}
+
+fn main() {}

frame/support/test/tests/benchmark_ui/question_op_in_non_result.stderr

@@ -0,0 +1,5 @@
+error: The `?` operator can only be used in benchmark functions that specify `Result<(), BenchmarkError>`as their return type
+ --> tests/benchmark_ui/question_op_in_non_result.rs:9:14
+  |
+9 |         something()?;
+  |                    ^