1use proc_macro::TokenStream;
2use quote::{quote, ToTokens};
3use syn::{
4 parse::{Parse, Parser},
5 spanned::Spanned,
6 Expr, ExprArray, Ident, Token, Type,
7};
8
9use crate::{tokens, Macro};
10
11/// Values from parsed options shared between `#[divan::bench]` and
12/// `#[divan::bench_group]`.
13///
14/// The `crate` option is not included because it is only needed to get proper
15/// access to `__private`.
16pub(crate) struct AttrOptions {
17 /// `divan::__private`.
18 pub private_mod: proc_macro2::TokenStream,
19
20 /// Custom name for the benchmark or group.
21 pub name_expr: Option<Expr>,
22
23 /// `IntoIterator` from which to provide runtime arguments.
24 pub args_expr: Option<Expr>,
25
26 /// Options for generic functions.
27 pub generic: GenericOptions,
28
29 /// The `BenchOptions.counters` field and its value, followed by a comma.
30 pub counters: proc_macro2::TokenStream,
31
32 /// Options used directly as `BenchOptions` fields.
33 ///
34 /// Option reuse is handled by the compiler ensuring `BenchOptions` fields
35 /// are not repeated.
36 pub bench_options: Vec<(Ident, Expr)>,
37}
38
39impl AttrOptions {
40 pub fn parse(tokens: TokenStream, target_macro: Macro) -> Result<Self, TokenStream> {
41 let macro_name = target_macro.name();
42
43 let mut divan_crate = None::<syn::Path>;
44 let mut name_expr = None::<Expr>;
45 let mut args_expr = None::<Expr>;
46 let mut bench_options = Vec::new();
47
48 let mut counters = Vec::<(proc_macro2::TokenStream, Option<&str>)>::new();
49 let mut counters_ident = None::<Ident>;
50
51 let mut seen_bytes_count = false;
52 let mut seen_chars_count = false;
53 let mut seen_cycles_count = false;
54 let mut seen_items_count = false;
55
56 let mut generic = GenericOptions::default();
57
58 let attr_parser = syn::meta::parser(|meta| {
59 macro_rules! error {
60 ($($t:tt)+) => {
61 return Err(meta.error(format_args!($($t)+)))
62 };
63 }
64
65 let Some(ident) = meta.path.get_ident() else {
66 error!("unsupported '{macro_name}' option");
67 };
68
69 let ident_name = ident.to_string();
70 let ident_name = ident_name.strip_prefix("r#").unwrap_or(&ident_name);
71
72 let repeat_error = || error!("repeated '{macro_name}' option '{ident_name}'");
73 let unsupported_error = || error!("unsupported '{macro_name}' option '{ident_name}'");
74
75 macro_rules! parse {
76 ($storage:expr) => {
77 if $storage.is_none() {
78 $storage = Some(meta.value()?.parse()?);
79 } else {
80 return repeat_error();
81 }
82 };
83 }
84
85 match ident_name {
86 "crate" => parse!(divan_crate),
87 "name" => parse!(name_expr),
88 "types" => {
89 match target_macro {
90 Macro::Bench { fn_sig } => {
91 if fn_sig.generics.type_params().next().is_none() {
92 error!("generic type required for '{macro_name}' option '{ident_name}'");
93 }
94 }
95 _ => return unsupported_error(),
96 }
97
98 parse!(generic.types);
99 }
100 "consts" => {
101 match target_macro {
102 Macro::Bench { fn_sig } => {
103 if fn_sig.generics.const_params().next().is_none() {
104 error!("generic const required for '{macro_name}' option '{ident_name}'");
105 }
106 }
107 _ => return unsupported_error(),
108 }
109
110 parse!(generic.consts);
111 }
112 "args" => {
113 match target_macro {
114 Macro::Bench { fn_sig } => {
115 if !matches!(fn_sig.inputs.len(), 1 | 2) {
116 return Err(meta.error(format_args!("function argument required for '{macro_name}' option '{ident_name}'")));
117 }
118 }
119 _ => return unsupported_error(),
120 }
121
122 parse!(args_expr);
123 }
124 "counter" => {
125 if counters_ident.is_some() {
126 return repeat_error();
127 }
128 let value: Expr = meta.value()?.parse()?;
129 counters.push((value.into_token_stream(), None));
130 counters_ident = Some(Ident::new("counters", ident.span()));
131 }
132 "counters" => {
133 if counters_ident.is_some() {
134 return repeat_error();
135 }
136 let values: ExprArray = meta.value()?.parse()?;
137 counters.extend(
138 values.elems.into_iter().map(|elem| (elem.into_token_stream(), None)),
139 );
140 counters_ident = Some(ident.clone());
141 }
142
143 "bytes_count" if seen_bytes_count => return repeat_error(),
144 "chars_count" if seen_chars_count => return repeat_error(),
145 "cycles_count" if seen_cycles_count => return repeat_error(),
146 "items_count" if seen_items_count => return repeat_error(),
147
148 "bytes_count" | "chars_count" | "cycles_count" | "items_count" => {
149 let name = match ident_name {
150 "bytes_count" => {
151 seen_bytes_count = true;
152 "BytesCount"
153 }
154 "chars_count" => {
155 seen_chars_count = true;
156 "CharsCount"
157 }
158 "cycles_count" => {
159 seen_cycles_count = true;
160 "CyclesCount"
161 }
162 "items_count" => {
163 seen_items_count = true;
164 "ItemsCount"
165 }
166 _ => unreachable!(),
167 };
168
169 let value: Expr = meta.value()?.parse()?;
170 counters.push((value.into_token_stream(), Some(name)));
171 counters_ident = Some(Ident::new("counters", proc_macro2::Span::call_site()));
172 }
173
174 _ => {
175 let value: Expr = match meta.value() {
176 Ok(value) => value.parse()?,
177
178 // If the option is missing `=`, use a `true` literal.
179 Err(_) => Expr::Lit(syn::ExprLit {
180 lit: syn::LitBool::new(true, meta.path.span()).into(),
181 attrs: Vec::new(),
182 }),
183 };
184
185 bench_options.push((ident.clone(), value));
186 }
187 }
188
189 Ok(())
190 });
191
192 match attr_parser.parse(tokens) {
193 Ok(()) => {}
194 Err(error) => return Err(error.into_compile_error().into()),
195 }
196
197 let divan_crate = divan_crate.unwrap_or_else(|| syn::parse_quote!(::divan));
198 let private_mod = quote! { #divan_crate::__private };
199
200 let counters = counters.iter().map(|(expr, type_name)| match type_name {
201 Some(type_name) => {
202 let type_name = Ident::new(type_name, proc_macro2::Span::call_site());
203 quote! {
204 // We do a scoped import for the expression to override any
205 // local `From` trait.
206 {
207 use ::std::convert::From as _;
208
209 #divan_crate::counter::#type_name::from(#expr)
210 }
211 }
212 }
213 None => expr.to_token_stream(),
214 });
215
216 let counters = counters_ident
217 .map(|ident| {
218 quote! {
219 #ident: #private_mod::new_counter_set() #(.with(#counters))* ,
220 }
221 })
222 .unwrap_or_default();
223
224 Ok(Self { private_mod, name_expr, args_expr, generic, counters, bench_options })
225 }
226
227 /// Produces a function expression for creating `LazyLock<BenchOptions>`.
228 ///
229 /// If the `#[ignore]` attribute is specified, this be provided its
230 /// identifier to set `BenchOptions` using its span. Doing this instead of
231 /// creating the `ignore` identifier ourselves improves compiler error
232 /// diagnostics.
233 pub fn bench_options_fn(
234 &self,
235 ignore_attr_ident: Option<&syn::Path>,
236 ) -> proc_macro2::TokenStream {
237 fn is_lit_array(expr: &Expr) -> bool {
238 let Expr::Array(expr) = expr else {
239 return false;
240 };
241 expr.elems.iter().all(|elem| matches!(elem, Expr::Lit { .. }))
242 }
243
244 let private_mod = &self.private_mod;
245 let option_some = tokens::option_some();
246
247 // Directly set fields on `BenchOptions`. This simplifies things by:
248 // - Having a single source of truth
249 // - Making unknown options a compile error
250 //
251 // We use `..` (struct update syntax) to ensure that no option is set
252 // twice, even if raw identifiers are used. This also has the accidental
253 // benefit of Rust Analyzer recognizing fields and emitting suggestions
254 // with docs and type info.
255 if self.bench_options.is_empty() && self.counters.is_empty() && ignore_attr_ident.is_none()
256 {
257 tokens::option_none()
258 } else {
259 let options_iter = self.bench_options.iter().map(|(option, value)| {
260 let option_name = option.to_string();
261 let option_name = option_name.strip_prefix("r#").unwrap_or(&option_name);
262
263 let wrapped_value: proc_macro2::TokenStream;
264 let value: &dyn ToTokens = match option_name {
265 "threads" => {
266 wrapped_value = if is_lit_array(value) {
267 // If array of literals, just use `&[...]`.
268 quote! { ::std::borrow::Cow::Borrowed(&#value) }
269 } else {
270 quote! { #private_mod::IntoThreads::into_threads(#value) }
271 };
272
273 &wrapped_value
274 }
275
276 // If the option is a `Duration`, use `IntoDuration` to be
277 // polymorphic over `Duration` or `u64`/`f64` seconds.
278 "min_time" | "max_time" => {
279 wrapped_value =
280 quote! { #private_mod::IntoDuration::into_duration(#value) };
281 &wrapped_value
282 }
283
284 _ => value,
285 };
286
287 quote! { #option: #option_some(#value), }
288 });
289
290 let ignore = match ignore_attr_ident {
291 Some(ignore_attr_ident) => quote! { #ignore_attr_ident: #option_some(true), },
292 None => Default::default(),
293 };
294
295 let counters = &self.counters;
296
297 quote! {
298 #option_some(::std::sync::LazyLock::new(|| {
299 #[allow(clippy::needless_update)]
300 #private_mod::BenchOptions {
301 #(#options_iter)*
302
303 // Ignore comes after options so that options take
304 // priority in compiler error diagnostics.
305 #ignore
306
307 #counters
308
309 ..::std::default::Default::default()
310 }
311 }))
312 }
313 }
314 }
315}
316
317/// Options for generic functions.
318#[derive(Default)]
319pub struct GenericOptions {
320 /// Generic types over which to instantiate benchmark functions.
321 pub types: Option<GenericTypes>,
322
323 /// `const` array/slice over which to instantiate benchmark functions.
324 pub consts: Option<Expr>,
325}
326
327impl GenericOptions {
328 /// Returns `true` if set exclusively to either:
329 /// - `types = []`
330 /// - `consts = []`
331 pub fn is_empty(&self) -> bool {
332 match (&self.types, &self.consts) {
333 (Some(types: &GenericTypes), None) => types.is_empty(),
334 (None, Some(Expr::Array(consts: &ExprArray))) => consts.elems.is_empty(),
335 _ => false,
336 }
337 }
338
339 /// Returns an iterator of multiple `Some` for types, or a single `None` if
340 /// there are no types.
341 pub fn types_iter(&self) -> Box<dyn Iterator<Item = Option<&dyn ToTokens>> + '_> {
342 match &self.types {
343 None => Box::new(std::iter::once(None)),
344 Some(GenericTypes::List(types: &Vec)) => {
345 Box::new(types.iter().map(|t: &TokenStream| Some(t as &dyn ToTokens)))
346 }
347 }
348 }
349}
350
351/// Generic types over which to instantiate benchmark functions.
352pub enum GenericTypes {
353 /// List of types, e.g. `[i32, String, ()]`.
354 List(Vec<proc_macro2::TokenStream>),
355}
356
357impl Parse for GenericTypes {
358 fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
359 let content: ParseBuffer<'_>;
360 syn::bracketed!(content in input);
361
362 Ok(Self::List(
363 contentimpl Iterator
364 .parse_terminated(parser:Type::parse, separator:Token![,])?
365 .into_iter()
366 .map(|ty: Type| ty.into_token_stream())
367 .collect(),
368 ))
369 }
370}
371
372impl GenericTypes {
373 pub fn is_empty(&self) -> bool {
374 match self {
375 Self::List(list: &Vec) => list.is_empty(),
376 }
377 }
378}
379