Improve diagnostic context for expression position macros and for parsing

internal/src/diagnostics.rs

@@ -1,44 +1,149 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 
+use std::cell::RefCell;
 use std::fmt::Display;
+use std::marker::PhantomData;
 
-use proc_macro2::TokenStream;
-use quote::quote_spanned;
+use proc_macro2::{Span, TokenStream};
+use quote::{quote, quote_spanned};
 use syn::{spanned::Spanned, Error};
 
-pub(crate) struct DiagCtxt(TokenStream);
+pub(crate) struct DiagCtxt(PhantomData<*mut ()>);
 pub(crate) struct ErrorGuaranteed(());
 
+struct DiagCtxtData {
+    diag: TokenStream,
+}
+
+thread_local! {
+    static DIAGNOSTICS: RefCell<Option<DiagCtxtData>> = const { RefCell::new(None) };
+}
+
+// Allows `syn::Error` to be emitted into the current diagnostic context with just `?`.
+impl From<syn::Error> for ErrorGuaranteed {
+    fn from(error: syn::Error) -> Self {
+        DIAGNOSTICS.with_borrow_mut(|data| {
+            data.as_mut()
+                .unwrap()
+                .diag
+                .extend(error.into_compile_error());
+        });
+        Self(())
+    }
+}
+
 impl DiagCtxt {
-    pub(crate) fn error(&mut self, span: impl Spanned, msg: impl Display) -> ErrorGuaranteed {
-        let error = Error::new(span.span(), msg);
-        self.0.extend(error.into_compile_error());
-        ErrorGuaranteed(())
+    pub(crate) fn error(&self, span: impl Spanned, msg: impl Display) -> ErrorGuaranteed {
+        Error::new(span.span(), msg).into()
     }
 
-    pub(crate) fn warn(&mut self, span: impl Spanned, msg: impl Display) {
+    pub(crate) fn warn(&self, span: impl Spanned, msg: impl Display) {
         // Have the message start on a new line for visual clarity.
         let msg = format!("\n{}", msg);
-        self.0.extend(quote_spanned!(span.span() =>
-            // Approximate using deprecated warning while `proc_macro_diagnostic` is unstable.
-            const _: () = {
-                #[deprecated = #msg]
-                const fn warn() {}
-                warn();
-            };
-        ));
+        DIAGNOSTICS.with_borrow_mut(|data| {
+            data.as_mut()
+                .unwrap()
+                .diag
+                .extend(quote_spanned!(span.span() =>
+                    // Approximate using deprecated warning while `proc_macro_diagnostic` is
+                    // unstable.
+                    const _: () = {
+                        #[deprecated = #msg]
+                        const fn warn() {}
+                        warn();
+                    };
+                ))
+        });
     }
 
-    pub(crate) fn with(
-        fun: impl FnOnce(&mut DiagCtxt) -> Result<TokenStream, ErrorGuaranteed>,
+    /// Execute the provided function with the current diagnostic context.
+    pub(crate) fn current<R>(f: impl FnOnce(&DiagCtxt) -> R) -> R {
+        DIAGNOSTICS.with_borrow(|data| {
+            assert!(data.is_some(), "No active `DiagCtxt`");
+        });
+
+        f(&DiagCtxt(PhantomData))
+    }
+
+    fn with(
+        f: impl FnOnce(&mut DiagCtxt) -> Result<TokenStream, ErrorGuaranteed>,
+        merge_diag: impl FnOnce(TokenStream, TokenStream) -> TokenStream,
+        convert_diag: impl FnOnce(TokenStream) -> TokenStream,
     ) -> TokenStream {
-        let mut dcx = Self(TokenStream::new());
-        match fun(&mut dcx) {
-            Ok(mut stream) => {
-                stream.extend(dcx.0);
-                stream
+        DIAGNOSTICS.with_borrow_mut(|data| {
+            assert!(data.is_none(), "`DiagCtxt` cannot be nested");
+            *data = Some(DiagCtxtData {
+                diag: TokenStream::new(),
+            });
+        });
+
+        let mut dcx = DiagCtxt(PhantomData);
+        let result = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| f(&mut dcx))) {
+            Ok(result) => result,
+            Err(payload) => {
+                // Robustness against panicking in macros.
+                //
+                // Ensure that any error messages are still emitted when this happens.
+                let message = if let Some(&s) = payload.downcast_ref::<&'static str>() {
+                    s
+                } else if let Some(s) = payload.downcast_ref::<String>() {
+                    s.as_str()
+                } else {
+                    "Box<dyn Any>"
+                };
+
+                Err(dcx.error(
+                    Span::mixed_site(),
+                    format!("proc macro panicked: {message}"),
+                ))
             }
-            Err(ErrorGuaranteed(())) => dcx.0,
+        };
+
+        let data = DIAGNOSTICS.with_borrow_mut(|data| data.take().unwrap());
+
+        match result {
+            Ok(stream) => {
+                if data.diag.is_empty() {
+                    stream
+                } else {
+                    merge_diag(stream, data.diag)
+                }
+            }
+            Err(ErrorGuaranteed(())) => convert_diag(data.diag),
         }
     }
+
+    pub(crate) fn for_item(
+        f: impl FnOnce(&mut DiagCtxt) -> Result<TokenStream, ErrorGuaranteed>,
+    ) -> TokenStream {
+        Self::with(
+            f,
+            |mut out, diag| {
+                out.extend(diag);
+                out
+            },
+            std::convert::identity,
+        )
+    }
+
+    pub(crate) fn for_expr(
+        f: impl FnOnce(&mut DiagCtxt) -> Result<TokenStream, ErrorGuaranteed>,
+    ) -> TokenStream {
+        Self::with(
+            f,
+            |out, diag| {
+                // Diagnostics that we generate are always items.
+                // So for expressions create a block to place diagnostics in item position.
+                quote!({
+                    #diag
+                    #out
+                })
+            },
+            |diag| {
+                quote!({
+                    #diag
+                })
+            },
+        )
+    }
 }

internal/src/init.rs

@@ -44,9 +44,6 @@ pub(crate) enum InitExprKind {
 
 struct InitTupleField {
     attrs: Vec<Attribute>,
-    /// `<-` is not valid in constructor syntax; it is parsed anyway so that it can be rejected
-    /// with a proper diagnostic instead of a parse error.
-    left_arrow_token: Option<Token![<-]>,
     value: Expr,
 }
 
@@ -84,20 +81,6 @@ impl InitExprTuple {
             rest: None,
         }
     }
-
-    fn validate(&self, dcx: &mut DiagCtxt) -> Result<(), ErrorGuaranteed> {
-        let mut result = Ok(());
-        for field in &self.fields {
-            if let Some(left_arrow_token) = &field.left_arrow_token {
-                result = Err(dcx.error(
-                    left_arrow_token,
-                    "`<-` is not supported in tuple constructor syntax; name the fields by index \
-                     instead, e.g. `Type { 0 <- initializer, 1: value }`",
-                ));
-            }
-        }
-        result
-    }
 }
 
 struct This {
@@ -153,8 +136,6 @@ pub(crate) fn expand_with_cfg(
 ) -> Result<TokenStream, ErrorGuaranteed> {
     let initializer = match initializer.kind {
         InitExprKind::Tuple(expr) => {
-            expr.validate(dcx)?;
-
             let mut initializer = Initializer {
                 attrs: initializer.attrs,
                 this: initializer.this,
@@ -579,9 +560,20 @@ impl InitExprTuple {
         let paren_token = parenthesized!(content in input);
         let mut fields = Punctuated::new();
         while !content.is_empty() {
+            let attrs = content.call(Attribute::parse_outer)?;
+
+            if let Some(left_arrow_token) = content.parse::<Option<Token![<-]>>()? {
+                DiagCtxt::current(|dcx| {
+                    dcx.error(
+                        left_arrow_token,
+                        "`<-` is not supported in tuple constructor syntax; name the fields by \
+                        index instead, e.g. `Type { 0 <- initializer, 1: value }`",
+                    )
+                });
+            }
+
             fields.push_value(InitTupleField {
-                attrs: content.call(Attribute::parse_outer)?,
-                left_arrow_token: content.parse()?,
+                attrs,
                 value: content.parse()?,
             });
             if content.is_empty() {
@@ -758,13 +750,8 @@ impl ToTokens for InitExprTuple {
 
 impl ToTokens for InitTupleField {
     fn to_tokens(&self, tokens: &mut TokenStream) {
-        let Self {
-            attrs,
-            left_arrow_token,
-            value,
-        } = self;
+        let Self { attrs, value } = self;
         tokens.append_all(attrs);
-        left_arrow_token.to_tokens(tokens);
         value.to_tokens(tokens);
     }
 }

internal/src/lib.rs

@@ -10,7 +10,6 @@
 #![allow(missing_docs)]
 
 use proc_macro::TokenStream;
-use syn::parse_macro_input;
 
 use crate::diagnostics::DiagCtxt;
 
@@ -23,43 +22,46 @@ mod zeroable;
 
 #[proc_macro_attribute]
 pub fn pin_data(args: TokenStream, input: TokenStream) -> TokenStream {
-    let args = parse_macro_input!(args);
-    let input = parse_macro_input!(input);
-    DiagCtxt::with(|dcx| pin_data::pin_data(args, input, dcx)).into()
+    DiagCtxt::for_item(|dcx| pin_data::pin_data(syn::parse(args)?, syn::parse(input)?, dcx)).into()
 }
 
 #[proc_macro_attribute]
 pub fn pinned_drop(args: TokenStream, input: TokenStream) -> TokenStream {
-    let args = parse_macro_input!(args);
-    let input = parse_macro_input!(input);
-    DiagCtxt::with(|dcx| pinned_drop::pinned_drop(args, input, dcx)).into()
+    DiagCtxt::for_item(|dcx| pinned_drop::pinned_drop(syn::parse(args)?, syn::parse(input)?, dcx))
+        .into()
 }
 
 #[proc_macro_derive(Zeroable)]
 pub fn derive_zeroable(input: TokenStream) -> TokenStream {
-    let input = parse_macro_input!(input);
-    DiagCtxt::with(|dcx| zeroable::derive(input, dcx)).into()
+    DiagCtxt::for_item(|dcx| zeroable::derive(syn::parse(input)?, dcx)).into()
 }
 
 #[proc_macro_derive(MaybeZeroable)]
 pub fn maybe_derive_zeroable(input: TokenStream) -> TokenStream {
-    let input = parse_macro_input!(input);
-    DiagCtxt::with(|dcx| zeroable::maybe_derive(input, dcx)).into()
+    DiagCtxt::for_item(|dcx| zeroable::maybe_derive(syn::parse(input)?, dcx)).into()
 }
 #[proc_macro]
 pub fn init(input: TokenStream) -> TokenStream {
-    let input = parse_macro_input!(input);
-    DiagCtxt::with(|dcx| {
-        init::expand_with_cfg(input, Some("::core::convert::Infallible"), false, dcx)
+    DiagCtxt::for_expr(|dcx| {
+        init::expand_with_cfg(
+            syn::parse(input)?,
+            Some("::core::convert::Infallible"),
+            false,
+            dcx,
+        )
     })
     .into()
 }
 
 #[proc_macro]
 pub fn pin_init(input: TokenStream) -> TokenStream {
-    let input = parse_macro_input!(input);
-    DiagCtxt::with(|dcx| {
-        init::expand_with_cfg(input, Some("::core::convert::Infallible"), true, dcx)
+    DiagCtxt::for_expr(|dcx| {
+        init::expand_with_cfg(
+            syn::parse(input)?,
+            Some("::core::convert::Infallible"),
+            true,
+            dcx,
+        )
     })
     .into()
 }

tests/ui/compile-fail/zeroable/invalid_spread.rs

@@ -0,0 +1,17 @@
+extern crate pin_init;
+use pin_init::*;
+
+use Zeroable as MyZeroable;
+
+#[derive(Zeroable)]
+struct Foo {
+    a: usize,
+    b: usize,
+}
+
+fn main() {
+    let _ = init!(Foo {
+        a: 0,
+        ..MyZeroable::init_zeroed()
+    });
+}

tests/ui/compile-fail/zeroable/invalid_spread.stderr

@@ -0,0 +1,19 @@
+error: expected nothing or `..Zeroable::init_zeroed()`.
+  --> tests/ui/compile-fail/zeroable/invalid_spread.rs:15:9
+   |
+15 |         ..MyZeroable::init_zeroed()
+   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+warning: unused import: `Zeroable as MyZeroable`
+ --> tests/ui/compile-fail/zeroable/invalid_spread.rs:4:5
+  |
+4 | use Zeroable as MyZeroable;
+  |     ^^^^^^^^^^^^^^^^^^^^^^
+  |
+  = note: `#[warn(unused_imports)]` (part of `#[warn(unused)]`) on by default
+
+error[E0063]: missing field `b` in initializer of `Foo`
+  --> tests/ui/compile-fail/zeroable/invalid_spread.rs:13:19
+   |
+13 |     let _ = init!(Foo {
+   |                   ^^^ missing `b`