src/librustc_macros/src/query.rs

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