src/librustc_macros/src/query.rs

   1 use proc_macro::TokenStream;
   2 use proc_macro2::{Delimiter, TokenTree};
   3 use quote::quote;
   4 use syn::parse::{Parse, ParseStream, Result};
   5 use syn::punctuated::Punctuated;
   6 use syn::spanned::Spanned;
   7 use syn::{
   8     braced, parenthesized, parse_macro_input, Attribute, Block, Error, Expr, Ident, ReturnType,
   9     Token, Type,
  10 };
  11
  12 #[allow(non_camel_case_types)]
  13 mod kw {
  14     syn::custom_keyword!(query);
  15 }
  16
  17 /// Ident or a wildcard `_`.
  18 struct IdentOrWild(Ident);
  19
  20 impl Parse for IdentOrWild {
  21     fn parse(input: ParseStream<'_>) -> Result<Self> {
  22         Ok(if input.peek(Token![_]) {
  23             let underscore = input.parse::<Token![_]>()?;
  24             IdentOrWild(Ident::new("_", underscore.span()))
  25         } else {
  26             IdentOrWild(input.parse()?)
  27         })
  28     }
  29 }
  30
  31 /// A modifier for a query
  32 enum QueryModifier {
  33     /// The description of the query.
  34     Desc(Option<Ident>, Punctuated<Expr, Token![,]>),
  35
  36     /// Use this type for the in-memory cache.
  37     Storage(Type),
  38
  39     /// Cache the query to disk if the `Expr` returns true.
  40     Cache(Option<(IdentOrWild, IdentOrWild)>, Block),
  41
  42     /// Custom code to load the query from disk.
  43     LoadCached(Ident, Ident, Block),
  44
  45     /// A cycle error for this query aborting the compilation with a fatal error.
  46     FatalCycle,
  47
  48     /// A cycle error results in a delay_bug call
  49     CycleDelayBug,
  50
  51     /// Don't hash the result, instead just mark a query red if it runs
  52     NoHash,
  53
  54     /// Don't force the query
  55     NoForce,
  56
  57     /// Generate a dep node based on the dependencies of the query
  58     Anon,
  59
  60     /// Always evaluate the query, ignoring its depdendencies
  61     EvalAlways,
  62 }
  63
  64 impl Parse for QueryModifier {
  65     fn parse(input: ParseStream<'_>) -> Result<Self> {
  66         let modifier: Ident = input.parse()?;
  67         if modifier == "desc" {
  68             // Parse a description modifier like:
  69             // `desc { |tcx| "foo {}", tcx.item_path(key) }`
  70             let attr_content;
  71             braced!(attr_content in input);
  72             let tcx = if attr_content.peek(Token![|]) {
  73                 attr_content.parse::<Token![|]>()?;
  74                 let tcx = attr_content.parse()?;
  75                 attr_content.parse::<Token![|]>()?;
  76                 Some(tcx)
  77             } else {
  78                 None
  79             };
  80             let desc = attr_content.parse_terminated(Expr::parse)?;
  81             Ok(QueryModifier::Desc(tcx, desc))
  82         } else if modifier == "cache_on_disk_if" {
  83             // Parse a cache modifier like:
  84             // `cache(tcx, value) { |tcx| key.is_local() }`
  85             let has_args = if let TokenTree::Group(group) = input.fork().parse()? {
  86                 group.delimiter() == Delimiter::Parenthesis
  87             } else {
  88                 false
  89             };
  90             let args = if has_args {
  91                 let args;
  92                 parenthesized!(args in input);
  93                 let tcx = args.parse()?;
  94                 args.parse::<Token![,]>()?;
  95                 let value = args.parse()?;
  96                 Some((tcx, value))
  97             } else {
  98                 None
  99             };
 100             let block = input.parse()?;
 101             Ok(QueryModifier::Cache(args, block))
 102         } else if modifier == "load_cached" {
 103             // Parse a load_cached modifier like:
 104             // `load_cached(tcx, id) { tcx.queries.on_disk_cache.try_load_query_result(tcx, id) }`
 105             let args;
 106             parenthesized!(args in input);
 107             let tcx = args.parse()?;
 108             args.parse::<Token![,]>()?;
 109             let id = args.parse()?;
 110             let block = input.parse()?;
 111             Ok(QueryModifier::LoadCached(tcx, id, block))
 112         } else if modifier == "storage" {
 113             let ty = input.parse()?;
 114             Ok(QueryModifier::Storage(ty))
 115         } else if modifier == "fatal_cycle" {
 116             Ok(QueryModifier::FatalCycle)
 117         } else if modifier == "cycle_delay_bug" {
 118             Ok(QueryModifier::CycleDelayBug)
 119         } else if modifier == "no_hash" {
 120             Ok(QueryModifier::NoHash)
 121         } else if modifier == "no_force" {
 122             Ok(QueryModifier::NoForce)
 123         } else if modifier == "anon" {
 124             Ok(QueryModifier::Anon)
 125         } else if modifier == "eval_always" {
 126             Ok(QueryModifier::EvalAlways)
 127         } else {
 128             Err(Error::new(modifier.span(), "unknown query modifier"))
 129         }
 130     }
 131 }
 132
 133 /// Ensures only doc comment attributes are used
 134 fn check_attributes(attrs: Vec<Attribute>) -> Result<()> {
 135     for attr in attrs {
 136         if !attr.path.is_ident("doc") {
 137             return Err(Error::new(attr.span(), "attributes not supported on queries"));
 138         }
 139     }
 140     Ok(())
 141 }
 142
 143 /// A compiler query. `query ... { ... }`
 144 struct Query {
 145     modifiers: List<QueryModifier>,
 146     name: Ident,
 147     key: IdentOrWild,
 148     arg: Type,
 149     result: ReturnType,
 150 }
 151
 152 impl Parse for Query {
 153     fn parse(input: ParseStream<'_>) -> Result<Self> {
 154         check_attributes(input.call(Attribute::parse_outer)?)?;
 155
 156         // Parse the query declaration. Like `query type_of(key: DefId) -> Ty<'tcx>`
 157         input.parse::<kw::query>()?;
 158         let name: Ident = input.parse()?;
 159         let arg_content;
 160         parenthesized!(arg_content in input);
 161         let key = arg_content.parse()?;
 162         arg_content.parse::<Token![:]>()?;
 163         let arg = arg_content.parse()?;
 164         let result = input.parse()?;
 165
 166         // Parse the query modifiers
 167         let content;
 168         braced!(content in input);
 169         let modifiers = content.parse()?;
 170
 171         Ok(Query { modifiers, name, key, arg, result })
 172     }
 173 }
 174
 175 /// A type used to greedily parse another type until the input is empty.
 176 struct List<T>(Vec<T>);
 177
 178 impl<T: Parse> Parse for List<T> {
 179     fn parse(input: ParseStream<'_>) -> Result<Self> {
 180         let mut list = Vec::new();
 181         while !input.is_empty() {
 182             list.push(input.parse()?);
 183         }
 184         Ok(List(list))
 185     }
 186 }
 187
 188 /// A named group containing queries.
 189 struct Group {
 190     name: Ident,
 191     queries: List<Query>,
 192 }
 193
 194 impl Parse for Group {
 195     fn parse(input: ParseStream<'_>) -> Result<Self> {
 196         let name: Ident = input.parse()?;
 197         let content;
 198         braced!(content in input);
 199         Ok(Group { name, queries: content.parse()? })
 200     }
 201 }
 202
 203 struct QueryModifiers {
 204     /// The description of the query.
 205     desc: Option<(Option<Ident>, Punctuated<Expr, Token![,]>)>,
 206
 207     /// Use this type for the in-memory cache.
 208     storage: Option<Type>,
 209
 210     /// Cache the query to disk if the `Block` returns true.
 211     cache: Option<(Option<(IdentOrWild, IdentOrWild)>, Block)>,
 212
 213     /// Custom code to load the query from disk.
 214     load_cached: Option<(Ident, Ident, Block)>,
 215
 216     /// A cycle error for this query aborting the compilation with a fatal error.
 217     fatal_cycle: bool,
 218
 219     /// A cycle error results in a delay_bug call
 220     cycle_delay_bug: bool,
 221
 222     /// Don't hash the result, instead just mark a query red if it runs
 223     no_hash: bool,
 224
 225     /// Don't force the query
 226     no_force: bool,
 227
 228     /// Generate a dep node based on the dependencies of the query
 229     anon: bool,
 230
 231     // Always evaluate the query, ignoring its depdendencies
 232     eval_always: bool,
 233 }
 234
 235 /// Process query modifiers into a struct, erroring on duplicates
 236 fn process_modifiers(query: &mut Query) -> QueryModifiers {
 237     let mut load_cached = None;
 238     let mut storage = None;
 239     let mut cache = None;
 240     let mut desc = None;
 241     let mut fatal_cycle = false;
 242     let mut cycle_delay_bug = false;
 243     let mut no_hash = false;
 244     let mut no_force = false;
 245     let mut anon = false;
 246     let mut eval_always = false;
 247     for modifier in query.modifiers.0.drain(..) {
 248         match modifier {
 249             QueryModifier::LoadCached(tcx, id, block) => {
 250                 if load_cached.is_some() {
 251                     panic!("duplicate modifier `load_cached` for query `{}`", query.name);
 252                 }
 253                 load_cached = Some((tcx, id, block));
 254             }
 255             QueryModifier::Storage(ty) => {
 256                 if storage.is_some() {
 257                     panic!("duplicate modifier `storage` for query `{}`", query.name);
 258                 }
 259                 storage = Some(ty);
 260             }
 261             QueryModifier::Cache(args, expr) => {
 262                 if cache.is_some() {
 263                     panic!("duplicate modifier `cache` for query `{}`", query.name);
 264                 }
 265                 cache = Some((args, expr));
 266             }
 267             QueryModifier::Desc(tcx, list) => {
 268                 if desc.is_some() {
 269                     panic!("duplicate modifier `desc` for query `{}`", query.name);
 270                 }
 271                 desc = Some((tcx, list));
 272             }
 273             QueryModifier::FatalCycle => {
 274                 if fatal_cycle {
 275                     panic!("duplicate modifier `fatal_cycle` for query `{}`", query.name);
 276                 }
 277                 fatal_cycle = true;
 278             }
 279             QueryModifier::CycleDelayBug => {
 280                 if cycle_delay_bug {
 281                     panic!("duplicate modifier `cycle_delay_bug` for query `{}`", query.name);
 282                 }
 283                 cycle_delay_bug = true;
 284             }
 285             QueryModifier::NoHash => {
 286                 if no_hash {
 287                     panic!("duplicate modifier `no_hash` for query `{}`", query.name);
 288                 }
 289                 no_hash = true;
 290             }
 291             QueryModifier::NoForce => {
 292                 if no_force {
 293                     panic!("duplicate modifier `no_force` for query `{}`", query.name);
 294                 }
 295                 no_force = true;
 296             }
 297             QueryModifier::Anon => {
 298                 if anon {
 299                     panic!("duplicate modifier `anon` for query `{}`", query.name);
 300                 }
 301                 anon = true;
 302             }
 303             QueryModifier::EvalAlways => {
 304                 if eval_always {
 305                     panic!("duplicate modifier `eval_always` for query `{}`", query.name);
 306                 }
 307                 eval_always = true;
 308             }
 309         }
 310     }
 311     QueryModifiers {
 312         load_cached,
 313         storage,
 314         cache,
 315         desc,
 316         fatal_cycle,
 317         cycle_delay_bug,
 318         no_hash,
 319         no_force,
 320         anon,
 321         eval_always,
 322     }
 323 }
 324
 325 /// Add the impl of QueryDescription for the query to `impls` if one is requested
 326 fn add_query_description_impl(
 327     query: &Query,
 328     modifiers: QueryModifiers,
 329     impls: &mut proc_macro2::TokenStream,
 330 ) {
 331     let name = &query.name;
 332     let arg = &query.arg;
 333     let key = &query.key.0;
 334
 335     // Find out if we should cache the query on disk
 336     let cache = modifiers.cache.as_ref().map(|(args, expr)| {
 337         let try_load_from_disk = if let Some((tcx, id, block)) = modifiers.load_cached.as_ref() {
 338             // Use custom code to load the query from disk
 339             quote! {
 340                 #[inline]
 341                 fn try_load_from_disk(
 342                     #tcx: TyCtxt<'tcx>,
 343                     #id: SerializedDepNodeIndex
 344                 ) -> Option<Self::Value> {
 345                     #block
 346                 }
 347             }
 348         } else {
 349             // Use the default code to load the query from disk
 350             quote! {
 351                 #[inline]
 352                 fn try_load_from_disk(
 353                     tcx: TyCtxt<'tcx>,
 354                     id: SerializedDepNodeIndex
 355                 ) -> Option<Self::Value> {
 356                     tcx.queries.on_disk_cache.try_load_query_result(tcx, id)
 357                 }
 358             }
 359         };
 360
 361         let tcx = args
 362             .as_ref()
 363             .map(|t| {
 364                 let t = &(t.0).0;
 365                 quote! { #t }
 366             })
 367             .unwrap_or(quote! { _ });
 368         let value = args
 369             .as_ref()
 370             .map(|t| {
 371                 let t = &(t.1).0;
 372                 quote! { #t }
 373             })
 374             .unwrap_or(quote! { _ });
 375         quote! {
 376             #[inline]
 377             #[allow(unused_variables)]
 378             fn cache_on_disk(
 379                 #tcx: TyCtxt<'tcx>,
 380                 #key: Self::Key,
 381                 #value: Option<&Self::Value>
 382             ) -> bool {
 383                 #expr
 384             }
 385
 386             #try_load_from_disk
 387         }
 388     });
 389
 390     if cache.is_none() && modifiers.load_cached.is_some() {
 391         panic!("load_cached modifier on query `{}` without a cache modifier", name);
 392     }
 393
 394     let desc = modifiers.desc.as_ref().map(|(tcx, desc)| {
 395         let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
 396         quote! {
 397             #[allow(unused_variables)]
 398             fn describe(
 399                 #tcx: TyCtxt<'_>,
 400                 #key: #arg,
 401             ) -> Cow<'static, str> {
 402                 format!(#desc).into()
 403             }
 404         }
 405     });
 406
 407     if desc.is_some() || cache.is_some() {
 408         let cache = cache.unwrap_or(quote! {});
 409         let desc = desc.unwrap_or(quote! {});
 410
 411         impls.extend(quote! {
 412             impl<'tcx> QueryDescription<'tcx> for queries::#name<'tcx> {
 413                 #desc
 414                 #cache
 415             }
 416         });
 417     }
 418 }
 419
 420 pub fn rustc_queries(input: TokenStream) -> TokenStream {
 421     let groups = parse_macro_input!(input as List<Group>);
 422
 423     let mut query_stream = quote! {};
 424     let mut query_description_stream = quote! {};
 425     let mut dep_node_def_stream = quote! {};
 426     let mut dep_node_force_stream = quote! {};
 427     let mut try_load_from_on_disk_cache_stream = quote! {};
 428     let mut no_force_queries = Vec::new();
 429     let mut cached_queries = quote! {};
 430
 431     for group in groups.0 {
 432         let mut group_stream = quote! {};
 433         for mut query in group.queries.0 {
 434             let modifiers = process_modifiers(&mut query);
 435             let name = &query.name;
 436             let arg = &query.arg;
 437             let result_full = &query.result;
 438             let result = match query.result {
 439                 ReturnType::Default => quote! { -> () },
 440                 _ => quote! { #result_full },
 441             };
 442
 443             if modifiers.cache.is_some() {
 444                 cached_queries.extend(quote! {
 445                     #name,
 446                 });
 447             }
 448
 449             if modifiers.cache.is_some() && !modifiers.no_force {
 450                 try_load_from_on_disk_cache_stream.extend(quote! {
 451                     DepKind::#name => {
 452                         debug_assert!(tcx.dep_graph
 453                                          .node_color(self)
 454                                          .map(|c| c.is_green())
 455                                          .unwrap_or(false));
 456
 457                         let key = RecoverKey::recover(tcx, self).unwrap();
 458                         if queries::#name::cache_on_disk(tcx, key, None) {
 459                             let _ = tcx.#name(key);
 460                         }
 461                     }
 462                 });
 463             }
 464
 465             let mut attributes = Vec::new();
 466
 467             // Pass on the fatal_cycle modifier
 468             if modifiers.fatal_cycle {
 469                 attributes.push(quote! { fatal_cycle });
 470             };
 471             // Pass on the storage modifier
 472             if let Some(ref ty) = modifiers.storage {
 473                 attributes.push(quote! { storage(#ty) });
 474             };
 475             // Pass on the cycle_delay_bug modifier
 476             if modifiers.cycle_delay_bug {
 477                 attributes.push(quote! { cycle_delay_bug });
 478             };
 479             // Pass on the no_hash modifier
 480             if modifiers.no_hash {
 481                 attributes.push(quote! { no_hash });
 482             };
 483             // Pass on the anon modifier
 484             if modifiers.anon {
 485                 attributes.push(quote! { anon });
 486             };
 487             // Pass on the eval_always modifier
 488             if modifiers.eval_always {
 489                 attributes.push(quote! { eval_always });
 490             };
 491
 492             let attribute_stream = quote! {#(#attributes),*};
 493
 494             // Add the query to the group
 495             group_stream.extend(quote! {
 496                 [#attribute_stream] fn #name: #name(#arg) #result,
 497             });
 498
 499             // Create a dep node for the query
 500             dep_node_def_stream.extend(quote! {
 501                 [#attribute_stream] #name(#arg),
 502             });
 503
 504             if modifiers.no_force {
 505                 no_force_queries.push(name.clone());
 506             } else {
 507                 // Add a match arm to force the query given the dep node
 508                 dep_node_force_stream.extend(quote! {
 509                     DepKind::#name => {
 510                         if let Some(key) = RecoverKey::recover($tcx, $dep_node) {
 511                             $tcx.force_query::<crate::ty::query::queries::#name<'_>>(
 512                                 key,
 513                                 DUMMY_SP,
 514                                 *$dep_node
 515                             );
 516                         } else {
 517                             return false;
 518                         }
 519                     }
 520                 });
 521             }
 522
 523             add_query_description_impl(&query, modifiers, &mut query_description_stream);
 524         }
 525         let name = &group.name;
 526         query_stream.extend(quote! {
 527             #name { #group_stream },
 528         });
 529     }
 530
 531     // Add an arm for the no force queries to panic when trying to force them
 532     for query in no_force_queries {
 533         dep_node_force_stream.extend(quote! {
 534             DepKind::#query |
 535         });
 536     }
 537     dep_node_force_stream.extend(quote! {
 538         DepKind::Null => {
 539             bug!("Cannot force dep node: {:?}", $dep_node)
 540         }
 541     });
 542
 543     TokenStream::from(quote! {
 544         macro_rules! rustc_query_append {
 545             ([$($macro:tt)*][$($other:tt)*]) => {
 546                 $($macro)* {
 547                     $($other)*
 548
 549                     #query_stream
 550
 551                 }
 552             }
 553         }
 554         macro_rules! rustc_dep_node_append {
 555             ([$($macro:tt)*][$($other:tt)*]) => {
 556                 $($macro)*(
 557                     $($other)*
 558
 559                     #dep_node_def_stream
 560                 );
 561             }
 562         }
 563         macro_rules! rustc_dep_node_force {
 564             ([$dep_node:expr, $tcx:expr] $($other:tt)*) => {
 565                 match $dep_node.kind {
 566                     $($other)*
 567
 568                     #dep_node_force_stream
 569                 }
 570             }
 571         }
 572         macro_rules! rustc_cached_queries {
 573             ($($macro:tt)*) => {
 574                 $($macro)*(#cached_queries);
 575             }
 576         }
 577
 578         #query_description_stream
 579
 580         impl DepNode {
 581             /// Check whether the query invocation corresponding to the given
 582             /// DepNode is eligible for on-disk-caching. If so, this is method
 583             /// will execute the query corresponding to the given DepNode.
 584             /// Also, as a sanity check, it expects that the corresponding query
 585             /// invocation has been marked as green already.
 586             pub fn try_load_from_on_disk_cache(&self, tcx: TyCtxt<'_>) {
 587                 match self.kind {
 588                     #try_load_from_on_disk_cache_stream
 589                     _ => (),
 590                 }
 591             }
 592         }
 593     })
 594 }