Refactor query modifier parsing

compiler/rustc_macros/src/query.rs

@@ -1,139 +1,17 @@
 use proc_macro::TokenStream;
-use proc_macro2::{Delimiter, TokenTree};
 use quote::{quote, quote_spanned};
 use syn::parse::{Parse, ParseStream, Result};
 use syn::punctuated::Punctuated;
 use syn::spanned::Spanned;
 use syn::{
-    braced, parenthesized, parse_macro_input, parse_quote, AttrStyle, Attribute, Block, Error,
-    Expr, Ident, ReturnType, Token, Type,
+    braced, parenthesized, parse_macro_input, parse_quote, token, AttrStyle, Attribute, Block,
+    Error, Expr, Ident, Pat, ReturnType, Token, Type,
 };
 
 mod kw {
     syn::custom_keyword!(query);
 }
 
-/// Ident or a wildcard `_`.
-struct IdentOrWild(Ident);
-
-impl Parse for IdentOrWild {
-    fn parse(input: ParseStream<'_>) -> Result<Self> {
-        Ok(if input.peek(Token![_]) {
-            let underscore = input.parse::<Token![_]>()?;
-            IdentOrWild(Ident::new("_", underscore.span()))
-        } else {
-            IdentOrWild(input.parse()?)
-        })
-    }
-}
-
-/// A modifier for a query
-enum QueryModifier {
-    /// The description of the query.
-    Desc(Option<Ident>, Punctuated<Expr, Token![,]>),
-
-    /// Use this type for the in-memory cache.
-    Storage(Type),
-
-    /// Cache the query to disk if the `Expr` returns true.
-    Cache(Option<IdentOrWild>, Block),
-
-    /// 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(Ident),
-
-    /// A cycle error results in a delay_bug call
-    CycleDelayBug(Ident),
-
-    /// Don't hash the result, instead just mark a query red if it runs
-    NoHash(Ident),
-
-    /// Generate a dep node based on the dependencies of the query
-    Anon(Ident),
-
-    /// Always evaluate the query, ignoring its dependencies
-    EvalAlways(Ident),
-
-    /// Use a separate query provider for local and extern crates
-    SeparateProvideExtern(Ident),
-
-    /// Always remap the ParamEnv's constness before hashing and passing to the query provider
-    RemapEnvConstness(Ident),
-}
-
-impl Parse for QueryModifier {
-    fn parse(input: ParseStream<'_>) -> Result<Self> {
-        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![|]) {
-                attr_content.parse::<Token![|]>()?;
-                let tcx = attr_content.parse()?;
-                attr_content.parse::<Token![|]>()?;
-                Some(tcx)
-            } else {
-                None
-            };
-            let desc = attr_content.parse_terminated(Expr::parse)?;
-            Ok(QueryModifier::Desc(tcx, desc))
-        } else if modifier == "cache_on_disk_if" {
-            // Parse a cache modifier like:
-            // `cache(tcx, value) { |tcx| key.is_local() }`
-            let has_args = if let TokenTree::Group(group) = input.fork().parse()? {
-                group.delimiter() == Delimiter::Parenthesis
-            } else {
-                false
-            };
-            let args = if has_args {
-                let args;
-                parenthesized!(args in input);
-                let tcx = args.parse()?;
-                Some(tcx)
-            } else {
-                None
-            };
-            let block = input.parse()?;
-            Ok(QueryModifier::Cache(args, block))
-        } else if modifier == "load_cached" {
-            // Parse a load_cached modifier like:
-            // `load_cached(tcx, id) { tcx.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()?;
-            let block = input.parse()?;
-            Ok(QueryModifier::LoadCached(tcx, id, block))
-        } else if modifier == "storage" {
-            let args;
-            parenthesized!(args in input);
-            let ty = args.parse()?;
-            Ok(QueryModifier::Storage(ty))
-        } else if modifier == "fatal_cycle" {
-            Ok(QueryModifier::FatalCycle(modifier))
-        } else if modifier == "cycle_delay_bug" {
-            Ok(QueryModifier::CycleDelayBug(modifier))
-        } else if modifier == "no_hash" {
-            Ok(QueryModifier::NoHash(modifier))
-        } else if modifier == "anon" {
-            Ok(QueryModifier::Anon(modifier))
-        } else if modifier == "eval_always" {
-            Ok(QueryModifier::EvalAlways(modifier))
-        } else if modifier == "separate_provide_extern" {
-            Ok(QueryModifier::SeparateProvideExtern(modifier))
-        } else if modifier == "remap_env_constness" {
-            Ok(QueryModifier::RemapEnvConstness(modifier))
-        } else {
-            Err(Error::new(modifier.span(), "unknown query modifier"))
-        }
-    }
-}
-
 /// Ensures only doc comment attributes are used
 fn check_attributes(attrs: Vec<Attribute>) -> Result<Vec<Attribute>> {
     let inner = |attr: Attribute| {
@@ -154,16 +32,16 @@ fn check_attributes(attrs: Vec<Attribute>) -> Result<Vec<Attribute>> {
 /// A compiler query. `query ... { ... }`
 struct Query {
     doc_comments: Vec<Attribute>,
-    modifiers: List<QueryModifier>,
+    modifiers: QueryModifiers,
     name: Ident,
-    key: IdentOrWild,
+    key: Pat,
     arg: Type,
     result: ReturnType,
 }
 
 impl Parse for Query {
     fn parse(input: ParseStream<'_>) -> Result<Self> {
-        let doc_comments = check_attributes(input.call(Attribute::parse_outer)?)?;
+        let mut doc_comments = check_attributes(input.call(Attribute::parse_outer)?)?;
 
         // Parse the query declaration. Like `query type_of(key: DefId) -> Ty<'tcx>`
         input.parse::<kw::query>()?;
@@ -178,7 +56,13 @@ impl Parse for Query {
         // Parse the query modifiers
         let content;
         braced!(content in input);
-        let modifiers = content.parse()?;
+        let modifiers = parse_query_modifiers(&content)?;
+
+        // If there are no doc-comments, give at least some idea of what
+        // it does by showing the query description.
+        if doc_comments.is_empty() {
+            doc_comments.push(doc_comment_from_desc(&modifiers.desc.1)?);
+        }
 
         Ok(Query { doc_comments, modifiers, name, key, arg, result })
     }
@@ -205,7 +89,7 @@ struct QueryModifiers {
     storage: Option<Type>,
 
     /// Cache the query to disk if the `Block` returns true.
-    cache: Option<(Option<IdentOrWild>, Block)>,
+    cache: Option<(Option<Pat>, Block)>,
 
     /// Custom code to load the query from disk.
     load_cached: Option<(Ident, Ident, Block)>,
@@ -232,8 +116,7 @@ struct QueryModifiers {
     remap_env_constness: Option<Ident>,
 }
 
-/// Process query modifiers into a struct, erroring on duplicates
-fn process_modifiers(query: &mut Query) -> QueryModifiers {
+fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
     let mut load_cached = None;
     let mut storage = None;
     let mut cache = None;
@@ -245,117 +128,84 @@ fn process_modifiers(query: &mut Query) -> QueryModifiers {
     let mut eval_always = None;
     let mut separate_provide_extern = None;
     let mut remap_env_constness = None;
-    for modifier in query.modifiers.0.drain(..) {
-        match modifier {
-            QueryModifier::LoadCached(tcx, id, block) => {
-                if load_cached.is_some() {
-                    panic!("duplicate modifier `load_cached` for query `{}`", query.name);
+
+    while !input.is_empty() {
+        let modifier: Ident = input.parse()?;
+
+        macro_rules! try_insert {
+            ($name:ident = $expr:expr) => {
+                if $name.is_some() {
+                    return Err(Error::new(modifier.span(), "duplicate modifier"));
                 }
-                load_cached = Some((tcx, id, block));
-            }
-            QueryModifier::Storage(ty) => {
-                if storage.is_some() {
-                    panic!("duplicate modifier `storage` for query `{}`", query.name);
-                }
-                storage = Some(ty);
-            }
-            QueryModifier::Cache(args, expr) => {
-                if cache.is_some() {
-                    panic!("duplicate modifier `cache` for query `{}`", query.name);
-                }
-                cache = Some((args, expr));
-            }
-            QueryModifier::Desc(tcx, list) => {
-                if desc.is_some() {
-                    panic!("duplicate modifier `desc` for query `{}`", query.name);
-                }
-                // If there are no doc-comments, give at least some idea of what
-                // it does by showing the query description.
-                if query.doc_comments.is_empty() {
-                    use ::syn::*;
-                    let mut list = list.iter();
-                    let format_str: String = match list.next() {
-                        Some(&Expr::Lit(ExprLit { lit: Lit::Str(ref lit_str), .. })) => {
-                            lit_str.value().replace("`{}`", "{}") // We add them later anyways for consistency
-                        }
-                        _ => panic!("Expected a string literal"),
-                    };
-                    let mut fmt_fragments = format_str.split("{}");
-                    let mut doc_string = fmt_fragments.next().unwrap().to_string();
-                    list.map(::quote::ToTokens::to_token_stream).zip(fmt_fragments).for_each(
-                        |(tts, next_fmt_fragment)| {
-                            use ::core::fmt::Write;
-                            write!(
-                                &mut doc_string,
-                                " `{}` {}",
-                                tts.to_string().replace(" . ", "."),
-                                next_fmt_fragment,
-                            )
-                            .unwrap();
-                        },
-                    );
-                    let doc_string = format!(
-                        "[query description - consider adding a doc-comment!] {}",
-                        doc_string
-                    );
-                    let comment = parse_quote! {
-                        #[doc = #doc_string]
-                    };
-                    query.doc_comments.push(comment);
-                }
-                desc = Some((tcx, list));
-            }
-            QueryModifier::FatalCycle(ident) => {
-                if fatal_cycle.is_some() {
-                    panic!("duplicate modifier `fatal_cycle` for query `{}`", query.name);
-                }
-                fatal_cycle = Some(ident);
-            }
-            QueryModifier::CycleDelayBug(ident) => {
-                if cycle_delay_bug.is_some() {
-                    panic!("duplicate modifier `cycle_delay_bug` for query `{}`", query.name);
-                }
-                cycle_delay_bug = Some(ident);
-            }
-            QueryModifier::NoHash(ident) => {
-                if no_hash.is_some() {
-                    panic!("duplicate modifier `no_hash` for query `{}`", query.name);
-                }
-                no_hash = Some(ident);
-            }
-            QueryModifier::Anon(ident) => {
-                if anon.is_some() {
-                    panic!("duplicate modifier `anon` for query `{}`", query.name);
-                }
-                anon = Some(ident);
-            }
-            QueryModifier::EvalAlways(ident) => {
-                if eval_always.is_some() {
-                    panic!("duplicate modifier `eval_always` for query `{}`", query.name);
-                }
-                eval_always = Some(ident);
-            }
-            QueryModifier::SeparateProvideExtern(ident) => {
-                if separate_provide_extern.is_some() {
-                    panic!(
-                        "duplicate modifier `separate_provide_extern` for query `{}`",
-                        query.name
-                    );
-                }
-                separate_provide_extern = Some(ident);
-            }
-            QueryModifier::RemapEnvConstness(ident) => {
-                if remap_env_constness.is_some() {
-                    panic!("duplicate modifier `remap_env_constness` for query `{}`", query.name);
-                }
-                remap_env_constness = Some(ident)
-            }
+                $name = Some($expr);
+            };
+        }
+
+        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![|]) {
+                attr_content.parse::<Token![|]>()?;
+                let tcx = attr_content.parse()?;
+                attr_content.parse::<Token![|]>()?;
+                Some(tcx)
+            } else {
+                None
+            };
+            let list = attr_content.parse_terminated(Expr::parse)?;
+            try_insert!(desc = (tcx, list));
+        } else if modifier == "cache_on_disk_if" {
+            // Parse a cache modifier like:
+            // `cache(tcx) { |tcx| key.is_local() }`
+            let args = if input.peek(token::Paren) {
+                let args;
+                parenthesized!(args in input);
+                let tcx = args.parse()?;
+                Some(tcx)
+            } else {
+                None
+            };
+            let block = input.parse()?;
+            try_insert!(cache = (args, block));
+        } else if modifier == "load_cached" {
+            // Parse a load_cached modifier like:
+            // `load_cached(tcx, id) { tcx.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()?;
+            let block = input.parse()?;
+            try_insert!(load_cached = (tcx, id, block));
+        } else if modifier == "storage" {
+            let args;
+            parenthesized!(args in input);
+            let ty = args.parse()?;
+            try_insert!(storage = ty);
+        } else if modifier == "fatal_cycle" {
+            try_insert!(fatal_cycle = modifier);
+        } else if modifier == "cycle_delay_bug" {
+            try_insert!(cycle_delay_bug = modifier);
+        } else if modifier == "no_hash" {
+            try_insert!(no_hash = modifier);
+        } else if modifier == "anon" {
+            try_insert!(anon = modifier);
+        } else if modifier == "eval_always" {
+            try_insert!(eval_always = modifier);
+        } else if modifier == "separate_provide_extern" {
+            try_insert!(separate_provide_extern = modifier);
+        } else if modifier == "remap_env_constness" {
+            try_insert!(remap_env_constness = modifier);
+        } else {
+            return Err(Error::new(modifier.span(), "unknown query modifier"));
         }
     }
-    let desc = desc.unwrap_or_else(|| {
-        panic!("no description provided for query `{}`", query.name);
-    });
-    QueryModifiers {
+    let Some(desc) = desc else {
+        return Err(input.error("no description provided"));
+    };
+    Ok(QueryModifiers {
         load_cached,
         storage,
         cache,
@@ -367,17 +217,41 @@ fn process_modifiers(query: &mut Query) -> QueryModifiers {
         eval_always,
         separate_provide_extern,
         remap_env_constness,
-    }
+    })
+}
+
+fn doc_comment_from_desc(list: &Punctuated<Expr, token::Comma>) -> Result<Attribute> {
+    use ::syn::*;
+    let mut iter = list.iter();
+    let format_str: String = match iter.next() {
+        Some(&Expr::Lit(ExprLit { lit: Lit::Str(ref lit_str), .. })) => {
+            lit_str.value().replace("`{}`", "{}") // We add them later anyways for consistency
+        }
+        _ => return Err(Error::new(list.span(), "Expected a string literal")),
+    };
+    let mut fmt_fragments = format_str.split("{}");
+    let mut doc_string = fmt_fragments.next().unwrap().to_string();
+    iter.map(::quote::ToTokens::to_token_stream).zip(fmt_fragments).for_each(
+        |(tts, next_fmt_fragment)| {
+            use ::core::fmt::Write;
+            write!(
+                &mut doc_string,
+                " `{}` {}",
+                tts.to_string().replace(" . ", "."),
+                next_fmt_fragment,
+            )
+            .unwrap();
+        },
+    );
+    let doc_string = format!("[query description - consider adding a doc-comment!] {}", doc_string);
+    Ok(parse_quote! { #[doc = #doc_string] })
 }
 
 /// Add the impl of QueryDescription for the query to `impls` if one is requested
-fn add_query_description_impl(
-    query: &Query,
-    modifiers: QueryModifiers,
-    impls: &mut proc_macro2::TokenStream,
-) {
+fn add_query_description_impl(query: &Query, impls: &mut proc_macro2::TokenStream) {
     let name = &query.name;
-    let key = &query.key.0;
+    let key = &query.key;
+    let modifiers = &query.modifiers;
 
     // Find out if we should cache the query on disk
     let cache = if let Some((args, expr)) = modifiers.cache.as_ref() {
@@ -395,13 +269,7 @@ fn add_query_description_impl(
             }
         };
 
-        let tcx = args
-            .as_ref()
-            .map(|t| {
-                let t = &t.0;
-                quote! { #t }
-            })
-            .unwrap_or_else(|| quote! { _ });
+        let tcx = args.as_ref().map(|t| quote! { #t }).unwrap_or_else(|| quote! { _ });
         // expr is a `Block`, meaning that `{ #expr }` gets expanded
         // to `{ { stmts... } }`, which triggers the `unused_braces` lint.
         quote! {
@@ -427,7 +295,7 @@ fn add_query_description_impl(
         }
     };
 
-    let (tcx, desc) = modifiers.desc;
+    let (tcx, desc) = &modifiers.desc;
     let tcx = tcx.as_ref().map_or_else(|| quote! { _ }, |t| quote! { #t });
 
     let desc = quote! {
@@ -456,10 +324,8 @@ pub fn rustc_queries(input: TokenStream) -> TokenStream {
     let mut dep_node_def_stream = quote! {};
     let mut cached_queries = quote! {};
 
-    for mut query in queries.0 {
-        let modifiers = process_modifiers(&mut query);
-        let name = &query.name;
-        let arg = &query.arg;
+    for query in queries.0 {
+        let Query { name, arg, modifiers, .. } = &query;
         let result_full = &query.result;
         let result = match query.result {
             ReturnType::Default => quote! { -> () },
@@ -528,7 +394,7 @@ pub fn rustc_queries(input: TokenStream) -> TokenStream {
             [#attribute_stream] #name(#arg),
         });
 
-        add_query_description_impl(&query, modifiers, &mut query_description_stream);
+        add_query_description_impl(&query, &mut query_description_stream);
     }
 
     TokenStream::from(quote! {
@@ -539,7 +405,6 @@ pub fn rustc_queries(input: TokenStream) -> TokenStream {
                     $($other)*
 
                     #query_stream
-
                 }
             }
         }