[PATCH 2/3] rust: pin-init: internal: make `DiagCtxt` available inside parser

[Gary Guo <gary@kernel.org>]

From: Gary Guo <gary@garyguo.net>

Currently, there is a split between `syn::Error` and pin-init's custom
`DiagCtxt`. Parsing has no access to `DiagCtxt` and therefore cannot report
diagnostics without failing the parse. Bridge the gap by making the
`DiagCtxt` available inside parsers using `thread_local!`.

`parse_macro_input!` returns error as token stream directly, so it is no
longer suitable when the parsing is moved inside `DiagCtxt` closure. Add a
`From<syn::Error> for ErrorGuaranteed` implementation so `?` can be used to
add a `syn::Error` into the current diagnostic context and obtain a
`ErrorGuaranteed`.

Clean up the handler of using `<-` inside tuple expression as a
demonstration of the usefulness of having `DiagCtxt` access inside `Parse`.

Signed-off-by: Gary Guo <gary@garyguo.net>
---
 rust/pin-init/internal/src/diagnostics.rs | 84 ++++++++++++++++++-----
 rust/pin-init/internal/src/init.rs        | 41 ++++-------
 rust/pin-init/internal/src/lib.rs         | 32 +++++----
 3 files changed, 96 insertions(+), 61 deletions(-)

diff --git a/rust/pin-init/internal/src/diagnostics.rs b/rust/pin-init/internal/src/diagnostics.rs
index c42f1095ab10..efdcf45f40cf 100644
--- a/rust/pin-init/internal/src/diagnostics.rs
+++ b/rust/pin-init/internal/src/diagnostics.rs
@@ -1,32 +1,68 @@
 // 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, 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();
+                    };
+                ))
+        });
+    }
+
+    /// 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(
@@ -34,16 +70,26 @@ fn with(
         merge_diag: impl FnOnce(TokenStream, TokenStream) -> TokenStream,
         convert_diag: impl FnOnce(TokenStream) -> TokenStream,
     ) -> TokenStream {
-        let mut dcx = Self(TokenStream::new());
-        match f(&mut dcx) {
+        DIAGNOSTICS.with_borrow_mut(|data| {
+            assert!(data.is_none(), "`DiagCtxt` cannot be nested");
+            *data = Some(DiagCtxtData {
+                diag: TokenStream::new(),
+            });
+        });
+
+        let result = f(&mut DiagCtxt(PhantomData));
+
+        let data = DIAGNOSTICS.with_borrow_mut(|data| data.take().unwrap());
+
+        match result {
             Ok(stream) => {
-                if dcx.0.is_empty() {
+                if data.diag.is_empty() {
                     stream
                 } else {
-                    merge_diag(stream, dcx.0)
+                    merge_diag(stream, data.diag)
                 }
             }
-            Err(ErrorGuaranteed(())) => convert_diag(dcx.0),
+            Err(ErrorGuaranteed(())) => convert_diag(data.diag),
         }
     }
 
diff --git a/rust/pin-init/internal/src/init.rs b/rust/pin-init/internal/src/init.rs
index 1d2db93dd33d..b60feb68b65f 100644
--- a/rust/pin-init/internal/src/init.rs
+++ b/rust/pin-init/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 @@ fn normalize(self) -> InitExprStruct {
             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,
@@ -578,9 +559,20 @@ fn parse_with_path(path: Path, input: syn::parse::ParseStream<'_>) -> syn::Resul
         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() {
@@ -757,13 +749,8 @@ fn to_tokens(&self, tokens: &mut TokenStream) {
 
 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);
     }
 }
diff --git a/rust/pin-init/internal/src/lib.rs b/rust/pin-init/internal/src/lib.rs
index 0410024ba1e6..07b3d33d2282 100644
--- a/rust/pin-init/internal/src/lib.rs
+++ b/rust/pin-init/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 @@
 
 #[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::for_item(|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::for_item(|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::for_item(|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::for_item(|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::for_expr(|dcx| {
-        init::expand_with_cfg(input, Some("::core::convert::Infallible"), false, 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::for_expr(|dcx| {
-        init::expand_with_cfg(input, Some("::core::convert::Infallible"), true, dcx)
+        init::expand_with_cfg(
+            syn::parse(input)?,
+            Some("::core::convert::Infallible"),
+            true,
+            dcx,
+        )
     })
     .into()
 }
-- 
2.54.0