tenant.openagents/omega

1use proc_macro::TokenStream;
2use quote::{format_ident, quote};
3use syn::{Expr, ItemFn, LitStr, parse::Parser, spanned::Spanned};
4
5pub fn bench(args: TokenStream, function: TokenStream) -> TokenStream {
6    let mut fps: Option<u64> = None;
7    let mut inputs: Option<Expr> = None;
8    let mut input_name: Option<LitStr> = None;
9    let mut group_name: Option<LitStr> = None;
10    let mut sample_size: Option<usize> = None;
11    if !args.is_empty() {
12        let parser = syn::meta::parser(|meta| {
13            if meta.path.is_ident("fps") {
14                let value: syn::LitInt = meta.value()?.parse()?;
15                let value = value.base10_parse::<u64>()?;
16                if value == 0 {
17                    return Err(meta.error("#[gpui::bench] `fps` must be greater than zero"));
18                }
19                fps = Some(value);
20                Ok(())
21            } else if meta.path.is_ident("inputs") {
22                inputs = Some(meta.value()?.parse()?);
23                Ok(())
24            } else if meta.path.is_ident("input_name") {
25                input_name = Some(meta.value()?.parse()?);
26                Ok(())
27            } else if meta.path.is_ident("group") {
28                group_name = Some(meta.value()?.parse()?);
29                Ok(())
30            } else if meta.path.is_ident("sample_size") {
31                let value: syn::LitInt = meta.value()?.parse()?;
32                let value = value.base10_parse::<usize>()?;
33                if value == 0 {
34                    return Err(
35                        meta.error("#[gpui::bench] `sample_size` must be greater than zero")
36                    );
37                }
38                sample_size = Some(value);
39                Ok(())
40            } else {
41                Err(meta.error(
42                    "#[gpui::bench] only accepts `fps = N`, `inputs = EXPR`, `input_name = \"...\"`, `group = \"...\"`, and `sample_size = N`",
43                ))
44            }
45        });
46        if let Err(error) = parser.parse(args) {
47            return error_to_stream(error);
48        }
49    }
50
51    // The frame budget math lives in `BenchReport` so `bench_context` is the
52    // single source of truth; `default()` supplies the default frame rate.
53    let report_expr = match fps {
54        Some(fps) => quote! { gpui::BenchReport::with_fps(#fps) },
55        None => quote! { gpui::BenchReport::default() },
56    };
57
58    let mut inner_fn = match syn::parse::<ItemFn>(function) {
59        Ok(function) => function,
60        Err(error) => return error_to_stream(error),
61    };
62
63    if let Some(asyncness) = &inner_fn.sig.asyncness {
64        return error_to_stream(syn::Error::new(
65            asyncness.span(),
66            "#[gpui::bench] does not support async benchmark functions yet",
67        ));
68    }
69
70    let outer_fn_name = inner_fn.sig.ident.clone();
71    let inner_fn_name = format_ident!("__gpui_bench_{}", outer_fn_name);
72    inner_fn.sig.ident = inner_fn_name.clone();
73
74    let benchmark = if let Some(inputs) = inputs {
75        let input_name = match input_name {
76            Some(input_name) => quote! { #input_name },
77            None => quote! { stringify!(#outer_fn_name) },
78        };
79        let group_name = match group_name {
80            Some(group_name) => quote! { #group_name },
81            None => quote! { stringify!(#outer_fn_name) },
82        };
83        let sample_size =
84            sample_size.map(|sample_size| quote! { group.sample_size(#sample_size); });
85        quote! {
86            let report = #report_expr;
87            let mut group = criterion.benchmark_group(#group_name);
88            #sample_size
89            for input in #inputs {
90                group.bench_with_input(criterion::BenchmarkId::new(#input_name, &input), &input, {
91                    let report = report.clone();
92                    move |bencher, input| {
93                        let mut cx = gpui::BenchAppContext::new_with_platform_and_report(
94                            gpui::bench_platform(
95                                Some(Box::new(|| {
96                                    gpui_platform::current_headless_renderer()
97                                })),
98                                gpui_platform::current_platform(true).text_system(),
99                            ),
100                            Some(stringify!(#outer_fn_name)),
101                            bencher,
102                            report.clone(),
103                        );
104                        #inner_fn_name(input, &mut cx);
105                        cx.teardown();
106                    }
107                });
108            }
109            group.finish();
110            report.print(Some(stringify!(#outer_fn_name)));
111        }
112    } else {
113        if let Some(input_name) = input_name {
114            return error_to_stream(syn::Error::new(
115                input_name.span(),
116                "#[gpui::bench] `input_name` requires `inputs`",
117            ));
118        }
119        if let Some(group_name) = group_name {
120            return error_to_stream(syn::Error::new(
121                group_name.span(),
122                "#[gpui::bench] `group` requires `inputs`",
123            ));
124        }
125        if sample_size.is_some() {
126            return error_to_stream(syn::Error::new(
127                proc_macro2::Span::call_site(),
128                "#[gpui::bench] `sample_size` requires `inputs`",
129            ));
130        }
131        quote! {
132            let report = #report_expr;
133            criterion.bench_function(stringify!(#outer_fn_name), {
134                let report = report.clone();
135                move |bencher| {
136                    let mut cx = gpui::BenchAppContext::new_with_platform_and_report(
137                        gpui::bench_platform(
138                            Some(Box::new(|| {
139                                gpui_platform::current_headless_renderer()
140                            })),
141                            gpui_platform::current_platform(true).text_system(),
142                        ),
143                        Some(stringify!(#outer_fn_name)),
144                        bencher,
145                        report.clone(),
146                    );
147                    #inner_fn_name(&mut cx);
148                    cx.teardown();
149                }
150            });
151            report.print(Some(stringify!(#outer_fn_name)));
152        }
153    };
154
155    TokenStream::from(quote! {
156        #inner_fn
157
158        fn #outer_fn_name(criterion: &mut criterion::Criterion) {
159            #benchmark
160        }
161
162    })
163}
164
165fn error_to_stream(error: syn::Error) -> TokenStream {
166    TokenStream::from(error.into_compile_error())
167}
168