frame_support_procedural/
benchmark.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! Home of the parsing and expansion code for the new pallet benchmarking syntax
19
20use derive_syn_parse::Parse;
21use frame_support_procedural_tools::generate_access_from_frame_or_crate;
22use proc_macro::TokenStream;
23use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
24use quote::{quote, ToTokens};
25use syn::{
26	parse::{Nothing, ParseStream},
27	parse_quote,
28	punctuated::Punctuated,
29	spanned::Spanned,
30	token::{Comma, Gt, Lt, PathSep},
31	Attribute, Error, Expr, ExprBlock, ExprCall, ExprPath, FnArg, Item, ItemFn, ItemMod, Pat, Path,
32	PathArguments, PathSegment, Result, ReturnType, Signature, Stmt, Token, Type, TypePath,
33	Visibility, WhereClause,
34};
35
36mod keywords {
37	use syn::custom_keyword;
38
39	custom_keyword!(benchmark);
40	custom_keyword!(benchmarks);
41	custom_keyword!(block);
42	custom_keyword!(extra);
43	custom_keyword!(pov_mode);
44	custom_keyword!(extrinsic_call);
45	custom_keyword!(skip_meta);
46	custom_keyword!(BenchmarkError);
47	custom_keyword!(Result);
48	custom_keyword!(MaxEncodedLen);
49	custom_keyword!(Measured);
50	custom_keyword!(Ignored);
51
52	pub const BENCHMARK_TOKEN: &str = stringify!(benchmark);
53	pub const BENCHMARKS_TOKEN: &str = stringify!(benchmarks);
54}
55
56/// This represents the raw parsed data for a param definition such as `x: Linear<10, 20>`.
57#[derive(Clone)]
58struct ParamDef {
59	name: String,
60	_typ: Type,
61	start: syn::GenericArgument,
62	end: syn::GenericArgument,
63}
64
65/// Allows easy parsing of the `<10, 20>` component of `x: Linear<10, 20>`.
66#[derive(Parse)]
67struct RangeArgs {
68	_lt_token: Lt,
69	start: syn::GenericArgument,
70	_comma: Comma,
71	end: syn::GenericArgument,
72	_trailing_comma: Option<Comma>,
73	_gt_token: Gt,
74}
75
76#[derive(Clone, Debug)]
77struct BenchmarkAttrs {
78	skip_meta: bool,
79	extra: bool,
80	pov_mode: Option<PovModeAttr>,
81}
82
83/// Represents a single benchmark option
84enum BenchmarkAttr {
85	Extra,
86	SkipMeta,
87	/// How the PoV should be measured.
88	PoV(PovModeAttr),
89}
90
91impl syn::parse::Parse for PovModeAttr {
92	fn parse(input: ParseStream) -> Result<Self> {
93		let _pov: keywords::pov_mode = input.parse()?;
94		let _eq: Token![=] = input.parse()?;
95		let root = PovEstimationMode::parse(input)?;
96
97		let mut maybe_content = None;
98		let _ = || -> Result<()> {
99			let content;
100			syn::braced!(content in input);
101			maybe_content = Some(content);
102			Ok(())
103		}();
104
105		let per_key = match maybe_content {
106			Some(content) => {
107				let per_key = Punctuated::<PovModeKeyAttr, Token![,]>::parse_terminated(&content)?;
108				per_key.into_iter().collect()
109			},
110			None => Vec::new(),
111		};
112
113		Ok(Self { root, per_key })
114	}
115}
116
117impl syn::parse::Parse for BenchmarkAttr {
118	fn parse(input: ParseStream) -> Result<Self> {
119		let lookahead = input.lookahead1();
120		if lookahead.peek(keywords::extra) {
121			let _extra: keywords::extra = input.parse()?;
122			Ok(BenchmarkAttr::Extra)
123		} else if lookahead.peek(keywords::skip_meta) {
124			let _skip_meta: keywords::skip_meta = input.parse()?;
125			Ok(BenchmarkAttr::SkipMeta)
126		} else if lookahead.peek(keywords::pov_mode) {
127			PovModeAttr::parse(input).map(BenchmarkAttr::PoV)
128		} else {
129			Err(lookahead.error())
130		}
131	}
132}
133
134/// A `#[pov_mode = .. { .. }]` attribute.
135#[derive(Debug, Clone)]
136struct PovModeAttr {
137	/// The root mode for this benchmarks.
138	root: PovEstimationMode,
139	/// The pov-mode for a specific key. This overwrites `root` for this key.
140	per_key: Vec<PovModeKeyAttr>,
141}
142
143/// A single key-value pair inside the `{}` of a `#[pov_mode = .. { .. }]` attribute.
144#[derive(Debug, Clone, derive_syn_parse::Parse)]
145struct PovModeKeyAttr {
146	/// A specific storage key for which to set the PoV mode.
147	key: Path,
148	_underscore: Token![:],
149	/// The PoV mode for this key.
150	mode: PovEstimationMode,
151}
152
153/// How the PoV should be estimated.
154#[derive(Debug, Eq, PartialEq, Clone, Copy)]
155pub enum PovEstimationMode {
156	/// Use the maximal encoded length as provided by [`codec::MaxEncodedLen`].
157	MaxEncodedLen,
158	/// Measure the accessed value size in the pallet benchmarking and add some trie overhead.
159	Measured,
160	/// Do not estimate the PoV size for this storage item or benchmark.
161	Ignored,
162}
163
164impl syn::parse::Parse for PovEstimationMode {
165	fn parse(input: ParseStream) -> Result<Self> {
166		let lookahead = input.lookahead1();
167		if lookahead.peek(keywords::MaxEncodedLen) {
168			let _max_encoded_len: keywords::MaxEncodedLen = input.parse()?;
169			return Ok(PovEstimationMode::MaxEncodedLen)
170		} else if lookahead.peek(keywords::Measured) {
171			let _measured: keywords::Measured = input.parse()?;
172			return Ok(PovEstimationMode::Measured)
173		} else if lookahead.peek(keywords::Ignored) {
174			let _ignored: keywords::Ignored = input.parse()?;
175			return Ok(PovEstimationMode::Ignored)
176		} else {
177			return Err(lookahead.error())
178		}
179	}
180}
181
182impl ToString for PovEstimationMode {
183	fn to_string(&self) -> String {
184		match self {
185			PovEstimationMode::MaxEncodedLen => "MaxEncodedLen".into(),
186			PovEstimationMode::Measured => "Measured".into(),
187			PovEstimationMode::Ignored => "Ignored".into(),
188		}
189	}
190}
191
192impl quote::ToTokens for PovEstimationMode {
193	fn to_tokens(&self, tokens: &mut TokenStream2) {
194		match self {
195			PovEstimationMode::MaxEncodedLen => tokens.extend(quote!(MaxEncodedLen)),
196			PovEstimationMode::Measured => tokens.extend(quote!(Measured)),
197			PovEstimationMode::Ignored => tokens.extend(quote!(Ignored)),
198		}
199	}
200}
201
202impl syn::parse::Parse for BenchmarkAttrs {
203	fn parse(input: ParseStream) -> syn::Result<Self> {
204		let mut extra = false;
205		let mut skip_meta = false;
206		let mut pov_mode = None;
207		let args = Punctuated::<BenchmarkAttr, Token![,]>::parse_terminated(&input)?;
208
209		for arg in args.into_iter() {
210			match arg {
211				BenchmarkAttr::Extra => {
212					if extra {
213						return Err(input.error("`extra` can only be specified once"))
214					}
215					extra = true;
216				},
217				BenchmarkAttr::SkipMeta => {
218					if skip_meta {
219						return Err(input.error("`skip_meta` can only be specified once"))
220					}
221					skip_meta = true;
222				},
223				BenchmarkAttr::PoV(mode) => {
224					if pov_mode.is_some() {
225						return Err(input.error("`pov_mode` can only be specified once"))
226					}
227					pov_mode = Some(mode);
228				},
229			}
230		}
231		Ok(BenchmarkAttrs { extra, skip_meta, pov_mode })
232	}
233}
234
235/// Represents the parsed extrinsic call for a benchmark
236#[derive(Clone)]
237enum BenchmarkCallDef {
238	ExtrinsicCall { origin: Expr, expr_call: ExprCall, attr_span: Span }, // #[extrinsic_call]
239	Block { block: ExprBlock, attr_span: Span },                          // #[block]
240}
241
242impl BenchmarkCallDef {
243	/// Returns the `span()` for attribute
244	fn attr_span(&self) -> Span {
245		match self {
246			BenchmarkCallDef::ExtrinsicCall { origin: _, expr_call: _, attr_span } => *attr_span,
247			BenchmarkCallDef::Block { block: _, attr_span } => *attr_span,
248		}
249	}
250}
251
252/// Represents a parsed `#[benchmark]` or `#[instance_benchmark]` item.
253#[derive(Clone)]
254struct BenchmarkDef {
255	params: Vec<ParamDef>,
256	setup_stmts: Vec<Stmt>,
257	call_def: BenchmarkCallDef,
258	verify_stmts: Vec<Stmt>,
259	last_stmt: Option<Stmt>,
260	fn_sig: Signature,
261	fn_vis: Visibility,
262	fn_attrs: Vec<Attribute>,
263}
264
265/// used to parse something compatible with `Result<T, E>`
266#[derive(Parse)]
267struct ResultDef {
268	_result_kw: keywords::Result,
269	_lt: Token![<],
270	unit: Type,
271	_comma: Comma,
272	e_type: TypePath,
273	_gt: Token![>],
274}
275
276/// Ensures that `ReturnType` is a `Result<(), BenchmarkError>`, if specified
277fn ensure_valid_return_type(item_fn: &ItemFn) -> Result<()> {
278	if let ReturnType::Type(_, typ) = &item_fn.sig.output {
279		let non_unit = |span| return Err(Error::new(span, "expected `()`"));
280		let Type::Path(TypePath { path, qself: _ }) = &**typ else {
281			return Err(Error::new(
282					typ.span(),
283					"Only `Result<(), BenchmarkError>` or a blank return type is allowed on benchmark function definitions",
284				))
285		};
286		let seg = path
287			.segments
288			.last()
289			.expect("to be parsed as a TypePath, it must have at least one segment; qed");
290		let res: ResultDef = syn::parse2(seg.to_token_stream())?;
291		// ensure T in Result<T, E> is ()
292		let Type::Tuple(tup) = res.unit else { return non_unit(res.unit.span()) };
293		if !tup.elems.is_empty() {
294			return non_unit(tup.span())
295		}
296		let TypePath { path, qself: _ } = res.e_type;
297		let seg = path
298			.segments
299			.last()
300			.expect("to be parsed as a TypePath, it must have at least one segment; qed");
301		syn::parse2::<keywords::BenchmarkError>(seg.to_token_stream())?;
302	}
303	Ok(())
304}
305
306/// Ensure that the passed statements do not contain any forbidden variable names
307fn ensure_no_forbidden_variable_names(stmts: &[Stmt]) -> Result<()> {
308	const FORBIDDEN_VAR_NAMES: [&str; 2] = ["recording", "verify"];
309	for stmt in stmts {
310		let Stmt::Local(l) = stmt else { continue };
311		let Pat::Ident(ident) = &l.pat else { continue };
312		if FORBIDDEN_VAR_NAMES.contains(&ident.ident.to_string().as_str()) {
313			return Err(Error::new(
314				ident.span(),
315				format!(
316					"Variables {FORBIDDEN_VAR_NAMES:?} are reserved for benchmarking internals.",
317				),
318			));
319		}
320	}
321	Ok(())
322}
323
324/// Parses params such as `x: Linear<0, 1>`
325fn parse_params(item_fn: &ItemFn) -> Result<Vec<ParamDef>> {
326	let mut params: Vec<ParamDef> = Vec::new();
327	for arg in &item_fn.sig.inputs {
328		let invalid_param = |span| {
329			return Err(Error::new(
330				span,
331				"Invalid benchmark function param. A valid example would be `x: Linear<5, 10>`.",
332			))
333		};
334
335		let FnArg::Typed(arg) = arg else { return invalid_param(arg.span()) };
336		let Pat::Ident(ident) = &*arg.pat else { return invalid_param(arg.span()) };
337
338		// check param name
339		let var_span = ident.span();
340		let invalid_param_name = || {
341			return Err(Error::new(
342					var_span,
343					"Benchmark parameter names must consist of a single lowercase letter (a-z) and no other characters.",
344				));
345		};
346		let name = ident.ident.to_token_stream().to_string();
347		if name.len() > 1 {
348			return invalid_param_name()
349		};
350		let Some(name_char) = name.chars().next() else { return invalid_param_name() };
351		if !name_char.is_alphabetic() || !name_char.is_lowercase() {
352			return invalid_param_name()
353		}
354
355		// parse type
356		let typ = &*arg.ty;
357		let Type::Path(tpath) = typ else { return invalid_param(typ.span()) };
358		let Some(segment) = tpath.path.segments.last() else { return invalid_param(typ.span()) };
359		let args = segment.arguments.to_token_stream().into();
360		let Ok(args) = syn::parse::<RangeArgs>(args) else { return invalid_param(typ.span()) };
361
362		params.push(ParamDef { name, _typ: typ.clone(), start: args.start, end: args.end });
363	}
364	Ok(params)
365}
366
367/// Used in several places where the `#[extrinsic_call]` or `#[body]` annotation is missing
368fn missing_call<T>(item_fn: &ItemFn) -> Result<T> {
369	return Err(Error::new(
370		item_fn.block.brace_token.span.join(),
371		"No valid #[extrinsic_call] or #[block] annotation could be found in benchmark function body."
372	));
373}
374
375/// Finds the `BenchmarkCallDef` and its index (within the list of stmts for the fn) and
376/// returns them. Also handles parsing errors for invalid / extra call defs. AKA this is
377/// general handling for `#[extrinsic_call]` and `#[block]`
378fn parse_call_def(item_fn: &ItemFn) -> Result<(usize, BenchmarkCallDef)> {
379	// #[extrinsic_call] / #[block] handling
380	let call_defs = item_fn.block.stmts.iter().enumerate().filter_map(|(i, child)| {
381			if let Stmt::Expr(Expr::Call(expr_call), _semi) = child {
382				// #[extrinsic_call] case
383				expr_call.attrs.iter().enumerate().find_map(|(k, attr)| {
384					let segment = attr.path().segments.last()?;
385					let _: keywords::extrinsic_call = syn::parse(segment.ident.to_token_stream().into()).ok()?;
386					let mut expr_call = expr_call.clone();
387
388					// consume #[extrinsic_call] tokens
389					expr_call.attrs.remove(k);
390
391					// extract origin from expr_call
392					let Some(origin) = expr_call.args.first().cloned() else {
393						return Some(Err(Error::new(expr_call.span(), "Single-item extrinsic calls must specify their origin as the first argument.")))
394					};
395
396					Some(Ok((i, BenchmarkCallDef::ExtrinsicCall { origin, expr_call, attr_span: attr.span() })))
397				})
398			} else if let Stmt::Expr(Expr::Block(block), _) = child {
399				// #[block] case
400				block.attrs.iter().enumerate().find_map(|(k, attr)| {
401					let segment = attr.path().segments.last()?;
402					let _: keywords::block = syn::parse(segment.ident.to_token_stream().into()).ok()?;
403					let mut block = block.clone();
404
405					// consume #[block] tokens
406					block.attrs.remove(k);
407
408					Some(Ok((i, BenchmarkCallDef::Block { block, attr_span: attr.span() })))
409				})
410			} else {
411				None
412			}
413		}).collect::<Result<Vec<_>>>()?;
414	Ok(match &call_defs[..] {
415		[(i, call_def)] => (*i, call_def.clone()), // = 1
416		[] => return missing_call(item_fn),
417		_ =>
418			return Err(Error::new(
419				call_defs[1].1.attr_span(),
420				"Only one #[extrinsic_call] or #[block] attribute is allowed per benchmark.",
421			)),
422	})
423}
424
425impl BenchmarkDef {
426	/// Constructs a [`BenchmarkDef`] by traversing an existing [`ItemFn`] node.
427	pub fn from(item_fn: &ItemFn) -> Result<BenchmarkDef> {
428		let params = parse_params(item_fn)?;
429		ensure_valid_return_type(item_fn)?;
430		let (i, call_def) = parse_call_def(&item_fn)?;
431
432		let (verify_stmts, last_stmt) = match item_fn.sig.output {
433			ReturnType::Default =>
434			// no return type, last_stmt should be None
435				(Vec::from(&item_fn.block.stmts[(i + 1)..item_fn.block.stmts.len()]), None),
436			ReturnType::Type(_, _) => {
437				// defined return type, last_stmt should be Result<(), BenchmarkError>
438				// compatible and should not be included in verify_stmts
439				if i + 1 >= item_fn.block.stmts.len() {
440					return Err(Error::new(
441						item_fn.block.span(),
442						"Benchmark `#[block]` or `#[extrinsic_call]` item cannot be the \
443						last statement of your benchmark function definition if you have \
444						defined a return type. You should return something compatible \
445						with Result<(), BenchmarkError> (i.e. `Ok(())`) as the last statement \
446						or change your signature to a blank return type.",
447					))
448				}
449				let Some(stmt) = item_fn.block.stmts.last() else { return missing_call(item_fn) };
450				(
451					Vec::from(&item_fn.block.stmts[(i + 1)..item_fn.block.stmts.len() - 1]),
452					Some(stmt.clone()),
453				)
454			},
455		};
456
457		let setup_stmts = Vec::from(&item_fn.block.stmts[0..i]);
458		ensure_no_forbidden_variable_names(&setup_stmts)?;
459
460		Ok(BenchmarkDef {
461			params,
462			setup_stmts,
463			call_def,
464			verify_stmts,
465			last_stmt,
466			fn_sig: item_fn.sig.clone(),
467			fn_vis: item_fn.vis.clone(),
468			fn_attrs: item_fn.attrs.clone(),
469		})
470	}
471}
472
473/// Parses and expands a `#[benchmarks]` or `#[instance_benchmarks]` invocation
474pub fn benchmarks(
475	attrs: TokenStream,
476	tokens: TokenStream,
477	instance: bool,
478) -> syn::Result<TokenStream> {
479	let krate = generate_access_from_frame_or_crate("frame-benchmarking")?;
480	// gather module info
481	let module: ItemMod = syn::parse(tokens)?;
482	let mod_span = module.span();
483	let where_clause = match syn::parse::<Nothing>(attrs.clone()) {
484		Ok(_) => quote!(),
485		Err(_) => syn::parse::<WhereClause>(attrs)?.predicates.to_token_stream(),
486	};
487	let mod_vis = module.vis;
488	let mod_name = module.ident;
489
490	// consume #[benchmarks] attribute by excluding it from mod_attrs
491	let mod_attrs: Vec<&Attribute> = module
492		.attrs
493		.iter()
494		.filter(|attr| !attr.path().is_ident(keywords::BENCHMARKS_TOKEN))
495		.collect();
496
497	let mut benchmark_names: Vec<Ident> = Vec::new();
498	let mut extra_benchmark_names: Vec<Ident> = Vec::new();
499	let mut skip_meta_benchmark_names: Vec<Ident> = Vec::new();
500	// Map benchmarks to PoV modes.
501	let mut pov_modes = Vec::new();
502
503	let (_brace, mut content) =
504		module.content.ok_or(syn::Error::new(mod_span, "Module cannot be empty!"))?;
505
506	// find all function defs marked with #[benchmark]
507	let benchmark_fn_metas = content.iter_mut().filter_map(|stmt| {
508		// parse as a function def first
509		let Item::Fn(func) = stmt else { return None };
510
511		// find #[benchmark] attribute on function def
512		let benchmark_attr =
513			func.attrs.iter().find(|attr| attr.path().is_ident(keywords::BENCHMARK_TOKEN))?;
514
515		Some((benchmark_attr.clone(), func.clone(), stmt))
516	});
517
518	// parse individual benchmark defs and args
519	for (benchmark_attr, func, stmt) in benchmark_fn_metas {
520		// parse benchmark def
521		let benchmark_def = BenchmarkDef::from(&func)?;
522
523		// record benchmark name
524		let name = &func.sig.ident;
525		benchmark_names.push(name.clone());
526
527		// Check if we need to parse any args
528		if benchmark_attr.meta.require_path_only().is_err() {
529			// parse any args provided to #[benchmark]
530			let benchmark_attrs: BenchmarkAttrs = benchmark_attr.parse_args()?;
531
532			// record name sets
533			if benchmark_attrs.extra {
534				extra_benchmark_names.push(name.clone());
535			} else if benchmark_attrs.skip_meta {
536				skip_meta_benchmark_names.push(name.clone());
537			}
538
539			if let Some(mode) = benchmark_attrs.pov_mode {
540				let mut modes = Vec::new();
541				// We cannot expand strings here since it is no-std, but syn does not expand bytes.
542				let name = name.to_string();
543				let m = mode.root.to_string();
544				modes.push(quote!(("ALL".as_bytes().to_vec(), #m.as_bytes().to_vec())));
545
546				for attr in mode.per_key.iter() {
547					// syn always puts spaces in quoted paths:
548					let key = attr.key.clone().into_token_stream().to_string().replace(" ", "");
549					let mode = attr.mode.to_string();
550					modes.push(quote!((#key.as_bytes().to_vec(), #mode.as_bytes().to_vec())));
551				}
552
553				pov_modes.push(
554					quote!((#name.as_bytes().to_vec(), #krate::__private::vec![#(#modes),*])),
555				);
556			}
557		}
558
559		// expand benchmark
560		let expanded = expand_benchmark(benchmark_def, name, instance, where_clause.clone());
561
562		// replace original function def with expanded code
563		*stmt = Item::Verbatim(expanded);
564	}
565
566	// generics
567	let type_use_generics = match instance {
568		false => quote!(T),
569		true => quote!(T, I),
570	};
571	let type_impl_generics = match instance {
572		false => quote!(T: Config),
573		true => quote!(T: Config<I>, I: 'static),
574	};
575
576	let frame_system = generate_access_from_frame_or_crate("frame-system")?;
577
578	// benchmark name variables
579	let benchmark_names_str: Vec<String> = benchmark_names.iter().map(|n| n.to_string()).collect();
580	let extra_benchmark_names_str: Vec<String> =
581		extra_benchmark_names.iter().map(|n| n.to_string()).collect();
582	let skip_meta_benchmark_names_str: Vec<String> =
583		skip_meta_benchmark_names.iter().map(|n| n.to_string()).collect();
584	let mut selected_benchmark_mappings: Vec<TokenStream2> = Vec::new();
585	let mut benchmarks_by_name_mappings: Vec<TokenStream2> = Vec::new();
586	let test_idents: Vec<Ident> = benchmark_names_str
587		.iter()
588		.map(|n| Ident::new(format!("test_benchmark_{}", n).as_str(), Span::call_site()))
589		.collect();
590	for i in 0..benchmark_names.len() {
591		let name_ident = &benchmark_names[i];
592		let name_str = &benchmark_names_str[i];
593		let test_ident = &test_idents[i];
594		selected_benchmark_mappings.push(quote!(#name_str => SelectedBenchmark::#name_ident));
595		benchmarks_by_name_mappings.push(quote!(#name_str => Self::#test_ident()))
596	}
597
598	let impl_test_function = content
599		.iter_mut()
600		.find_map(|item| {
601			let Item::Macro(item_macro) = item else {
602				return None;
603			};
604
605			if !item_macro
606				.mac
607				.path
608				.segments
609				.iter()
610				.any(|s| s.ident == "impl_benchmark_test_suite")
611			{
612				return None;
613			}
614
615			let tokens = item_macro.mac.tokens.clone();
616			*item = Item::Verbatim(quote! {});
617
618			Some(quote! {
619				impl_test_function!(
620					(#( {} #benchmark_names )*)
621					(#( #extra_benchmark_names )*)
622					(#( #skip_meta_benchmark_names )*)
623					#tokens
624				);
625			})
626		})
627		.unwrap_or(quote! {});
628
629	// emit final quoted tokens
630	let res = quote! {
631		#(#mod_attrs)
632		*
633		#mod_vis mod #mod_name {
634			#(#content)
635			*
636
637			#[allow(non_camel_case_types)]
638			enum SelectedBenchmark {
639				#(#benchmark_names),
640				*
641			}
642
643			impl<#type_impl_generics> #krate::BenchmarkingSetup<#type_use_generics> for SelectedBenchmark where #where_clause {
644				fn components(&self) -> #krate::__private::Vec<(#krate::BenchmarkParameter, u32, u32)> {
645					match self {
646						#(
647							Self::#benchmark_names => {
648								<#benchmark_names as #krate::BenchmarkingSetup<#type_use_generics>>::components(&#benchmark_names)
649							}
650						)
651						*
652					}
653				}
654
655				fn instance(
656					&self,
657					recording: &mut impl #krate::Recording,
658					components: &[(#krate::BenchmarkParameter, u32)],
659					verify: bool,
660				) -> Result<(), #krate::BenchmarkError> {
661					match self {
662						#(
663							Self::#benchmark_names => {
664								<#benchmark_names as #krate::BenchmarkingSetup<
665									#type_use_generics
666								>>::instance(&#benchmark_names, recording, components, verify)
667							}
668						)
669						*
670					}
671				}
672			}
673			#[cfg(any(feature = "runtime-benchmarks", test))]
674			impl<#type_impl_generics> #krate::Benchmarking for Pallet<#type_use_generics>
675			where T: #frame_system::Config, #where_clause
676			{
677				fn benchmarks(
678					extra: bool,
679				) -> #krate::__private::Vec<#krate::BenchmarkMetadata> {
680					let mut all_names = #krate::__private::vec![
681						#(#benchmark_names_str),
682						*
683					];
684					if !extra {
685						let extra = [
686							#(#extra_benchmark_names_str),
687							*
688						];
689						all_names.retain(|x| !extra.contains(x));
690					}
691					let pov_modes:
692						#krate::__private::Vec<(
693							#krate::__private::Vec<u8>,
694							#krate::__private::Vec<(
695								#krate::__private::Vec<u8>,
696								#krate::__private::Vec<u8>
697							)>,
698						)> = #krate::__private::vec![
699						#( #pov_modes ),*
700					];
701					all_names.into_iter().map(|benchmark| {
702						let selected_benchmark = match benchmark {
703							#(#selected_benchmark_mappings),
704							*,
705							_ => panic!("all benchmarks should be selectable")
706						};
707						let components = <SelectedBenchmark as #krate::BenchmarkingSetup<#type_use_generics>>::components(&selected_benchmark);
708						let name = benchmark.as_bytes().to_vec();
709						let modes = pov_modes.iter().find(|p| p.0 == name).map(|p| p.1.clone());
710
711						#krate::BenchmarkMetadata {
712							name: benchmark.as_bytes().to_vec(),
713							components,
714							pov_modes: modes.unwrap_or_default(),
715						}
716					}).collect::<#krate::__private::Vec<_>>()
717				}
718
719				fn run_benchmark(
720					extrinsic: &[u8],
721					c: &[(#krate::BenchmarkParameter, u32)],
722					whitelist: &[#krate::__private::TrackedStorageKey],
723					verify: bool,
724					internal_repeats: u32,
725				) -> Result<#krate::__private::Vec<#krate::BenchmarkResult>, #krate::BenchmarkError> {
726					let extrinsic = #krate::__private::str::from_utf8(extrinsic).map_err(|_| "`extrinsic` is not a valid utf-8 string!")?;
727					let selected_benchmark = match extrinsic {
728						#(#selected_benchmark_mappings),
729						*,
730						_ => return Err("Could not find extrinsic.".into()),
731					};
732					let mut whitelist = whitelist.to_vec();
733					let whitelisted_caller_key = <#frame_system::Account<
734						T,
735					> as #krate::__private::storage::StorageMap<_, _,>>::hashed_key_for(
736						#krate::whitelisted_caller::<T::AccountId>()
737					);
738					whitelist.push(whitelisted_caller_key.into());
739					let transactional_layer_key = #krate::__private::TrackedStorageKey::new(
740						#krate::__private::storage::transactional::TRANSACTION_LEVEL_KEY.into(),
741					);
742					whitelist.push(transactional_layer_key);
743					// Whitelist the `:extrinsic_index`.
744					let extrinsic_index = #krate::__private::TrackedStorageKey::new(
745						#krate::__private::well_known_keys::EXTRINSIC_INDEX.into()
746					);
747					whitelist.push(extrinsic_index);
748					// Whitelist the `:intrablock_entropy`.
749					let intrablock_entropy = #krate::__private::TrackedStorageKey::new(
750						#krate::__private::well_known_keys::INTRABLOCK_ENTROPY.into()
751					);
752					whitelist.push(intrablock_entropy);
753
754					#krate::benchmarking::set_whitelist(whitelist.clone());
755					let mut results: #krate::__private::Vec<#krate::BenchmarkResult> = #krate::__private::Vec::new();
756
757					let on_before_start = || {
758						// Set the block number to at least 1 so events are deposited.
759						if #krate::__private::Zero::is_zero(&#frame_system::Pallet::<T>::block_number()) {
760							#frame_system::Pallet::<T>::set_block_number(1u32.into());
761						}
762
763						// Commit the externalities to the database, flushing the DB cache.
764						// This will enable worst case scenario for reading from the database.
765						#krate::benchmarking::commit_db();
766
767						// Access all whitelisted keys to get them into the proof recorder since the
768						// recorder does now have a whitelist.
769						for key in &whitelist {
770							#krate::__private::storage::unhashed::get_raw(&key.key);
771						}
772
773						// Reset the read/write counter so we don't count operations in the setup process.
774						#krate::benchmarking::reset_read_write_count();
775					};
776
777					// Always do at least one internal repeat...
778					for _ in 0 .. internal_repeats.max(1) {
779						// Always reset the state after the benchmark.
780						#krate::__private::defer!(#krate::benchmarking::wipe_db());
781
782						// Time the extrinsic logic.
783						#krate::__private::log::trace!(
784							target: "benchmark",
785							"Start Benchmark: {} ({:?})",
786							extrinsic,
787							c
788						);
789
790						let mut recording = #krate::BenchmarkRecording::new(&on_before_start);
791						<SelectedBenchmark as #krate::BenchmarkingSetup<#type_use_generics>>::instance(&selected_benchmark, &mut recording, c, verify)?;
792
793						// Calculate the diff caused by the benchmark.
794						let elapsed_extrinsic = recording.elapsed_extrinsic().expect("elapsed time should be recorded");
795						let diff_pov = recording.diff_pov().unwrap_or_default();
796
797						// Commit the changes to get proper write count
798						#krate::benchmarking::commit_db();
799						#krate::__private::log::trace!(
800							target: "benchmark",
801							"End Benchmark: {} ns", elapsed_extrinsic
802						);
803						let read_write_count = #krate::benchmarking::read_write_count();
804						#krate::__private::log::trace!(
805							target: "benchmark",
806							"Read/Write Count {:?}", read_write_count
807						);
808
809						// Time the storage root recalculation.
810						let start_storage_root = #krate::benchmarking::current_time();
811						#krate::__private::storage_root(#krate::__private::StateVersion::V1);
812						let finish_storage_root = #krate::benchmarking::current_time();
813						let elapsed_storage_root = finish_storage_root - start_storage_root;
814
815						let skip_meta = [ #(#skip_meta_benchmark_names_str),* ];
816						let read_and_written_keys = if skip_meta.contains(&extrinsic) {
817							#krate::__private::vec![(b"Skipped Metadata".to_vec(), 0, 0, false)]
818						} else {
819							#krate::benchmarking::get_read_and_written_keys()
820						};
821
822						results.push(#krate::BenchmarkResult {
823							components: c.to_vec(),
824							extrinsic_time: elapsed_extrinsic,
825							storage_root_time: elapsed_storage_root,
826							reads: read_write_count.0,
827							repeat_reads: read_write_count.1,
828							writes: read_write_count.2,
829							repeat_writes: read_write_count.3,
830							proof_size: diff_pov,
831							keys: read_and_written_keys,
832						});
833					}
834
835					return Ok(results);
836				}
837			}
838
839			#[cfg(test)]
840			impl<#type_impl_generics> Pallet<#type_use_generics> where T: #frame_system::Config, #where_clause {
841				/// Test a particular benchmark by name.
842				///
843				/// This isn't called `test_benchmark_by_name` just in case some end-user eventually
844				/// writes a benchmark, itself called `by_name`; the function would be shadowed in
845				/// that case.
846				///
847				/// This is generally intended to be used by child test modules such as those created
848				/// by the `impl_benchmark_test_suite` macro. However, it is not an error if a pallet
849				/// author chooses not to implement benchmarks.
850				#[allow(unused)]
851				fn test_bench_by_name(name: &[u8]) -> Result<(), #krate::BenchmarkError> {
852					let name = #krate::__private::str::from_utf8(name)
853						.map_err(|_| -> #krate::BenchmarkError { "`name` is not a valid utf8 string!".into() })?;
854					match name {
855						#(#benchmarks_by_name_mappings),
856						*,
857						_ => Err("Could not find test for requested benchmark.".into()),
858					}
859				}
860			}
861
862			#impl_test_function
863		}
864		#mod_vis use #mod_name::*;
865	};
866	Ok(res.into())
867}
868
869/// Prepares a [`Vec<ParamDef>`] to be interpolated by [`quote!`] by creating easily-iterable
870/// arrays formatted in such a way that they can be interpolated directly.
871struct UnrolledParams {
872	param_ranges: Vec<TokenStream2>,
873	param_names: Vec<TokenStream2>,
874}
875
876impl UnrolledParams {
877	/// Constructs an [`UnrolledParams`] from a [`Vec<ParamDef>`]
878	fn from(params: &Vec<ParamDef>) -> UnrolledParams {
879		let param_ranges: Vec<TokenStream2> = params
880			.iter()
881			.map(|p| {
882				let name = Ident::new(&p.name, Span::call_site());
883				let start = &p.start;
884				let end = &p.end;
885				quote!(#name, #start, #end)
886			})
887			.collect();
888		let param_names: Vec<TokenStream2> = params
889			.iter()
890			.map(|p| {
891				let name = Ident::new(&p.name, Span::call_site());
892				quote!(#name)
893			})
894			.collect();
895		UnrolledParams { param_ranges, param_names }
896	}
897}
898
899/// Performs expansion of an already-parsed [`BenchmarkDef`].
900fn expand_benchmark(
901	benchmark_def: BenchmarkDef,
902	name: &Ident,
903	is_instance: bool,
904	where_clause: TokenStream2,
905) -> TokenStream2 {
906	// set up variables needed during quoting
907	let krate = match generate_access_from_frame_or_crate("frame-benchmarking") {
908		Ok(ident) => ident,
909		Err(err) => return err.to_compile_error().into(),
910	};
911	let frame_system = match generate_access_from_frame_or_crate("frame-system") {
912		Ok(path) => path,
913		Err(err) => return err.to_compile_error().into(),
914	};
915	let codec = quote!(#krate::__private::codec);
916	let traits = quote!(#krate::__private::traits);
917	let setup_stmts = benchmark_def.setup_stmts;
918	let verify_stmts = benchmark_def.verify_stmts;
919	let last_stmt = benchmark_def.last_stmt;
920	let test_ident =
921		Ident::new(format!("test_benchmark_{}", name.to_string()).as_str(), Span::call_site());
922
923	// unroll params (prepare for quoting)
924	let unrolled = UnrolledParams::from(&benchmark_def.params);
925	let param_names = unrolled.param_names;
926	let param_ranges = unrolled.param_ranges;
927
928	let type_use_generics = match is_instance {
929		false => quote!(T),
930		true => quote!(T, I),
931	};
932
933	let type_impl_generics = match is_instance {
934		false => quote!(T: Config),
935		true => quote!(T: Config<I>, I: 'static),
936	};
937
938	// used in the benchmarking impls
939	let (pre_call, post_call, fn_call_body) = match &benchmark_def.call_def {
940		BenchmarkCallDef::ExtrinsicCall { origin, expr_call, attr_span: _ } => {
941			let mut expr_call = expr_call.clone();
942
943			// remove first arg from expr_call
944			let mut final_args = Punctuated::<Expr, Comma>::new();
945			let args: Vec<&Expr> = expr_call.args.iter().collect();
946			for arg in &args[1..] {
947				final_args.push((*(*arg)).clone());
948			}
949			expr_call.args = final_args;
950
951			let origin = match origin {
952				Expr::Cast(t) => {
953					let ty = t.ty.clone();
954					quote! {
955						<<T as #frame_system::Config>::RuntimeOrigin as From<#ty>>::from(#origin);
956					}
957				},
958				_ => quote! {
959					#origin.into();
960				},
961			};
962
963			// determine call name (handles `_` and normal call syntax)
964			let expr_span = expr_call.span();
965			let call_err = || {
966				syn::Error::new(expr_span, "Extrinsic call must be a function call or `_`")
967					.to_compile_error()
968			};
969			let call_name = match *expr_call.func {
970				Expr::Path(expr_path) => {
971					// normal function call
972					let Some(segment) = expr_path.path.segments.last() else { return call_err() };
973					segment.ident.to_string()
974				},
975				Expr::Infer(_) => {
976					// `_` style
977					// replace `_` with fn name
978					name.to_string()
979				},
980				_ => return call_err(),
981			};
982
983			// modify extrinsic call to be prefixed with "new_call_variant"
984			let call_name = format!("new_call_variant_{}", call_name);
985			let mut punct: Punctuated<PathSegment, PathSep> = Punctuated::new();
986			punct.push(PathSegment {
987				arguments: PathArguments::None,
988				ident: Ident::new(call_name.as_str(), Span::call_site()),
989			});
990			*expr_call.func = Expr::Path(ExprPath {
991				attrs: vec![],
992				qself: None,
993				path: Path { leading_colon: None, segments: punct },
994			});
995			let pre_call = quote! {
996				let __call = Call::<#type_use_generics>::#expr_call;
997				let __benchmarked_call_encoded = #codec::Encode::encode(&__call);
998			};
999			let post_call = quote! {
1000				let __call_decoded = <Call<#type_use_generics> as #codec::Decode>
1001					::decode(&mut &__benchmarked_call_encoded[..])
1002					.expect("call is encoded above, encoding must be correct");
1003				let __origin = #origin;
1004				<Call<#type_use_generics> as #traits::UnfilteredDispatchable>::dispatch_bypass_filter(
1005					__call_decoded,
1006					__origin,
1007				)
1008			};
1009			(
1010				// (pre_call, post_call, fn_call_body):
1011				pre_call.clone(),
1012				quote!(#post_call?;),
1013				quote! {
1014					#pre_call
1015					#post_call.unwrap();
1016				},
1017			)
1018		},
1019		BenchmarkCallDef::Block { block, attr_span: _ } =>
1020			(quote!(), quote!(#block), quote!(#block)),
1021	};
1022
1023	let vis = benchmark_def.fn_vis;
1024
1025	// remove #[benchmark] attribute
1026	let fn_attrs = benchmark_def
1027		.fn_attrs
1028		.iter()
1029		.filter(|attr| !attr.path().is_ident(keywords::BENCHMARK_TOKEN));
1030
1031	// modify signature generics, ident, and inputs, e.g:
1032	// before: `fn bench(u: Linear<1, 100>) -> Result<(), BenchmarkError>`
1033	// after: `fn _bench <T: Config<I>, I: 'static>(u: u32, verify: bool) -> Result<(),
1034	// BenchmarkError>`
1035	let mut sig = benchmark_def.fn_sig;
1036	sig.generics = parse_quote!(<#type_impl_generics>);
1037	if !where_clause.is_empty() {
1038		sig.generics.where_clause = parse_quote!(where #where_clause);
1039	}
1040	sig.ident =
1041		Ident::new(format!("_{}", name.to_token_stream().to_string()).as_str(), Span::call_site());
1042	let mut fn_param_inputs: Vec<TokenStream2> =
1043		param_names.iter().map(|name| quote!(#name: u32)).collect();
1044	fn_param_inputs.push(quote!(verify: bool));
1045	sig.inputs = parse_quote!(#(#fn_param_inputs),*);
1046
1047	// used in instance() impl
1048	let impl_last_stmt = match &last_stmt {
1049		Some(stmt) => quote!(#stmt),
1050		None => quote!(Ok(())),
1051	};
1052	let fn_attrs_clone = fn_attrs.clone();
1053
1054	let fn_def = quote! {
1055		#(
1056			#fn_attrs_clone
1057		)*
1058		#vis #sig {
1059			#(
1060				#setup_stmts
1061			)*
1062			#fn_call_body
1063			if verify {
1064				#(
1065					#verify_stmts
1066				)*
1067			}
1068			#last_stmt
1069		}
1070	};
1071
1072	// generate final quoted tokens
1073	let res = quote! {
1074		// benchmark function definition
1075		#fn_def
1076
1077		#[allow(non_camel_case_types)]
1078		#(
1079			#fn_attrs
1080		)*
1081		struct #name;
1082
1083		#[allow(unused_variables)]
1084		impl<#type_impl_generics> #krate::BenchmarkingSetup<#type_use_generics>
1085		for #name where #where_clause {
1086			fn components(&self) -> #krate::__private::Vec<(#krate::BenchmarkParameter, u32, u32)> {
1087				#krate::__private::vec! [
1088					#(
1089						(#krate::BenchmarkParameter::#param_ranges)
1090					),*
1091				]
1092			}
1093
1094			fn instance(
1095				&self,
1096				recording: &mut impl #krate::Recording,
1097				components: &[(#krate::BenchmarkParameter, u32)],
1098				verify: bool
1099			) -> Result<(), #krate::BenchmarkError> {
1100				#(
1101					// prepare instance #param_names
1102					let #param_names = components.iter()
1103						.find(|&c| c.0 == #krate::BenchmarkParameter::#param_names)
1104						.ok_or("Could not find component during benchmark preparation.")?
1105						.1;
1106				)*
1107
1108				// benchmark setup code
1109				#(
1110					#setup_stmts
1111				)*
1112				#pre_call
1113				recording.start();
1114				#post_call
1115				recording.stop();
1116				if verify {
1117					#(
1118						#verify_stmts
1119					)*
1120				}
1121				#impl_last_stmt
1122			}
1123		}
1124
1125		#[cfg(test)]
1126		impl<#type_impl_generics> Pallet<#type_use_generics> where T: #frame_system::Config, #where_clause {
1127			#[allow(unused)]
1128			fn #test_ident() -> Result<(), #krate::BenchmarkError> {
1129				let selected_benchmark = SelectedBenchmark::#name;
1130				let components = <
1131					SelectedBenchmark as #krate::BenchmarkingSetup<T, _>
1132				>::components(&selected_benchmark);
1133				let execute_benchmark = |
1134					c: #krate::__private::Vec<(#krate::BenchmarkParameter, u32)>
1135				| -> Result<(), #krate::BenchmarkError> {
1136					// Always reset the state after the benchmark.
1137					#krate::__private::defer!(#krate::benchmarking::wipe_db());
1138
1139					let on_before_start = || {
1140						// Set the block number to at least 1 so events are deposited.
1141						if #krate::__private::Zero::is_zero(&#frame_system::Pallet::<T>::block_number()) {
1142							#frame_system::Pallet::<T>::set_block_number(1u32.into());
1143						}
1144					};
1145
1146					// Run execution + verification
1147					<SelectedBenchmark as #krate::BenchmarkingSetup<T, _>>::test_instance(&selected_benchmark,  &c, &on_before_start)
1148				};
1149
1150				if components.is_empty() {
1151					execute_benchmark(Default::default())?;
1152				} else {
1153					let num_values: u32 = if let Ok(ev) = std::env::var("VALUES_PER_COMPONENT") {
1154						ev.parse().map_err(|_| {
1155							#krate::BenchmarkError::Stop(
1156								"Could not parse env var `VALUES_PER_COMPONENT` as u32."
1157							)
1158						})?
1159					} else {
1160						6
1161					};
1162
1163					if num_values < 2 {
1164						return Err("`VALUES_PER_COMPONENT` must be at least 2".into());
1165					}
1166
1167					for (name, low, high) in components.clone().into_iter() {
1168						// Test the lowest, highest (if its different from the lowest)
1169						// and up to num_values-2 more equidistant values in between.
1170						// For 0..10 and num_values=6 this would mean: [0, 2, 4, 6, 8, 10]
1171						if high < low {
1172							return Err("The start of a `ParamRange` must be less than or equal to the end".into());
1173						}
1174
1175						let mut values = #krate::__private::vec![low];
1176						let diff = (high - low).min(num_values - 1);
1177						let slope = (high - low) as f32 / diff as f32;
1178
1179						for i in 1..=diff {
1180							let value = ((low as f32 + slope * i as f32) as u32)
1181											.clamp(low, high);
1182							values.push(value);
1183						}
1184
1185						for component_value in values {
1186							// Select the max value for all the other components.
1187							let c: #krate::__private::Vec<(#krate::BenchmarkParameter, u32)> = components
1188								.iter()
1189								.map(|(n, _, h)|
1190									if *n == name {
1191										(*n, component_value)
1192									} else {
1193										(*n, *h)
1194									}
1195								)
1196								.collect();
1197
1198							execute_benchmark(c)?;
1199						}
1200					}
1201				}
1202				return Ok(());
1203			}
1204		}
1205	};
1206	res
1207}