1use proc_macro2::TokenStream;
2use quote::{format_ident, quote, quote_spanned};
3use syn::{
4 FnArg, Ident, ItemFn, Type, parse2, punctuated::Punctuated, spanned::Spanned, token::Comma,
5};
6
7pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
8 let item_span = item.span();
9 let Ok(func) = parse2::<ItemFn>(item) else {
10 return quote_spanned! { item_span =>
11 compile_error!("#[gpui::property_test] must be placed on a function");
12 };
13 };
14
15 let test_name = func.sig.ident.clone();
16 let inner_fn_name = format_ident!("__{test_name}");
17 let outer_fn_attributes = &func.attrs;
18
19 let parsed_args = parse_args(func.sig.inputs, &test_name);
20
21 let inner_body = func.block;
22 let inner_arg_decls = parsed_args.inner_fn_decl_args;
23 let asyncness = func.sig.asyncness;
24
25 let inner_fn = quote! {
26 let #inner_fn_name = #asyncness move |#inner_arg_decls| #inner_body;
27 };
28
29 let arg_errors = parsed_args.errors;
30 let proptest_args = parsed_args.proptest_args;
31 let inner_args = parsed_args.inner_fn_args;
32 let cx_vars = parsed_args.cx_vars;
33 let cx_teardowns = parsed_args.cx_teardowns;
34
35 let proptest_args = quote! {
36 #[strategy = ::gpui::seed_strategy()] __seed: u64,
37 #proptest_args
38 };
39
40 let run_test_body = match &asyncness {
41 None => quote! {
42 #cx_vars
43 #inner_fn_name(#inner_args);
44 #cx_teardowns
45 },
46 Some(_) => quote! {
47 let foreground_executor = gpui::ForegroundExecutor::new(std::sync::Arc::new(dispatcher.clone()));
48 #cx_vars
49 foreground_executor.block_test(#inner_fn_name(#inner_args));
50 #cx_teardowns
51 },
52 };
53
54 quote! {
55 #arg_errors
56
57 #[::gpui::proptest::property_test(proptest_path = "::gpui::proptest", #args)]
58 #(#outer_fn_attributes)*
59 fn #test_name(#proptest_args) {
60 #inner_fn
61
62 ::gpui::run_test_once(
63 __seed,
64 Box::new(move |dispatcher| {
65 #run_test_body
66 }),
67 )
68 }
69 }
70}
71
72#[derive(Default)]
73struct ParsedArgs {
74 cx_vars: TokenStream,
75 cx_teardowns: TokenStream,
76 proptest_args: TokenStream,
77 errors: TokenStream,
78
79 // exprs passed at the call-site
80 inner_fn_args: TokenStream,
81 // args in the declaration
82 inner_fn_decl_args: TokenStream,
83}
84
85fn parse_args(args: Punctuated<FnArg, Comma>, test_name: &Ident) -> ParsedArgs {
86 let mut parsed = ParsedArgs::default();
87 let mut args = args.into_iter().collect();
88
89 remove_cxs(&mut parsed, &mut args, test_name);
90 remove_std_rng(&mut parsed, &mut args);
91 remove_background_executor(&mut parsed, &mut args);
92
93 // all remaining args forwarded to proptest's macro
94 parsed.proptest_args = quote!( #(#args),* );
95
96 parsed
97}
98
99fn remove_cxs(parsed: &mut ParsedArgs, args: &mut Vec<FnArg>, test_name: &Ident) {
100 let mut ix = 0;
101 args.retain_mut(|arg| {
102 if !is_test_cx(arg) {
103 return true;
104 }
105
106 let cx_varname = format_ident!("cx_{ix}");
107 ix += 1;
108
109 parsed.cx_vars.extend(quote!(
110 let mut #cx_varname = gpui::TestAppContext::build(
111 dispatcher.clone(),
112 Some(stringify!(#test_name)),
113 );
114 ));
115 parsed.cx_teardowns.extend(quote!(
116 dispatcher.run_until_parked();
117 #cx_varname.executor().forbid_parking();
118 #cx_varname.quit();
119 dispatcher.run_until_parked();
120 ));
121
122 parsed.inner_fn_decl_args.extend(quote!(#arg,));
123 parsed.inner_fn_args.extend(quote!(&mut #cx_varname,));
124
125 false
126 });
127}
128
129fn remove_std_rng(parsed: &mut ParsedArgs, args: &mut Vec<FnArg>) {
130 args.retain_mut(|arg| {
131 if !is_std_rng(arg) {
132 return true;
133 }
134
135 parsed.errors.extend(quote_spanned! { arg.span() =>
136 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");
137 });
138
139 false
140 });
141}
142
143fn remove_background_executor(parsed: &mut ParsedArgs, args: &mut Vec<FnArg>) {
144 args.retain_mut(|arg| {
145 if !is_background_executor(arg) {
146 return true;
147 }
148
149 parsed.inner_fn_decl_args.extend(quote!(#arg,));
150 parsed
151 .inner_fn_args
152 .extend(quote!(gpui::BackgroundExecutor::new(std::sync::Arc::new(
153 dispatcher.clone()
154 )),));
155
156 false
157 });
158}
159
160// Matches `&TestAppContext` or `&foo::bar::baz::TestAppContext`
161fn is_test_cx(arg: &FnArg) -> bool {
162 let FnArg::Typed(arg) = arg else {
163 return false;
164 };
165
166 let Type::Reference(ty) = &*arg.ty else {
167 return false;
168 };
169
170 let Type::Path(ty) = &*ty.elem else {
171 return false;
172 };
173
174 ty.path
175 .segments
176 .last()
177 .is_some_and(|seg| seg.ident == "TestAppContext")
178}
179
180fn is_std_rng(arg: &FnArg) -> bool {
181 is_path_with_last_segment(arg, "StdRng")
182}
183
184fn is_background_executor(arg: &FnArg) -> bool {
185 is_path_with_last_segment(arg, "BackgroundExecutor")
186}
187
188fn is_path_with_last_segment(arg: &FnArg, last_segment: &str) -> bool {
189 let FnArg::Typed(arg) = arg else {
190 return false;
191 };
192
193 let Type::Path(ty) = &*arg.ty else {
194 return false;
195 };
196
197 ty.path
198 .segments
199 .last()
200 .is_some_and(|seg| seg.ident == last_segment)
201}