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