Clean up parsing code and split out codegen for the QueryDescription impl

Author: Zoxc

src/librustc_macros/src/lib.rs

@@ -9,7 +9,6 @@ use proc_macro::TokenStream;
 
 mod hash_stable;
 mod query;
-mod tt;
 
 #[proc_macro]
 pub fn rustc_queries(input: TokenStream) -> TokenStream {

src/librustc_macros/src/query.rs

@@ -1,14 +1,21 @@
 use proc_macro::TokenStream;
 use proc_macro2::Span;
 use syn::{
-    Token, Ident, Type, Attribute, ReturnType, Expr, Block,
+    Token, Ident, Type, Attribute, ReturnType, Expr, Block, Error,
     braced, parenthesized, parse_macro_input,
 };
+use syn::spanned::Spanned;
 use syn::parse::{Result, Parse, ParseStream};
 use syn::punctuated::Punctuated;
+use syn;
 use quote::quote;
-use crate::tt::TS;
 
+#[allow(non_camel_case_types)]
+mod kw {
+    syn::custom_keyword!(query);
+}
+
+/// Ident or a wildcard `_`.
 struct IdentOrWild(Ident);
 
 impl Parse for IdentOrWild {
@@ -22,17 +29,27 @@ impl Parse for IdentOrWild {
     }
 }
 
-enum QueryAttribute {
+/// A modifier for a query
+enum QueryModifier {
+    /// The description of the query
     Desc(Option<Ident>, Punctuated<Expr, Token![,]>),
+
+    /// Cache the query to disk if the `Expr` returns true.
     Cache(Option<Ident>, Expr),
+
+    /// Custom code to load the query from disk.
     LoadCached(Ident, Ident, Block),
+
+    /// A cycle error for this query aborting the compilation with a fatal error.
     FatalCycle,
 }
 
-impl Parse for QueryAttribute {
+impl Parse for QueryModifier {
     fn parse(input: ParseStream<'_>) -> Result<Self> {
-        let attr: Ident = input.parse()?;
-        if attr == "desc" {
+        let modifier: Ident = input.parse()?;
+        if modifier == "desc" {
+            // Parse a description modifier like:
+            // `desc { |tcx| "foo {}", tcx.item_path(key) }`
             let attr_content;
             braced!(attr_content in input);
             let tcx = if attr_content.peek(Token![|]) {
@@ -44,11 +61,10 @@ impl Parse for QueryAttribute {
                 None
             };
             let desc = attr_content.parse_terminated(Expr::parse)?;
-            if !attr_content.is_empty() {
-                panic!("unexpected tokens in block");
-            };
-            Ok(QueryAttribute::Desc(tcx, desc))
-        } else if attr == "cache" {
+            Ok(QueryModifier::Desc(tcx, desc))
+        } else if modifier == "cache" {
+            // Parse a cache modifier like:
+            // `cache { |tcx| key.is_local() }`
             let attr_content;
             braced!(attr_content in input);
             let tcx = if attr_content.peek(Token![|]) {
@@ -60,68 +76,59 @@ impl Parse for QueryAttribute {
                 None
             };
             let expr = attr_content.parse()?;
-            if !attr_content.is_empty() {
-                panic!("unexpected tokens in block");
-            };
-            Ok(QueryAttribute::Cache(tcx, expr))
-        } else if attr == "load_cached" {
+            Ok(QueryModifier::Cache(tcx, expr))
+        } else if modifier == "load_cached" {
+            // Parse a load_cached modifier like:
+            // `load_cached(tcx, id) { tcx.queries.on_disk_cache.try_load_query_result(tcx, id) }`
             let args;
             parenthesized!(args in input);
             let tcx = args.parse()?;
             args.parse::<Token![,]>()?;
             let id = args.parse()?;
-            if !args.is_empty() {
-                panic!("unexpected tokens in arguments");
-            };
             let block = input.parse()?;
-            Ok(QueryAttribute::LoadCached(tcx, id, block))
-        } else if attr == "fatal_cycle" {
-            Ok(QueryAttribute::FatalCycle)
+            Ok(QueryModifier::LoadCached(tcx, id, block))
+        } else if modifier == "fatal_cycle" {
+            Ok(QueryModifier::FatalCycle)
         } else {
-            panic!("unknown query modifier {}", attr)
+            Err(Error::new(modifier.span(), "unknown query modifier"))
         }
     }
 }
 
+/// Ensures only doc comment attributes are used
+fn check_attributes(attrs: Vec<Attribute>) -> Result<()> {
+    for attr in attrs {
+        if !attr.path.is_ident("doc") {
+            return Err(Error::new(attr.span(), "attributes not supported on queries"));
+        }
+    }
+    Ok(())
+}
+
+/// A compiler query. `query ... { ... }`
 struct Query {
-    attrs: List<QueryAttribute>,
+    attrs: List<QueryModifier>,
     name: Ident,
     key: IdentOrWild,
     arg: Type,
     result: ReturnType,
 }
 
-fn check_attributes(attrs: Vec<Attribute>) {
-    for attr in attrs {
-        let path = attr.path;
-        let path = quote! { #path };
-        let path = TS(&path);
-
-        if path != TS(&quote! { doc }) {
-            panic!("attribute `{}` not supported on queries", path.0)
-        }
-    }
-}
-
 impl Parse for Query {
     fn parse(input: ParseStream<'_>) -> Result<Self> {
-        check_attributes(input.call(Attribute::parse_outer)?);
+        check_attributes(input.call(Attribute::parse_outer)?)?;
 
-        let query: Ident = input.parse()?;
-        if query != "query" {
-            panic!("expected `query`");
-        }
+        // Parse the query declaration. Like `query type_of(key: DefId) -> Ty<'tcx>`
+        input.parse::<kw::query>()?;
         let name: Ident = input.parse()?;
         let arg_content;
         parenthesized!(arg_content in input);
         let key = arg_content.parse()?;
         arg_content.parse::<Token![:]>()?;
         let arg = arg_content.parse()?;
-        if !arg_content.is_empty() {
-            panic!("expected only one query argument");
-        };
         let result = input.parse()?;
 
+        // Parse the query modifiers
         let content;
         braced!(content in input);
         let attrs = content.parse()?;
@@ -136,6 +143,7 @@ impl Parse for Query {
     }
 }
 
+/// A type used to greedily parse another type until the input is empty.
 struct List<T>(Vec<T>);
 
 impl<T: Parse> Parse for List<T> {
@@ -148,6 +156,7 @@ impl<T: Parse> Parse for List<T> {
     }
 }
 
+/// A named group containing queries.
 struct Group {
     name: Ident,
     queries: List<Query>,
@@ -165,6 +174,88 @@ impl Parse for Group {
     }
 }
 
+/// Add the impl of QueryDescription for the query to `impls` if one is requested
+fn add_query_description_impl(query: &Query, impls: &mut proc_macro2::TokenStream) {
+    let name = &query.name;
+    let arg = &query.arg;
+    let key = &query.key.0;
+
+    // Find custom code to load the query from disk
+    let load_cached = query.attrs.0.iter().find_map(|attr| match attr {
+        QueryModifier::LoadCached(tcx, id, block) => Some((tcx, id, block)),
+        _ => None,
+    });
+
+    // Find out if we should cache the query on disk
+    let cache = query.attrs.0.iter().find_map(|attr| match attr {
+        QueryModifier::Cache(tcx, expr) => Some((tcx, expr)),
+        _ => None,
+    }).map(|(tcx, expr)| {
+        let try_load_from_disk = if let Some((tcx, id, block)) = load_cached {
+            quote! {
+                #[inline]
+                fn try_load_from_disk(
+                    #tcx: TyCtxt<'_, 'tcx, 'tcx>,
+                    #id: SerializedDepNodeIndex
+                ) -> Option<Self::Value> {
+                    #block
+                }
+            }
+        } else {
+            quote! {
+                #[inline]
+                fn try_load_from_disk(
+                    tcx: TyCtxt<'_, 'tcx, 'tcx>,
+                    id: SerializedDepNodeIndex
+                ) -> Option<Self::Value> {
+                    tcx.queries.on_disk_cache.try_load_query_result(tcx, id)
+                }
+            }
+        };
+
+        let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
+        quote! {
+            #[inline]
+            fn cache_on_disk(#tcx: TyCtxt<'_, 'tcx, 'tcx>, #key: Self::Key) -> bool {
+                #expr
+            }
+
+            #try_load_from_disk
+        }
+    });
+
+    if cache.is_none() && load_cached.is_some() {
+        panic!("load_cached modifier on query `{}` without a cache modifier", name);
+    }
+
+    let desc = query.attrs.0.iter().find_map(|attr| match attr {
+        QueryModifier::Desc(tcx, desc) => Some((tcx, desc)),
+        _ => None,
+    }).map(|(tcx, desc)| {
+        let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
+        quote! {
+            fn describe(
+                #tcx: TyCtxt<'_, '_, '_>,
+                #key: #arg,
+            ) -> Cow<'static, str> {
+                format!(#desc).into()
+            }
+        }
+    });
+
+    if desc.is_some() || cache.is_some() {
+        let cache = cache.unwrap_or(quote! {});
+        let desc = desc.unwrap_or(quote! {});
+
+        impls.extend(quote! {
+            impl<'tcx> QueryDescription<'tcx> for queries::#name<'tcx> {
+                #desc
+                #cache
+            }
+        });
+    }
+}
+
 pub fn rustc_queries(input: TokenStream) -> TokenStream {
     let groups = parse_macro_input!(input as List<Group>);
 
@@ -178,99 +269,31 @@ pub fn rustc_queries(input: TokenStream) -> TokenStream {
         for query in &group.queries.0 {
             let name = &query.name;
             let arg = &query.arg;
-            let key = &query.key.0;
             let result_full = &query.result;
             let result = match query.result {
                 ReturnType::Default => quote! { -> () },
                 _ => quote! { #result_full },
             };
 
-            let load_cached = query.attrs.0.iter().find_map(|attr| match attr {
-                QueryAttribute::LoadCached(tcx, id, block) => Some((tcx, id, block)),
-                _ => None,
-            });
-
-            // Find out if we should cache the query on disk
-            let cache = query.attrs.0.iter().find_map(|attr| match attr {
-                QueryAttribute::Cache(tcx, expr) => Some((tcx, expr)),
-                _ => None,
-            }).map(|(tcx, expr)| {
-                let try_load_from_disk = if let Some((tcx, id, block)) = load_cached {
-                    quote! {
-                        #[inline]
-                        fn try_load_from_disk(
-                            #tcx: TyCtxt<'_, 'tcx, 'tcx>,
-                            #id: SerializedDepNodeIndex
-                        ) -> Option<Self::Value> {
-                            #block
-                        }
-                    }
-                } else {
-                    quote! {
-                        #[inline]
-                        fn try_load_from_disk(
-                            tcx: TyCtxt<'_, 'tcx, 'tcx>,
-                            id: SerializedDepNodeIndex
-                        ) -> Option<Self::Value> {
-                            tcx.queries.on_disk_cache.try_load_query_result(tcx, id)
-                        }
-                    }
-                };
-
-                let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
-                quote! {
-                    #[inline]
-                    fn cache_on_disk(#tcx: TyCtxt<'_, 'tcx, 'tcx>, #key: Self::Key) -> bool {
-                        #expr
-                    }
-
-                    #try_load_from_disk
-                }
-            });
-
-            if cache.is_none() && load_cached.is_some() {
-                panic!("load_cached modifier on query `{}` without a cache modifier", name);
-            }
-
+            // Look for a fatal_cycle modifier to pass on
             let fatal_cycle = query.attrs.0.iter().find_map(|attr| match attr {
-                QueryAttribute::FatalCycle => Some(()),
+                QueryModifier::FatalCycle => Some(()),
                 _ => None,
             }).map(|_| quote! { fatal_cycle }).unwrap_or(quote! {});
 
+            // Add the query to the group
             group_stream.extend(quote! {
                 [#fatal_cycle] fn #name: #name(#arg) #result,
             });
 
-            let desc = query.attrs.0.iter().find_map(|attr| match attr {
-                QueryAttribute::Desc(tcx, desc) => Some((tcx, desc)),
-                _ => None,
-            }).map(|(tcx, desc)| {
-                let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
-                quote! {
-                    fn describe(
-                        #tcx: TyCtxt<'_, '_, '_>,
-                        #key: #arg,
-                    ) -> Cow<'static, str> {
-                        format!(#desc).into()
-                    }
-                }
-            });
-
-            if desc.is_some() || cache.is_some() {
-                let cache = cache.unwrap_or(quote! {});
-                let desc = desc.unwrap_or(quote! {});
-
-                query_description_stream.extend(quote! {
-                    impl<'tcx> QueryDescription<'tcx> for queries::#name<'tcx> {
-                        #desc
-                        #cache
-                    }
-                });
-            }
+            add_query_description_impl(query, &mut query_description_stream);
 
+            // Create a dep node for the query
             dep_node_def_stream.extend(quote! {
                 [] #name(#arg),
             });
+
+            // Add a match arm to force the query given the dep node
             dep_node_force_stream.extend(quote! {
                 DepKind::#name => {
                     if let Some(key) = RecoverKey::recover($tcx, $dep_node) {