1use proc_macro::TokenStream;
2use quote::{quote, quote_spanned};
3use syn::parse::{Parse, ParseStream, Result};
4use syn::punctuated::Punctuated;
5use syn::spanned::Spanned;
6use syn::{
7 AttrStyle, Attribute, Block, Error, Expr, Ident, Pat, ReturnType, Token, Type, braced,
8 parenthesized, parse_macro_input, parse_quote, token,
9};
10
11mod kw {
12 syn::custom_keyword!(query);
13}
14
15fn check_attributes(attrs: Vec<Attribute>) -> Result<Vec<Attribute>> {
17 let inner = |attr: Attribute| {
18 if !attr.path().is_ident("doc") {
19 Err(Error::new(attr.span(), "attributes not supported on queries"))
20 } else if attr.style != AttrStyle::Outer {
21 Err(Error::new(
22 attr.span(),
23 "attributes must be outer attributes (`///`), not inner attributes",
24 ))
25 } else {
26 Ok(attr)
27 }
28 };
29 attrs.into_iter().map(inner).collect()
30}
31
32struct Query {
34 doc_comments: Vec<Attribute>,
35 modifiers: QueryModifiers,
36 name: Ident,
37 key: Pat,
38 arg: Type,
39 result: ReturnType,
40}
41
42impl Parse for Query {
43 fn parse(input: ParseStream<'_>) -> Result<Self> {
44 let mut doc_comments = check_attributes(input.call(Attribute::parse_outer)?)?;
45
46 input.parse::<kw::query>()?;
48 let name: Ident = input.parse()?;
49 let arg_content;
50 parenthesized!(arg_content in input);
51 let key = Pat::parse_single(&arg_content)?;
52 arg_content.parse::<Token![:]>()?;
53 let arg = arg_content.parse()?;
54 let _ = arg_content.parse::<Option<Token![,]>>()?;
55 let result = input.parse()?;
56
57 let content;
59 braced!(content in input);
60 let modifiers = parse_query_modifiers(&content)?;
61
62 if doc_comments.is_empty() {
65 doc_comments.push(doc_comment_from_desc(&modifiers.desc.1)?);
66 }
67
68 Ok(Query { doc_comments, modifiers, name, key, arg, result })
69 }
70}
71
72struct List<T>(Vec<T>);
74
75impl<T: Parse> Parse for List<T> {
76 fn parse(input: ParseStream<'_>) -> Result<Self> {
77 let mut list = Vec::new();
78 while !input.is_empty() {
79 list.push(input.parse()?);
80 }
81 Ok(List(list))
82 }
83}
84
85struct QueryModifiers {
86 desc: (Option<Ident>, Punctuated<Expr, Token![,]>),
88
89 arena_cache: Option<Ident>,
91
92 cache: Option<(Option<Pat>, Block)>,
94
95 fatal_cycle: Option<Ident>,
97
98 cycle_delay_bug: Option<Ident>,
100
101 cycle_stash: Option<Ident>,
103
104 no_hash: Option<Ident>,
106
107 anon: Option<Ident>,
109
110 eval_always: Option<Ident>,
112
113 depth_limit: Option<Ident>,
115
116 separate_provide_extern: Option<Ident>,
118
119 feedable: Option<Ident>,
121
122 return_result_from_ensure_ok: Option<Ident>,
133}
134
135fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
136 let mut arena_cache = None;
137 let mut cache = None;
138 let mut desc = None;
139 let mut fatal_cycle = None;
140 let mut cycle_delay_bug = None;
141 let mut cycle_stash = None;
142 let mut no_hash = None;
143 let mut anon = None;
144 let mut eval_always = None;
145 let mut depth_limit = None;
146 let mut separate_provide_extern = None;
147 let mut feedable = None;
148 let mut return_result_from_ensure_ok = None;
149
150 while !input.is_empty() {
151 let modifier: Ident = input.parse()?;
152
153 macro_rules! try_insert {
154 ($name:ident = $expr:expr) => {
155 if $name.is_some() {
156 return Err(Error::new(modifier.span(), "duplicate modifier"));
157 }
158 $name = Some($expr);
159 };
160 }
161
162 if modifier == "desc" {
163 let attr_content;
166 braced!(attr_content in input);
167 let tcx = if attr_content.peek(Token![|]) {
168 attr_content.parse::<Token![|]>()?;
169 let tcx = attr_content.parse()?;
170 attr_content.parse::<Token![|]>()?;
171 Some(tcx)
172 } else {
173 None
174 };
175 let list = attr_content.parse_terminated(Expr::parse, Token![,])?;
176 try_insert!(desc = (tcx, list));
177 } else if modifier == "cache_on_disk_if" {
178 let args = if input.peek(token::Paren) {
181 let args;
182 parenthesized!(args in input);
183 let tcx = Pat::parse_single(&args)?;
184 Some(tcx)
185 } else {
186 None
187 };
188 let block = input.parse()?;
189 try_insert!(cache = (args, block));
190 } else if modifier == "arena_cache" {
191 try_insert!(arena_cache = modifier);
192 } else if modifier == "fatal_cycle" {
193 try_insert!(fatal_cycle = modifier);
194 } else if modifier == "cycle_delay_bug" {
195 try_insert!(cycle_delay_bug = modifier);
196 } else if modifier == "cycle_stash" {
197 try_insert!(cycle_stash = modifier);
198 } else if modifier == "no_hash" {
199 try_insert!(no_hash = modifier);
200 } else if modifier == "anon" {
201 try_insert!(anon = modifier);
202 } else if modifier == "eval_always" {
203 try_insert!(eval_always = modifier);
204 } else if modifier == "depth_limit" {
205 try_insert!(depth_limit = modifier);
206 } else if modifier == "separate_provide_extern" {
207 try_insert!(separate_provide_extern = modifier);
208 } else if modifier == "feedable" {
209 try_insert!(feedable = modifier);
210 } else if modifier == "return_result_from_ensure_ok" {
211 try_insert!(return_result_from_ensure_ok = modifier);
212 } else {
213 return Err(Error::new(modifier.span(), "unknown query modifier"));
214 }
215 }
216 let Some(desc) = desc else {
217 return Err(input.error("no description provided"));
218 };
219 Ok(QueryModifiers {
220 arena_cache,
221 cache,
222 desc,
223 fatal_cycle,
224 cycle_delay_bug,
225 cycle_stash,
226 no_hash,
227 anon,
228 eval_always,
229 depth_limit,
230 separate_provide_extern,
231 feedable,
232 return_result_from_ensure_ok,
233 })
234}
235
236fn doc_comment_from_desc(list: &Punctuated<Expr, token::Comma>) -> Result<Attribute> {
237 use ::syn::*;
238 let mut iter = list.iter();
239 let format_str: String = match iter.next() {
240 Some(&Expr::Lit(ExprLit { lit: Lit::Str(ref lit_str), .. })) => {
241 lit_str.value().replace("`{}`", "{}") }
243 _ => return Err(Error::new(list.span(), "Expected a string literal")),
244 };
245 let mut fmt_fragments = format_str.split("{}");
246 let mut doc_string = fmt_fragments.next().unwrap().to_string();
247 iter.map(::quote::ToTokens::to_token_stream).zip(fmt_fragments).for_each(
248 |(tts, next_fmt_fragment)| {
249 use ::core::fmt::Write;
250 write!(
251 &mut doc_string,
252 " `{}` {}",
253 tts.to_string().replace(" . ", "."),
254 next_fmt_fragment,
255 )
256 .unwrap();
257 },
258 );
259 let doc_string = format!("[query description - consider adding a doc-comment!] {doc_string}");
260 Ok(parse_quote! { #[doc = #doc_string] })
261}
262
263fn add_query_desc_cached_impl(
265 query: &Query,
266 descs: &mut proc_macro2::TokenStream,
267 cached: &mut proc_macro2::TokenStream,
268) {
269 let Query { name, key, modifiers, .. } = &query;
270
271 let ra_hint = quote! {
280 let crate::query::Providers { #name: _, .. };
281 };
282
283 let cache = if let Some((args, expr)) = modifiers.cache.as_ref() {
285 let tcx = args.as_ref().map(|t| quote! { #t }).unwrap_or_else(|| quote! { _ });
286 quote! {
290 #[allow(unused_variables, unused_braces, rustc::pass_by_value)]
291 #[inline]
292 pub fn #name<'tcx>(#tcx: TyCtxt<'tcx>, #key: &crate::query::queries::#name::Key<'tcx>) -> bool {
293 #ra_hint
294 #expr
295 }
296 }
297 } else {
298 quote! {
299 #[allow(rustc::pass_by_value)]
301 #[inline]
302 pub fn #name<'tcx>(_: TyCtxt<'tcx>, _: &crate::query::queries::#name::Key<'tcx>) -> bool {
303 #ra_hint
304 false
305 }
306 }
307 };
308
309 let (tcx, desc) = &modifiers.desc;
310 let tcx = tcx.as_ref().map_or_else(|| quote! { _ }, |t| quote! { #t });
311
312 let desc = quote! {
313 #[allow(unused_variables)]
314 pub fn #name<'tcx>(tcx: TyCtxt<'tcx>, key: crate::query::queries::#name::Key<'tcx>) -> String {
315 let (#tcx, #key) = (tcx, key);
316 ::rustc_middle::ty::print::with_no_trimmed_paths!(
317 format!(#desc)
318 )
319 }
320 };
321
322 descs.extend(quote! {
323 #desc
324 });
325
326 cached.extend(quote! {
327 #cache
328 });
329}
330
331pub(super) fn rustc_queries(input: TokenStream) -> TokenStream {
332 let queries = parse_macro_input!(input as List<Query>);
333
334 let mut query_stream = quote! {};
335 let mut query_description_stream = quote! {};
336 let mut query_cached_stream = quote! {};
337 let mut feedable_queries = quote! {};
338 let mut errors = quote! {};
339
340 macro_rules! assert {
341 ( $cond:expr, $span:expr, $( $tt:tt )+ ) => {
342 if !$cond {
343 errors.extend(
344 Error::new($span, format!($($tt)+)).into_compile_error(),
345 );
346 }
347 }
348 }
349
350 for query in queries.0 {
351 let Query { name, arg, modifiers, .. } = &query;
352 let result_full = &query.result;
353 let result = match query.result {
354 ReturnType::Default => quote! { -> () },
355 _ => quote! { #result_full },
356 };
357
358 let mut attributes = Vec::new();
359
360 macro_rules! passthrough {
361 ( $( $modifier:ident ),+ $(,)? ) => {
362 $( if let Some($modifier) = &modifiers.$modifier {
363 attributes.push(quote! { (#$modifier) });
364 }; )+
365 }
366 }
367
368 passthrough!(
369 fatal_cycle,
370 arena_cache,
371 cycle_delay_bug,
372 cycle_stash,
373 no_hash,
374 anon,
375 eval_always,
376 depth_limit,
377 separate_provide_extern,
378 return_result_from_ensure_ok,
379 );
380
381 if modifiers.cache.is_some() {
382 attributes.push(quote! { (cache) });
383 }
384 if modifiers.cache.is_some() {
386 attributes.push(quote! { (cache) });
387 }
388
389 let span = name.span();
396 let attribute_stream = quote_spanned! {span=> #(#attributes),*};
397 let doc_comments = &query.doc_comments;
398 query_stream.extend(quote! {
400 #(#doc_comments)*
401 [#attribute_stream] fn #name(#arg) #result,
402 });
403
404 if let Some(feedable) = &modifiers.feedable {
405 assert!(
406 modifiers.anon.is_none(),
407 feedable.span(),
408 "Query {name} cannot be both `feedable` and `anon`."
409 );
410 assert!(
411 modifiers.eval_always.is_none(),
412 feedable.span(),
413 "Query {name} cannot be both `feedable` and `eval_always`."
414 );
415 feedable_queries.extend(quote! {
416 #(#doc_comments)*
417 [#attribute_stream] fn #name(#arg) #result,
418 });
419 }
420
421 add_query_desc_cached_impl(&query, &mut query_description_stream, &mut query_cached_stream);
422 }
423
424 TokenStream::from(quote! {
425 #[macro_export]
431 macro_rules! rustc_with_all_queries {
432 (
433 $macro:ident!
435
436 $( [$($extra_fake_queries:tt)*] )?
439 ) => {
440 $macro! {
441 $( $($extra_fake_queries)* )?
442 #query_stream
443 }
444 }
445 }
446 macro_rules! rustc_feedable_queries {
447 ( $macro:ident! ) => {
448 $macro!(#feedable_queries);
449 }
450 }
451 pub mod descs {
452 use super::*;
453 #query_description_stream
454 }
455 pub mod cached {
456 use super::*;
457 #query_cached_stream
458 }
459 #errors
460 })
461}