compiler/rustc_macros/src/query.rs

   1 use proc_macro::TokenStream;
   2 use quote::{quote, quote_spanned};
   3 use syn::parse::{Parse, ParseStream, Result};
   4 use syn::punctuated::Punctuated;
   5 use syn::spanned::Spanned;
   6 use syn::{
   7     braced, parenthesized, parse_macro_input, parse_quote, token, AttrStyle, Attribute, Block,
   8     Error, Expr, Ident, Pat, ReturnType, Token, Type,
   9 };
  10
  11 mod kw {
  12     syn::custom_keyword!(query);
  13 }
  14
  15 /// Ensures only doc comment attributes are used
  16 fn 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
  32 /// A compiler query. `query ... { ... }`
  33 struct Query {
  34     doc_comments: Vec<Attribute>,
  35     modifiers: QueryModifiers,
  36     name: Ident,
  37     key: Pat,
  38     arg: Type,
  39     result: ReturnType,
  40 }
  41
  42 impl Parse for Query {
  43     fn parse(input: ParseStream<'_>) -> Result<Self> {
  44         let mut doc_comments = check_attributes(input.call(Attribute::parse_outer)?)?;
  45
  46         // Parse the query declaration. Like `query type_of(key: DefId) -> Ty<'tcx>`
  47         input.parse::<kw::query>()?;
  48         let name: Ident = input.parse()?;
  49         let arg_content;
  50         parenthesized!(arg_content in input);
  51         let key = arg_content.parse()?;
  52         arg_content.parse::<Token![:]>()?;
  53         let arg = arg_content.parse()?;
  54         let result = input.parse()?;
  55
  56         // Parse the query modifiers
  57         let content;
  58         braced!(content in input);
  59         let modifiers = parse_query_modifiers(&content)?;
  60
  61         // If there are no doc-comments, give at least some idea of what
  62         // it does by showing the query description.
  63         if doc_comments.is_empty() {
  64             doc_comments.push(doc_comment_from_desc(&modifiers.desc.1)?);
  65         }
  66
  67         Ok(Query { doc_comments, modifiers, name, key, arg, result })
  68     }
  69 }
  70
  71 /// A type used to greedily parse another type until the input is empty.
  72 struct List<T>(Vec<T>);
  73
  74 impl<T: Parse> Parse for List<T> {
  75     fn parse(input: ParseStream<'_>) -> Result<Self> {
  76         let mut list = Vec::new();
  77         while !input.is_empty() {
  78             list.push(input.parse()?);
  79         }
  80         Ok(List(list))
  81     }
  82 }
  83
  84 struct QueryModifiers {
  85     /// The description of the query.
  86     desc: (Option<Ident>, Punctuated<Expr, Token![,]>),
  87
  88     /// Use this type for the in-memory cache.
  89     storage: Option<Type>,
  90
  91     /// Cache the query to disk if the `Block` returns true.
  92     cache: Option<(Option<Pat>, Block)>,
  93
  94     /// Custom code to load the query from disk.
  95     load_cached: Option<(Ident, Ident, Block)>,
  96
  97     /// A cycle error for this query aborting the compilation with a fatal error.
  98     fatal_cycle: Option<Ident>,
  99
 100     /// A cycle error results in a delay_bug call
 101     cycle_delay_bug: Option<Ident>,
 102
 103     /// Don't hash the result, instead just mark a query red if it runs
 104     no_hash: Option<Ident>,
 105
 106     /// Generate a dep node based on the dependencies of the query
 107     anon: Option<Ident>,
 108
 109     /// Always evaluate the query, ignoring its dependencies
 110     eval_always: Option<Ident>,
 111
 112     /// Whether the query has a call depth limit
 113     depth_limit: Option<Ident>,
 114
 115     /// Use a separate query provider for local and extern crates
 116     separate_provide_extern: Option<Ident>,
 117
 118     /// Always remap the ParamEnv's constness before hashing.
 119     remap_env_constness: Option<Ident>,
 120 }
 121
 122 fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
 123     let mut load_cached = None;
 124     let mut storage = None;
 125     let mut cache = None;
 126     let mut desc = None;
 127     let mut fatal_cycle = None;
 128     let mut cycle_delay_bug = None;
 129     let mut no_hash = None;
 130     let mut anon = None;
 131     let mut eval_always = None;
 132     let mut depth_limit = None;
 133     let mut separate_provide_extern = None;
 134     let mut remap_env_constness = None;
 135
 136     while !input.is_empty() {
 137         let modifier: Ident = input.parse()?;
 138
 139         macro_rules! try_insert {
 140             ($name:ident = $expr:expr) => {
 141                 if $name.is_some() {
 142                     return Err(Error::new(modifier.span(), "duplicate modifier"));
 143                 }
 144                 $name = Some($expr);
 145             };
 146         }
 147
 148         if modifier == "desc" {
 149             // Parse a description modifier like:
 150             // `desc { |tcx| "foo {}", tcx.item_path(key) }`
 151             let attr_content;
 152             braced!(attr_content in input);
 153             let tcx = if attr_content.peek(Token![|]) {
 154                 attr_content.parse::<Token![|]>()?;
 155                 let tcx = attr_content.parse()?;
 156                 attr_content.parse::<Token![|]>()?;
 157                 Some(tcx)
 158             } else {
 159                 None
 160             };
 161             let list = attr_content.parse_terminated(Expr::parse)?;
 162             try_insert!(desc = (tcx, list));
 163         } else if modifier == "cache_on_disk_if" {
 164             // Parse a cache modifier like:
 165             // `cache(tcx) { |tcx| key.is_local() }`
 166             let args = if input.peek(token::Paren) {
 167                 let args;
 168                 parenthesized!(args in input);
 169                 let tcx = args.parse()?;
 170                 Some(tcx)
 171             } else {
 172                 None
 173             };
 174             let block = input.parse()?;
 175             try_insert!(cache = (args, block));
 176         } else if modifier == "load_cached" {
 177             // Parse a load_cached modifier like:
 178             // `load_cached(tcx, id) { tcx.on_disk_cache.try_load_query_result(tcx, id) }`
 179             let args;
 180             parenthesized!(args in input);
 181             let tcx = args.parse()?;
 182             args.parse::<Token![,]>()?;
 183             let id = args.parse()?;
 184             let block = input.parse()?;
 185             try_insert!(load_cached = (tcx, id, block));
 186         } else if modifier == "storage" {
 187             let args;
 188             parenthesized!(args in input);
 189             let ty = args.parse()?;
 190             try_insert!(storage = ty);
 191         } else if modifier == "fatal_cycle" {
 192             try_insert!(fatal_cycle = modifier);
 193         } else if modifier == "cycle_delay_bug" {
 194             try_insert!(cycle_delay_bug = modifier);
 195         } else if modifier == "no_hash" {
 196             try_insert!(no_hash = modifier);
 197         } else if modifier == "anon" {
 198             try_insert!(anon = modifier);
 199         } else if modifier == "eval_always" {
 200             try_insert!(eval_always = modifier);
 201         } else if modifier == "depth_limit" {
 202             try_insert!(depth_limit = modifier);
 203         } else if modifier == "separate_provide_extern" {
 204             try_insert!(separate_provide_extern = modifier);
 205         } else if modifier == "remap_env_constness" {
 206             try_insert!(remap_env_constness = modifier);
 207         } else {
 208             return Err(Error::new(modifier.span(), "unknown query modifier"));
 209         }
 210     }
 211     let Some(desc) = desc else {
 212         return Err(input.error("no description provided"));
 213     };
 214     Ok(QueryModifiers {
 215         load_cached,
 216         storage,
 217         cache,
 218         desc,
 219         fatal_cycle,
 220         cycle_delay_bug,
 221         no_hash,
 222         anon,
 223         eval_always,
 224         depth_limit,
 225         separate_provide_extern,
 226         remap_env_constness,
 227     })
 228 }
 229
 230 fn doc_comment_from_desc(list: &Punctuated<Expr, token::Comma>) -> Result<Attribute> {
 231     use ::syn::*;
 232     let mut iter = list.iter();
 233     let format_str: String = match iter.next() {
 234         Some(&Expr::Lit(ExprLit { lit: Lit::Str(ref lit_str), .. })) => {
 235             lit_str.value().replace("`{}`", "{}") // We add them later anyways for consistency
 236         }
 237         _ => return Err(Error::new(list.span(), "Expected a string literal")),
 238     };
 239     let mut fmt_fragments = format_str.split("{}");
 240     let mut doc_string = fmt_fragments.next().unwrap().to_string();
 241     iter.map(::quote::ToTokens::to_token_stream).zip(fmt_fragments).for_each(
 242         |(tts, next_fmt_fragment)| {
 243             use ::core::fmt::Write;
 244             write!(
 245                 &mut doc_string,
 246                 " `{}` {}",
 247                 tts.to_string().replace(" . ", "."),
 248                 next_fmt_fragment,
 249             )
 250             .unwrap();
 251         },
 252     );
 253     let doc_string = format!("[query description - consider adding a doc-comment!] {}", doc_string);
 254     Ok(parse_quote! { #[doc = #doc_string] })
 255 }
 256
 257 /// Add the impl of QueryDescription for the query to `impls` if one is requested
 258 fn add_query_description_impl(query: &Query, impls: &mut proc_macro2::TokenStream) {
 259     let name = &query.name;
 260     let key = &query.key;
 261     let modifiers = &query.modifiers;
 262
 263     // Find out if we should cache the query on disk
 264     let cache = if let Some((args, expr)) = modifiers.cache.as_ref() {
 265         let try_load_from_disk = if let Some((tcx, id, block)) = modifiers.load_cached.as_ref() {
 266             // Use custom code to load the query from disk
 267             quote! {
 268                 const TRY_LOAD_FROM_DISK: Option<fn(QueryCtxt<'tcx>, SerializedDepNodeIndex) -> Option<Self::Value>>
 269                     = Some(|#tcx, #id| { #block });
 270             }
 271         } else {
 272             // Use the default code to load the query from disk
 273             quote! {
 274                 const TRY_LOAD_FROM_DISK: Option<fn(QueryCtxt<'tcx>, SerializedDepNodeIndex) -> Option<Self::Value>>
 275                     = Some(|tcx, id| tcx.on_disk_cache().as_ref()?.try_load_query_result(*tcx, id));
 276             }
 277         };
 278
 279         let tcx = args.as_ref().map(|t| quote! { #t }).unwrap_or_else(|| quote! { _ });
 280         // expr is a `Block`, meaning that `{ #expr }` gets expanded
 281         // to `{ { stmts... } }`, which triggers the `unused_braces` lint.
 282         quote! {
 283             #[allow(unused_variables, unused_braces)]
 284             #[inline]
 285             fn cache_on_disk(#tcx: TyCtxt<'tcx>, #key: &Self::Key) -> bool {
 286                 #expr
 287             }
 288
 289             #try_load_from_disk
 290         }
 291     } else {
 292         if modifiers.load_cached.is_some() {
 293             panic!("load_cached modifier on query `{}` without a cache modifier", name);
 294         }
 295         quote! {
 296             #[inline]
 297             fn cache_on_disk(_: TyCtxt<'tcx>, _: &Self::Key) -> bool {
 298                 false
 299             }
 300
 301             const TRY_LOAD_FROM_DISK: Option<fn(QueryCtxt<'tcx>, SerializedDepNodeIndex) -> Option<Self::Value>> = None;
 302         }
 303     };
 304
 305     let (tcx, desc) = &modifiers.desc;
 306     let tcx = tcx.as_ref().map_or_else(|| quote! { _ }, |t| quote! { #t });
 307
 308     let desc = quote! {
 309         #[allow(unused_variables)]
 310         fn describe(tcx: QueryCtxt<'tcx>, key: Self::Key) -> String {
 311             let (#tcx, #key) = (*tcx, key);
 312             ::rustc_middle::ty::print::with_no_trimmed_paths!(
 313                 format!(#desc)
 314             )
 315         }
 316     };
 317
 318     impls.extend(quote! {
 319         (#name) => {
 320             #desc
 321             #cache
 322         };
 323     });
 324 }
 325
 326 pub fn rustc_queries(input: TokenStream) -> TokenStream {
 327     let queries = parse_macro_input!(input as List<Query>);
 328
 329     let mut query_stream = quote! {};
 330     let mut query_description_stream = quote! {};
 331     let mut dep_node_def_stream = quote! {};
 332     let mut cached_queries = quote! {};
 333
 334     for query in queries.0 {
 335         let Query { name, arg, modifiers, .. } = &query;
 336         let result_full = &query.result;
 337         let result = match query.result {
 338             ReturnType::Default => quote! { -> () },
 339             _ => quote! { #result_full },
 340         };
 341
 342         if modifiers.cache.is_some() {
 343             cached_queries.extend(quote! {
 344                 #name,
 345             });
 346         }
 347
 348         let mut attributes = Vec::new();
 349
 350         // Pass on the fatal_cycle modifier
 351         if let Some(fatal_cycle) = &modifiers.fatal_cycle {
 352             attributes.push(quote! { (#fatal_cycle) });
 353         };
 354         // Pass on the storage modifier
 355         if let Some(ref ty) = modifiers.storage {
 356             let span = ty.span();
 357             attributes.push(quote_spanned! {span=> (storage #ty) });
 358         };
 359         // Pass on the cycle_delay_bug modifier
 360         if let Some(cycle_delay_bug) = &modifiers.cycle_delay_bug {
 361             attributes.push(quote! { (#cycle_delay_bug) });
 362         };
 363         // Pass on the no_hash modifier
 364         if let Some(no_hash) = &modifiers.no_hash {
 365             attributes.push(quote! { (#no_hash) });
 366         };
 367         // Pass on the anon modifier
 368         if let Some(anon) = &modifiers.anon {
 369             attributes.push(quote! { (#anon) });
 370         };
 371         // Pass on the eval_always modifier
 372         if let Some(eval_always) = &modifiers.eval_always {
 373             attributes.push(quote! { (#eval_always) });
 374         };
 375         // Pass on the depth_limit modifier
 376         if let Some(depth_limit) = &modifiers.depth_limit {
 377             attributes.push(quote! { (#depth_limit) });
 378         };
 379         // Pass on the separate_provide_extern modifier
 380         if let Some(separate_provide_extern) = &modifiers.separate_provide_extern {
 381             attributes.push(quote! { (#separate_provide_extern) });
 382         }
 383         // Pass on the remap_env_constness modifier
 384         if let Some(remap_env_constness) = &modifiers.remap_env_constness {
 385             attributes.push(quote! { (#remap_env_constness) });
 386         }
 387
 388         // This uses the span of the query definition for the commas,
 389         // which can be important if we later encounter any ambiguity
 390         // errors with any of the numerous macro_rules! macros that
 391         // we use. Using the call-site span would result in a span pointing
 392         // at the entire `rustc_queries!` invocation, which wouldn't
 393         // be very useful.
 394         let span = name.span();
 395         let attribute_stream = quote_spanned! {span=> #(#attributes),*};
 396         let doc_comments = query.doc_comments.iter();
 397         // Add the query to the group
 398         query_stream.extend(quote! {
 399             #(#doc_comments)*
 400             [#attribute_stream] fn #name(#arg) #result,
 401         });
 402
 403         // Create a dep node for the query
 404         dep_node_def_stream.extend(quote! {
 405             [#attribute_stream] #name(#arg),
 406         });
 407
 408         add_query_description_impl(&query, &mut query_description_stream);
 409     }
 410
 411     TokenStream::from(quote! {
 412         #[macro_export]
 413         macro_rules! rustc_query_append {
 414             ($macro:ident !) => {
 415                 $macro! {
 416                     #query_stream
 417                 }
 418             }
 419         }
 420         macro_rules! rustc_dep_node_append {
 421             ($macro:ident! [$($other:tt)*]) => {
 422                 $macro!(
 423                     $($other)*
 424
 425                     #dep_node_def_stream
 426                 );
 427             }
 428         }
 429         #[macro_export]
 430         macro_rules! rustc_cached_queries {
 431             ( $macro:ident! ) => {
 432                 $macro!(#cached_queries);
 433             }
 434         }
 435         #[macro_export]
 436         macro_rules! rustc_query_description {
 437             #query_description_stream
 438         }
 439     })
 440 }