1use proc_macro::TokenStream;
2use quote::{format_ident, quote};
3use std::mem;
4use syn::{
5 parse_macro_input, parse_quote, spanned::Spanned as _, AttributeArgs, FnArg, ItemFn, Lit, Meta,
6 NestedMeta, Type,
7};
8
9#[proc_macro_attribute]
10pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
11 let mut namespace = format_ident!("gpui");
12
13 let args = syn::parse_macro_input!(args as AttributeArgs);
14 let mut max_retries = 0;
15 let mut num_iterations = 1;
16 let mut starting_seed = 0;
17
18 for arg in args {
19 match arg {
20 NestedMeta::Meta(Meta::Path(name))
21 if name.get_ident().map_or(false, |n| n == "self") =>
22 {
23 namespace = format_ident!("crate");
24 }
25 NestedMeta::Meta(Meta::NameValue(meta)) => {
26 let key_name = meta.path.get_ident().map(|i| i.to_string());
27 let result = (|| {
28 match key_name.as_ref().map(String::as_str) {
29 Some("retries") => max_retries = parse_int(&meta.lit)?,
30 Some("iterations") => num_iterations = parse_int(&meta.lit)?,
31 Some("seed") => starting_seed = parse_int(&meta.lit)?,
32 _ => {
33 return Err(TokenStream::from(
34 syn::Error::new(meta.path.span(), "invalid argument")
35 .into_compile_error(),
36 ))
37 }
38 }
39 Ok(())
40 })();
41
42 if let Err(tokens) = result {
43 return tokens;
44 }
45 }
46 other => {
47 return TokenStream::from(
48 syn::Error::new_spanned(other, "invalid argument").into_compile_error(),
49 )
50 }
51 }
52 }
53
54 let mut inner_fn = parse_macro_input!(function as ItemFn);
55 let inner_fn_attributes = mem::take(&mut inner_fn.attrs);
56 let inner_fn_name = format_ident!("_{}", inner_fn.sig.ident);
57 let outer_fn_name = mem::replace(&mut inner_fn.sig.ident, inner_fn_name.clone());
58
59 let mut outer_fn: ItemFn = if inner_fn.sig.asyncness.is_some() {
60 // Pass to the test function the number of app contexts that it needs,
61 // based on its parameter list.
62 let mut inner_fn_args = proc_macro2::TokenStream::new();
63 for (ix, arg) in inner_fn.sig.inputs.iter().enumerate() {
64 if let FnArg::Typed(arg) = arg {
65 if let Type::Path(ty) = &*arg.ty {
66 let last_segment = ty.path.segments.last();
67 match last_segment.map(|s| s.ident.to_string()).as_deref() {
68 Some("TestAppContext") => {
69 let first_entity_id = ix * 100_000;
70 inner_fn_args.extend(
71 quote!(#namespace::TestAppContext::new(foreground.clone(), background.clone(), #first_entity_id),)
72 );
73 }
74 Some("StdRng") => {
75 inner_fn_args.extend(quote!(rand::SeedableRng::seed_from_u64(seed)));
76 }
77 _ => {
78 return TokenStream::from(
79 syn::Error::new_spanned(arg, "invalid argument")
80 .into_compile_error(),
81 )
82 }
83 }
84 } else {
85 return TokenStream::from(
86 syn::Error::new_spanned(arg, "invalid argument").into_compile_error(),
87 );
88 }
89 } else {
90 return TokenStream::from(
91 syn::Error::new_spanned(arg, "invalid argument").into_compile_error(),
92 );
93 }
94 }
95
96 parse_quote! {
97 #[test]
98 fn #outer_fn_name() {
99 #inner_fn
100
101 let is_randomized = #num_iterations > 1;
102 let mut num_iterations = #num_iterations;
103 let mut starting_seed = #starting_seed;
104 if is_randomized {
105 if let Ok(value) = std::env::var("SEED") {
106 starting_seed = value.parse().expect("invalid SEED variable");
107 }
108 if let Ok(value) = std::env::var("ITERATIONS") {
109 num_iterations = value.parse().expect("invalid ITERATIONS variable");
110 }
111 }
112
113 let mut retries = 0;
114 let mut i = 0;
115 loop {
116 let seed = (starting_seed + i) as u64;
117 if is_randomized {
118 dbg!(seed);
119 }
120
121 let result = std::panic::catch_unwind(|| {
122 let (foreground, background) = #namespace::executor::deterministic(seed);
123 foreground.run(#inner_fn_name(#inner_fn_args));
124 });
125
126 match result {
127 Ok(result) => {
128 retries = 0;
129 i += 1;
130 if i == num_iterations {
131 return result
132 }
133 }
134 Err(error) => {
135 if retries < #max_retries {
136 retries += 1;
137 println!("retrying: attempt {}", retries);
138 } else {
139 if num_iterations > 1 {
140 eprintln!("failing seed: {}", seed);
141 }
142 std::panic::resume_unwind(error);
143 }
144 }
145 }
146 }
147 }
148 }
149 } else {
150 let mut inner_fn_args = proc_macro2::TokenStream::new();
151 for arg in inner_fn.sig.inputs.iter() {
152 if let FnArg::Typed(arg) = arg {
153 if let Type::Path(ty) = &*arg.ty {
154 let last_segment = ty.path.segments.last();
155 if let Some("StdRng") = last_segment.map(|s| s.ident.to_string()).as_deref() {
156 inner_fn_args.extend(quote!(rand::SeedableRng::seed_from_u64(seed),));
157 }
158 } else {
159 inner_fn_args.extend(quote!(cx,));
160 }
161 } else {
162 return TokenStream::from(
163 syn::Error::new_spanned(arg, "invalid argument").into_compile_error(),
164 );
165 }
166 }
167
168 parse_quote! {
169 #[test]
170 fn #outer_fn_name() {
171 #inner_fn
172
173 let is_randomized = #num_iterations > 1;
174 let mut num_iterations = #num_iterations;
175 let mut starting_seed = #starting_seed;
176 if is_randomized {
177 if let Ok(value) = std::env::var("SEED") {
178 starting_seed = value.parse().expect("invalid SEED variable");
179 }
180 if let Ok(value) = std::env::var("ITERATIONS") {
181 num_iterations = value.parse().expect("invalid ITERATIONS variable");
182 }
183 }
184
185 let mut retries = 0;
186 let mut i = 0;
187 loop {
188 let seed = (starting_seed + i) as u64;
189 if is_randomized {
190 dbg!(seed);
191 }
192
193 let result = std::panic::catch_unwind(|| {
194 #namespace::App::test(|cx| {
195 #inner_fn_name(#inner_fn_args);
196 });
197 });
198
199 match result {
200 Ok(result) => {
201 retries = 0;
202 i += 1;
203 if i == num_iterations {
204 return result
205 }
206 }
207 Err(error) => {
208 if retries < #max_retries {
209 retries += 1;
210 println!("retrying: attempt {}", retries);
211 } else {
212 if #num_iterations > 1 {
213 eprintln!("failing seed: {}", seed);
214 }
215 std::panic::resume_unwind(error);
216 }
217 }
218 }
219 }
220 }
221 }
222 };
223 outer_fn.attrs.extend(inner_fn_attributes);
224
225 TokenStream::from(quote!(#outer_fn))
226}
227
228fn parse_int(literal: &Lit) -> Result<usize, TokenStream> {
229 let result = if let Lit::Int(int) = &literal {
230 int.base10_parse()
231 } else {
232 Err(syn::Error::new(literal.span(), "must be an integer"))
233 };
234
235 result.map_err(|err| TokenStream::from(err.into_compile_error()))
236}