src/librustc_macros/src/query.rs

   1 use proc_macro::TokenStream;
   2 use syn::{
   3     Token, Ident, Type, Attribute, ReturnType, Expr, Block, Error,
   4     braced, parenthesized, parse_macro_input,
   5 };
   6 use syn::spanned::Spanned;
   7 use syn::parse::{Result, Parse, ParseStream};
   8 use syn::punctuated::Punctuated;
   9 use syn;
  10 use quote::quote;
  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     /// Cache the query to disk if the `Expr` returns true.
  37     Cache(Option<Ident>, Expr),
  38
  39     /// Custom code to load the query from disk.
  40     LoadCached(Ident, Ident, Block),
  41
  42     /// A cycle error for this query aborting the compilation with a fatal error.
  43     FatalCycle,
  44 }
  45
  46 impl Parse for QueryModifier {
  47     fn parse(input: ParseStream<'_>) -> Result<Self> {
  48         let modifier: Ident = input.parse()?;
  49         if modifier == "desc" {
  50             // Parse a description modifier like:
  51             // `desc { |tcx| "foo {}", tcx.item_path(key) }`
  52             let attr_content;
  53             braced!(attr_content in input);
  54             let tcx = if attr_content.peek(Token![|]) {
  55                 attr_content.parse::<Token![|]>()?;
  56                 let tcx = attr_content.parse()?;
  57                 attr_content.parse::<Token![|]>()?;
  58                 Some(tcx)
  59             } else {
  60                 None
  61             };
  62             let desc = attr_content.parse_terminated(Expr::parse)?;
  63             Ok(QueryModifier::Desc(tcx, desc))
  64         } else if modifier == "cache" {
  65             // Parse a cache modifier like:
  66             // `cache { |tcx| key.is_local() }`
  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 expr = attr_content.parse()?;
  78             Ok(QueryModifier::Cache(tcx, expr))
  79         } else if modifier == "load_cached" {
  80             // Parse a load_cached modifier like:
  81             // `load_cached(tcx, id) { tcx.queries.on_disk_cache.try_load_query_result(tcx, id) }`
  82             let args;
  83             parenthesized!(args in input);
  84             let tcx = args.parse()?;
  85             args.parse::<Token![,]>()?;
  86             let id = args.parse()?;
  87             let block = input.parse()?;
  88             Ok(QueryModifier::LoadCached(tcx, id, block))
  89         } else if modifier == "fatal_cycle" {
  90             Ok(QueryModifier::FatalCycle)
  91         } else {
  92             Err(Error::new(modifier.span(), "unknown query modifier"))
  93         }
  94     }
  95 }
  96
  97 /// Ensures only doc comment attributes are used
  98 fn check_attributes(attrs: Vec<Attribute>) -> Result<()> {
  99     for attr in attrs {
 100         if !attr.path.is_ident("doc") {
 101             return Err(Error::new(attr.span(), "attributes not supported on queries"));
 102         }
 103     }
 104     Ok(())
 105 }
 106
 107 /// A compiler query. `query ... { ... }`
 108 struct Query {
 109     modifiers: List<QueryModifier>,
 110     name: Ident,
 111     key: IdentOrWild,
 112     arg: Type,
 113     result: ReturnType,
 114 }
 115
 116 impl Parse for Query {
 117     fn parse(input: ParseStream<'_>) -> Result<Self> {
 118         check_attributes(input.call(Attribute::parse_outer)?)?;
 119
 120         // Parse the query declaration. Like `query type_of(key: DefId) -> Ty<'tcx>`
 121         input.parse::<kw::query>()?;
 122         let name: Ident = input.parse()?;
 123         let arg_content;
 124         parenthesized!(arg_content in input);
 125         let key = arg_content.parse()?;
 126         arg_content.parse::<Token![:]>()?;
 127         let arg = arg_content.parse()?;
 128         let result = input.parse()?;
 129
 130         // Parse the query modifiers
 131         let content;
 132         braced!(content in input);
 133         let modifiers = content.parse()?;
 134
 135         Ok(Query {
 136             modifiers,
 137             name,
 138             key,
 139             arg,
 140             result,
 141         })
 142     }
 143 }
 144
 145 /// A type used to greedily parse another type until the input is empty.
 146 struct List<T>(Vec<T>);
 147
 148 impl<T: Parse> Parse for List<T> {
 149     fn parse(input: ParseStream<'_>) -> Result<Self> {
 150         let mut list = Vec::new();
 151         while !input.is_empty() {
 152             list.push(input.parse()?);
 153         }
 154         Ok(List(list))
 155     }
 156 }
 157
 158 /// A named group containing queries.
 159 struct Group {
 160     name: Ident,
 161     queries: List<Query>,
 162 }
 163
 164 impl Parse for Group {
 165     fn parse(input: ParseStream<'_>) -> Result<Self> {
 166         let name: Ident = input.parse()?;
 167         let content;
 168         braced!(content in input);
 169         Ok(Group {
 170             name,
 171             queries: content.parse()?,
 172         })
 173     }
 174 }
 175
 176 struct QueryModifiers {
 177     /// The description of the query.
 178     desc: Option<(Option<Ident>, Punctuated<Expr, Token![,]>)>,
 179
 180     /// Cache the query to disk if the `Expr` returns true.
 181     cache: Option<(Option<Ident>, Expr)>,
 182
 183     /// Custom code to load the query from disk.
 184     load_cached: Option<(Ident, Ident, Block)>,
 185
 186     /// A cycle error for this query aborting the compilation with a fatal error.
 187     fatal_cycle: bool,
 188 }
 189
 190 /// Process query modifiers into a struct, erroring on duplicates
 191 fn process_modifiers(query: &mut Query) -> QueryModifiers {
 192     let mut load_cached = None;
 193     let mut cache = None;
 194     let mut desc = None;
 195     let mut fatal_cycle = false;
 196     for modifier in query.modifiers.0.drain(..) {
 197         match modifier {
 198             QueryModifier::LoadCached(tcx, id, block) => {
 199                 if load_cached.is_some() {
 200                     panic!("duplicate modifier `load_cached` for query `{}`", query.name);
 201                 }
 202                 load_cached = Some((tcx, id, block));
 203             }
 204             QueryModifier::Cache(tcx, expr) => {
 205                 if cache.is_some() {
 206                     panic!("duplicate modifier `cache` for query `{}`", query.name);
 207                 }
 208                 cache = Some((tcx, expr));
 209             }
 210             QueryModifier::Desc(tcx, list) => {
 211                 if desc.is_some() {
 212                     panic!("duplicate modifier `desc` for query `{}`", query.name);
 213                 }
 214                 desc = Some((tcx, list));
 215             }
 216             QueryModifier::FatalCycle => {
 217                 if fatal_cycle {
 218                     panic!("duplicate modifier `fatal_cycle` for query `{}`", query.name);
 219                 }
 220                 fatal_cycle = true;
 221             }
 222         }
 223     }
 224     QueryModifiers {
 225         load_cached,
 226         cache,
 227         desc,
 228         fatal_cycle,
 229     }
 230 }
 231
 232 /// Add the impl of QueryDescription for the query to `impls` if one is requested
 233 fn add_query_description_impl(
 234     query: &Query,
 235     modifiers: QueryModifiers,
 236     impls: &mut proc_macro2::TokenStream
 237 ) {
 238     let name = &query.name;
 239     let arg = &query.arg;
 240     let key = &query.key.0;
 241
 242     // Find out if we should cache the query on disk
 243     let cache = modifiers.cache.as_ref().map(|(tcx, expr)| {
 244         let try_load_from_disk = if let Some((tcx, id, block)) = modifiers.load_cached.as_ref() {
 245             // Use custom code to load the query from disk
 246             quote! {
 247                 #[inline]
 248                 fn try_load_from_disk(
 249                     #tcx: TyCtxt<'_, 'tcx, 'tcx>,
 250                     #id: SerializedDepNodeIndex
 251                 ) -> Option<Self::Value> {
 252                     #block
 253                 }
 254             }
 255         } else {
 256             // Use the default code to load the query from disk
 257             quote! {
 258                 #[inline]
 259                 fn try_load_from_disk(
 260                     tcx: TyCtxt<'_, 'tcx, 'tcx>,
 261                     id: SerializedDepNodeIndex
 262                 ) -> Option<Self::Value> {
 263                     tcx.queries.on_disk_cache.try_load_query_result(tcx, id)
 264                 }
 265             }
 266         };
 267
 268         let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
 269         quote! {
 270             #[inline]
 271             fn cache_on_disk(#tcx: TyCtxt<'_, 'tcx, 'tcx>, #key: Self::Key) -> bool {
 272                 #expr
 273             }
 274
 275             #try_load_from_disk
 276         }
 277     });
 278
 279     if cache.is_none() && modifiers.load_cached.is_some() {
 280         panic!("load_cached modifier on query `{}` without a cache modifier", name);
 281     }
 282
 283     let desc = modifiers.desc.as_ref().map(|(tcx, desc)| {
 284         let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
 285         quote! {
 286             fn describe(
 287                 #tcx: TyCtxt<'_, '_, '_>,
 288                 #key: #arg,
 289             ) -> Cow<'static, str> {
 290                 format!(#desc).into()
 291             }
 292         }
 293     });
 294
 295     if desc.is_some() || cache.is_some() {
 296         let cache = cache.unwrap_or(quote! {});
 297         let desc = desc.unwrap_or(quote! {});
 298
 299         impls.extend(quote! {
 300             impl<'tcx> QueryDescription<'tcx> for queries::#name<'tcx> {
 301                 #desc
 302                 #cache
 303             }
 304         });
 305     }
 306 }
 307
 308 pub fn rustc_queries(input: TokenStream) -> TokenStream {
 309     let groups = parse_macro_input!(input as List<Group>);
 310
 311     let mut query_stream = quote! {};
 312     let mut query_description_stream = quote! {};
 313     let mut dep_node_def_stream = quote! {};
 314     let mut dep_node_force_stream = quote! {};
 315
 316     for group in groups.0 {
 317         let mut group_stream = quote! {};
 318         for mut query in group.queries.0 {
 319             let modifiers = process_modifiers(&mut query);
 320             let name = &query.name;
 321             let arg = &query.arg;
 322             let result_full = &query.result;
 323             let result = match query.result {
 324                 ReturnType::Default => quote! { -> () },
 325                 _ => quote! { #result_full },
 326             };
 327
 328             // Pass on the fatal_cycle modifier
 329             let fatal_cycle = if modifiers.fatal_cycle {
 330                 quote! { fatal_cycle }
 331             } else {
 332                 quote! {}
 333             };
 334
 335             // Add the query to the group
 336             group_stream.extend(quote! {
 337                 [#fatal_cycle] fn #name: #name(#arg) #result,
 338             });
 339
 340             add_query_description_impl(&query, modifiers, &mut query_description_stream);
 341
 342             // Create a dep node for the query
 343             dep_node_def_stream.extend(quote! {
 344                 [] #name(#arg),
 345             });
 346
 347             // Add a match arm to force the query given the dep node
 348             dep_node_force_stream.extend(quote! {
 349                 DepKind::#name => {
 350                     if let Some(key) = RecoverKey::recover($tcx, $dep_node) {
 351                         force_ex!($tcx, #name, key);
 352                     } else {
 353                         return false;
 354                     }
 355                 }
 356             });
 357         }
 358         let name = &group.name;
 359         query_stream.extend(quote! {
 360             #name { #group_stream },
 361         });
 362     }
 363     TokenStream::from(quote! {
 364         macro_rules! rustc_query_append {
 365             ([$($macro:tt)*][$($other:tt)*]) => {
 366                 $($macro)* {
 367                     $($other)*
 368
 369                     #query_stream
 370
 371                 }
 372             }
 373         }
 374         macro_rules! rustc_dep_node_append {
 375             ([$($macro:tt)*][$($other:tt)*]) => {
 376                 $($macro)*(
 377                     $($other)*
 378
 379                     #dep_node_def_stream
 380                 );
 381             }
 382         }
 383         macro_rules! rustc_dep_node_force {
 384             ([$dep_node:expr, $tcx:expr] $($other:tt)*) => {
 385                 match $dep_node.kind {
 386                     $($other)*
 387
 388                     #dep_node_force_stream
 389                 }
 390             }
 391         }
 392         #query_description_stream
 393     })
 394 }