macros: allow `Vec` fields in diagnostic derive

Diagnostics can have multiple primary spans, or have subdiagnostics
repeated at multiple locations, so support `Vec<..>` fields in the
diagnostic derive which become loops in the generated code.

Signed-off-by: David Wood <david.wood@huawei.com>

3 files changed, 85 insertions, 39 deletions

diff --git a/compiler/rustc_macros/src/diagnostics/diagnostic.rs b/compiler/rustc_macros/src/diagnostics/diagnostic.rs
index f49166433fa..83fc7bcde8a 100644
--- a/compiler/rustc_macros/src/diagnostics/diagnostic.rs
+++ b/compiler/rustc_macros/src/diagnostics/diagnostic.rs
@@ -5,8 +5,8 @@ use crate::diagnostics::error::{
     SessionDiagnosticDeriveError,
 };
 use crate::diagnostics::utils::{
-    option_inner_ty, report_error_if_not_applied_to_span, type_matches_path, Applicability,
-    FieldInfo, HasFieldMap, SetOnce,
+    report_error_if_not_applied_to_span, type_matches_path, Applicability, FieldInfo, FieldInnerTy,
+    HasFieldMap, SetOnce,
 };
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
@@ -353,35 +353,40 @@ impl SessionDiagnosticDeriveBuilder {
         info: FieldInfo<'_>,
     ) -> Result<TokenStream, SessionDiagnosticDeriveError> {
         let field_binding = &info.binding.binding;
-        let option_ty = option_inner_ty(&info.ty);
-        let generated_code = self.generate_non_option_field_code(
+
+        let inner_ty = FieldInnerTy::from_type(&info.ty);
+        let name = attr.path.segments.last().unwrap().ident.to_string();
+        let (binding, needs_destructure) = match (name.as_str(), &inner_ty) {
+            // `primary_span` can accept a `Vec<Span>` so don't destructure that.
+            ("primary_span", FieldInnerTy::Vec(_)) => (quote! { #field_binding.clone() }, false),
+            _ => (quote! { *#field_binding }, true),
+        };
+
+        let generated_code = self.generate_inner_field_code(
             attr,
             FieldInfo {
                 vis: info.vis,
                 binding: info.binding,
-                ty: option_ty.unwrap_or(&info.ty),
+                ty: inner_ty.inner_type().unwrap_or(&info.ty),
                 span: info.span,
             },
+            binding,
         )?;
 
-        if option_ty.is_none() {
-            Ok(quote! { #generated_code })
+        if needs_destructure {
+            Ok(inner_ty.with(field_binding, generated_code))
         } else {
-            Ok(quote! {
-                if let Some(#field_binding) = #field_binding {
-                    #generated_code
-                }
-            })
+            Ok(generated_code)
         }
     }
 
-    fn generate_non_option_field_code(
+    fn generate_inner_field_code(
         &mut self,
         attr: &Attribute,
         info: FieldInfo<'_>,
+        binding: TokenStream,
     ) -> Result<TokenStream, SessionDiagnosticDeriveError> {
         let diag = &self.diag;
-        let field_binding = &info.binding.binding;
 
         let name = attr.path.segments.last().unwrap().ident.to_string();
         let name = name.as_str();
@@ -397,14 +402,14 @@ impl SessionDiagnosticDeriveBuilder {
                 "primary_span" => {
                     report_error_if_not_applied_to_span(attr, &info)?;
                     Ok(quote! {
-                        #diag.set_span(*#field_binding);
+                        #diag.set_span(#binding);
                     })
                 }
                 "label" | "note" | "help" => {
                     report_error_if_not_applied_to_span(attr, &info)?;
-                    Ok(self.add_subdiagnostic(field_binding, name, name))
+                    Ok(self.add_subdiagnostic(binding, name, name))
                 }
-                "subdiagnostic" => Ok(quote! { #diag.subdiagnostic(*#field_binding); }),
+                "subdiagnostic" => Ok(quote! { #diag.subdiagnostic(#binding); }),
                 _ => throw_invalid_attr!(attr, &meta, |diag| {
                     diag
                         .help("only `skip_arg`, `primary_span`, `label`, `note`, `help` and `subdiagnostic` are valid field attributes")
@@ -413,7 +418,7 @@ impl SessionDiagnosticDeriveBuilder {
             Meta::NameValue(MetaNameValue { lit: syn::Lit::Str(ref s), .. }) => match name {
                 "label" | "note" | "help" => {
                     report_error_if_not_applied_to_span(attr, &info)?;
-                    Ok(self.add_subdiagnostic(field_binding, name, &s.value()))
+                    Ok(self.add_subdiagnostic(binding, name, &s.value()))
                 }
                 _ => throw_invalid_attr!(attr, &meta, |diag| {
                     diag.help("only `label`, `note` and `help` are valid field attributes")
@@ -509,7 +514,7 @@ impl SessionDiagnosticDeriveBuilder {
     /// `fluent_attr_identifier`.
     fn add_subdiagnostic(
         &self,
-        field_binding: &proc_macro2::Ident,
+        field_binding: TokenStream,
         kind: &str,
         fluent_attr_identifier: &str,
     ) -> TokenStream {
@@ -520,7 +525,7 @@ impl SessionDiagnosticDeriveBuilder {
         let fn_name = format_ident!("span_{}", kind);
         quote! {
             #diag.#fn_name(
-                *#field_binding,
+                #field_binding,
                 rustc_errors::DiagnosticMessage::fluent_attr(#slug, #fluent_attr_identifier)
             );
         }
diff --git a/compiler/rustc_macros/src/diagnostics/subdiagnostic.rs b/compiler/rustc_macros/src/diagnostics/subdiagnostic.rs
index 961b42f424f..65b1328682f 100644
--- a/compiler/rustc_macros/src/diagnostics/subdiagnostic.rs
+++ b/compiler/rustc_macros/src/diagnostics/subdiagnostic.rs
@@ -5,8 +5,8 @@ use crate::diagnostics::error::{
     SessionDiagnosticDeriveError,
 };
 use crate::diagnostics::utils::{
-    option_inner_ty, report_error_if_not_applied_to_applicability,
-    report_error_if_not_applied_to_span, Applicability, FieldInfo, HasFieldMap, SetOnce,
+    report_error_if_not_applied_to_applicability, report_error_if_not_applied_to_span,
+    Applicability, FieldInfo, FieldInnerTy, HasFieldMap, SetOnce,
 };
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
@@ -301,11 +301,11 @@ impl<'a> SessionSubdiagnosticDeriveBuilder<'a> {
     ) -> Result<TokenStream, SessionDiagnosticDeriveError> {
         let ast = binding.ast();
 
-        let option_ty = option_inner_ty(&ast.ty);
+        let inner_ty = FieldInnerTy::from_type(&ast.ty);
         let info = FieldInfo {
             vis: &ast.vis,
             binding: binding,
-            ty: option_ty.unwrap_or(&ast.ty),
+            ty: inner_ty.inner_type().unwrap_or(&ast.ty),
             span: &ast.span(),
         };
 
@@ -353,15 +353,7 @@ impl<'a> SessionSubdiagnosticDeriveBuilder<'a> {
             );
         };
 
-        if option_ty.is_none() {
-            Ok(quote! { #generated })
-        } else {
-            Ok(quote! {
-                if let Some(#binding) = #binding {
-                    #generated
-                }
-            })
-        }
+        Ok(inner_ty.with(binding, generated))
     }
 
     fn into_tokens(&mut self) -> Result<TokenStream, SessionDiagnosticDeriveError> {
diff --git a/compiler/rustc_macros/src/diagnostics/utils.rs b/compiler/rustc_macros/src/diagnostics/utils.rs
index 1f36af0a20b..aba861fc6aa 100644
--- a/compiler/rustc_macros/src/diagnostics/utils.rs
+++ b/compiler/rustc_macros/src/diagnostics/utils.rs
@@ -1,7 +1,7 @@
 use crate::diagnostics::error::{span_err, throw_span_err, SessionDiagnosticDeriveError};
 use proc_macro::Span;
 use proc_macro2::TokenStream;
-use quote::{format_ident, quote};
+use quote::{format_ident, quote, ToTokens};
 use std::collections::BTreeSet;
 use std::str::FromStr;
 use syn::{spanned::Spanned, Attribute, Meta, Type, Visibility};
@@ -76,22 +76,71 @@ pub(crate) fn report_error_if_not_applied_to_span(
     report_error_if_not_applied_to_ty(attr, info, &["rustc_span", "Span"], "Span")
 }
 
-/// If `ty` is an Option, returns `Some(inner type)`, otherwise returns `None`.
-pub(crate) fn option_inner_ty(ty: &Type) -> Option<&Type> {
-    if type_matches_path(ty, &["std", "option", "Option"]) {
+/// Inner type of a field and type of wrapper.
+pub(crate) enum FieldInnerTy<'ty> {
+    /// Field is wrapped in a `Option<$inner>`.
+    Option(&'ty Type),
+    /// Field is wrapped in a `Vec<$inner>`.
+    Vec(&'ty Type),
+    /// Field isn't wrapped in an outer type.
+    None,
+}
+
+impl<'ty> FieldInnerTy<'ty> {
+    /// Returns inner type for a field, if there is one.
+    ///
+    /// - If `ty` is an `Option`, returns `FieldInnerTy::Option { inner: (inner type) }`.
+    /// - If `ty` is a `Vec`, returns `FieldInnerTy::Vec { inner: (inner type) }`.
+    /// - Otherwise returns `None`.
+    pub(crate) fn from_type(ty: &'ty Type) -> Self {
+        let variant: &dyn Fn(&'ty Type) -> FieldInnerTy<'ty> =
+            if type_matches_path(ty, &["std", "option", "Option"]) {
+                &FieldInnerTy::Option
+            } else if type_matches_path(ty, &["std", "vec", "Vec"]) {
+                &FieldInnerTy::Vec
+            } else {
+                return FieldInnerTy::None;
+            };
+
         if let Type::Path(ty_path) = ty {
             let path = &ty_path.path;
             let ty = path.segments.iter().last().unwrap();
             if let syn::PathArguments::AngleBracketed(bracketed) = &ty.arguments {
                 if bracketed.args.len() == 1 {
                     if let syn::GenericArgument::Type(ty) = &bracketed.args[0] {
-                        return Some(ty);
+                        return variant(ty);
                     }
                 }
             }
         }
+
+        unreachable!();
+    }
+
+    /// Returns `Option` containing inner type if there is one.
+    pub(crate) fn inner_type(&self) -> Option<&'ty Type> {
+        match self {
+            FieldInnerTy::Option(inner) | FieldInnerTy::Vec(inner) => Some(inner),
+            FieldInnerTy::None => None,
+        }
+    }
+
+    /// Surrounds `inner` with destructured wrapper type, exposing inner type as `binding`.
+    pub(crate) fn with(&self, binding: impl ToTokens, inner: impl ToTokens) -> TokenStream {
+        match self {
+            FieldInnerTy::Option(..) => quote! {
+                if let Some(#binding) = #binding {
+                    #inner
+                }
+            },
+            FieldInnerTy::Vec(..) => quote! {
+                for #binding in #binding {
+                    #inner
+                }
+            },
+            FieldInnerTy::None => quote! { #inner },
+        }
     }
-    None
 }
 
 /// Field information passed to the builder. Deliberately omits attrs to discourage the