tenant.openagents/omega

1use proc_macro2::TokenStream;
2use quote::{ToTokens, format_ident, quote, quote_spanned};
3use syn::{
4    Expr, FnArg, Ident, ItemFn, MetaNameValue, Token, Type,
5    parse::{Parse, ParseStream},
6    parse2,
7    punctuated::Punctuated,
8    spanned::Spanned,
9    token::Comma,
10};
11
12pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
13    let item_span = item.span();
14    let Ok(func) = parse2::<ItemFn>(item) else {
15        return quote_spanned! { item_span =>
16            compile_error!("#[gpui::property_test] must be placed on a function");
17        };
18    };
19
20    let args = match parse2::<Args>(args) {
21        Ok(args) => args,
22        Err(e) => return e.to_compile_error(),
23    };
24
25    let test_name = func.sig.ident.clone();
26    let test_ret_ty = func.sig.output.clone();
27    let inner_fn_name = format_ident!("__{test_name}");
28    let outer_fn_attributes = &func.attrs;
29
30    let parsed_args = parse_args(func.sig.inputs, &test_name);
31
32    let inner_body = func.block;
33    let inner_arg_decls = parsed_args.inner_fn_decl_args;
34    let asyncness = func.sig.asyncness;
35
36    let inner_fn = quote! {
37        let #inner_fn_name = #asyncness move |#inner_arg_decls| #inner_body;
38    };
39
40    let arg_errors = parsed_args.errors;
41    let proptest_args = parsed_args.proptest_args;
42    let inner_args = parsed_args.inner_fn_args;
43    let cx_vars = parsed_args.cx_vars;
44    let cx_teardowns = parsed_args.cx_teardowns;
45
46    let proptest_args = quote! {
47        #[strategy = ::gpui::seed_strategy()] __seed: u64,
48        #proptest_args
49    };
50
51    let run_test_body = match &asyncness {
52        None => quote! {
53            #cx_vars
54            let result = #inner_fn_name(#inner_args);
55            #cx_teardowns
56            result
57        },
58        Some(_) => quote! {
59            let foreground_executor = gpui::ForegroundExecutor::new(std::sync::Arc::new(dispatcher.clone()));
60            #cx_vars
61            let result = foreground_executor.block_test(#inner_fn_name(#inner_args));
62            #cx_teardowns
63            result
64        },
65    };
66
67    let fixed_macro_invocation = args.render();
68
69    quote! {
70        #arg_errors
71
72        #fixed_macro_invocation
73        #(#outer_fn_attributes)*
74        fn #test_name(#proptest_args) #test_ret_ty {
75            #inner_fn
76
77            ::gpui::run_test_once(
78                __seed,
79                Box::new(move |dispatcher| #test_ret_ty {
80                    #run_test_body
81                }),
82            )
83        }
84    }
85}
86
87struct Args {
88    config: Option<Expr>,
89    remaining_args: Vec<MetaNameValue>,
90    errors: TokenStream,
91}
92
93impl Args {
94    /// By default, proptest uses random seeds unless `$PROPTEST_SEED` is set.
95    /// Rather than managing both `$SEED` and `$PROPTEST_SEED`, we intercept
96    /// `config = ...` tokens and add a call to `gpui::apply_seed_to_config`.
97    fn render(&self) -> TokenStream {
98        let user_provided_config = match &self.config {
99            None => quote! { ::gpui::proptest::prelude::ProptestConfig::default() },
100            Some(config) => config.into_token_stream(),
101        };
102
103        let fixed_config = quote!(::gpui::apply_seed_to_proptest_config(#user_provided_config));
104        let remaining_args = &self.remaining_args;
105        let errors = &self.errors;
106
107        quote! {
108            #errors
109            #[::gpui::proptest::property_test(
110                proptest_path = "::gpui::proptest",
111                config = #fixed_config,
112                #(#remaining_args,)*
113            )]
114        }
115    }
116}
117
118impl Parse for Args {
119    fn parse(input: ParseStream) -> syn::Result<Self> {
120        let pairs = Punctuated::<MetaNameValue, Token![,]>::parse_terminated(input)?;
121
122        let mut config = None;
123        let mut remaining_args = vec![];
124        let mut errors = quote!();
125
126        for pair in pairs {
127            match pair.path.get_ident().map(Ident::to_string).as_deref() {
128                Some("config") => config = Some(pair.value),
129                Some("proptest_path") => errors.extend(quote_spanned! {pair.span() =>
130                    compile_error!("`gpui::property_test` overrides the `proptest_path` parameter")
131                }),
132                _ => remaining_args.push(pair),
133            }
134        }
135
136        Ok(Self {
137            config,
138            remaining_args,
139            errors,
140        })
141    }
142}
143
144#[derive(Default)]
145struct ParsedArgs {
146    cx_vars: TokenStream,
147    cx_teardowns: TokenStream,
148    proptest_args: TokenStream,
149    errors: TokenStream,
150
151    // exprs passed at the call-site
152    inner_fn_args: TokenStream,
153    // args in the declaration
154    inner_fn_decl_args: TokenStream,
155}
156
157fn parse_args(args: Punctuated<FnArg, Comma>, test_name: &Ident) -> ParsedArgs {
158    let mut parsed = ParsedArgs::default();
159    let mut args = args.into_iter().collect();
160
161    remove_cxs(&mut parsed, &mut args, test_name);
162    remove_std_rng(&mut parsed, &mut args);
163    remove_background_executor(&mut parsed, &mut args);
164
165    // all remaining args forwarded to proptest's macro
166    parsed.proptest_args = quote!( #(#args),* );
167
168    parsed
169}
170
171fn remove_cxs(parsed: &mut ParsedArgs, args: &mut Vec<FnArg>, test_name: &Ident) {
172    let mut ix = 0;
173    args.retain_mut(|arg| {
174        if !is_test_cx(arg) {
175            return true;
176        }
177
178        let cx_varname = format_ident!("cx_{ix}");
179        ix += 1;
180
181        parsed.cx_vars.extend(quote!(
182            let mut #cx_varname = gpui::TestAppContext::build(
183                dispatcher.clone(),
184                Some(stringify!(#test_name)),
185            );
186        ));
187        parsed.cx_teardowns.extend(quote!(
188            dispatcher.run_until_parked();
189            #cx_varname.executor().forbid_parking();
190            #cx_varname.quit();
191            dispatcher.run_until_parked();
192        ));
193
194        parsed.inner_fn_decl_args.extend(quote!(#arg,));
195        parsed.inner_fn_args.extend(quote!(&mut #cx_varname,));
196
197        false
198    });
199}
200
201fn remove_std_rng(parsed: &mut ParsedArgs, args: &mut Vec<FnArg>) {
202    args.retain_mut(|arg| {
203        if !is_std_rng(arg) {
204            return true;
205        }
206
207        parsed.errors.extend(quote_spanned! { arg.span() =>
208            compile_error!("`StdRng` is not allowed in a property test. Consider implementing `Arbitrary`, or implementing a custom `Strategy`. https://altsysrq.github.io/proptest-book/proptest/tutorial/strategy-basics.html");
209        });
210
211        false
212    });
213}
214
215fn remove_background_executor(parsed: &mut ParsedArgs, args: &mut Vec<FnArg>) {
216    args.retain_mut(|arg| {
217        if !is_background_executor(arg) {
218            return true;
219        }
220
221        parsed.inner_fn_decl_args.extend(quote!(#arg,));
222        parsed
223            .inner_fn_args
224            .extend(quote!(gpui::BackgroundExecutor::new(std::sync::Arc::new(
225                dispatcher.clone()
226            )),));
227
228        false
229    });
230}
231
232// Matches `&TestAppContext` or `&foo::bar::baz::TestAppContext`
233fn is_test_cx(arg: &FnArg) -> bool {
234    let FnArg::Typed(arg) = arg else {
235        return false;
236    };
237
238    let Type::Reference(ty) = &*arg.ty else {
239        return false;
240    };
241
242    let Type::Path(ty) = &*ty.elem else {
243        return false;
244    };
245
246    ty.path
247        .segments
248        .last()
249        .is_some_and(|seg| seg.ident == "TestAppContext")
250}
251
252fn is_std_rng(arg: &FnArg) -> bool {
253    is_path_with_last_segment(arg, "StdRng")
254}
255
256fn is_background_executor(arg: &FnArg) -> bool {
257    is_path_with_last_segment(arg, "BackgroundExecutor")
258}
259
260fn is_path_with_last_segment(arg: &FnArg, last_segment: &str) -> bool {
261    let FnArg::Typed(arg) = arg else {
262        return false;
263    };
264
265    let Type::Path(ty) = &*arg.ty else {
266        return false;
267    };
268
269    ty.path
270        .segments
271        .last()
272        .is_some_and(|seg| seg.ident == last_segment)
273}
274