[PATCH 1/3] rust: pin-init: internal: handle item and expr diagnostics differently

[Gary Guo]

From: Gary Guo <gary@xxxxxxxxxxx>

Currently `DiagCtxt` is used to allow error recovery in the pin-init
macros, by generating both a `compile_error!()` macro item and continue
expansion, and then merge token streams together. This works very well for
items, however for expressions, this will produce invalid expression and
result in a confusing error message.

    error: macro expansion ignores `::` and any tokens following
      --> tests/ui/compile-fail/zeroable/invalid_spread.rs:15:9
       |
    13 |       let _ = init!(Foo {
       |  _____________-
    14 | |         a: 0,
    15 | |         ..MyZeroable::init_zeroed()
       | |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^
    16 | |     });
       | |______- caused by the macro expansion here
       |
       = note: the usage of `init!` is likely invalid in expression context

Fix this by wrap the concatenated diagnostics in a block, with the
generated expression in its tail position. This results in the desired
error message.

    error: expected nothing or `..Zeroable::init_zeroed()`.
      --> tests/ui/compile-fail/zeroable/invalid_spread.rs:15:9
       |
    15 |         ..MyZeroable::init_zeroed()
       |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^

Signed-off-by: Gary Guo <gary@xxxxxxxxxxx>
---
 rust/pin-init/internal/src/diagnostics.rs | 55 +++++++++++++++++++----
 rust/pin-init/internal/src/lib.rs         | 12 ++---
 2 files changed, 53 insertions(+), 14 deletions(-)

diff --git a/rust/pin-init/internal/src/diagnostics.rs b/rust/pin-init/internal/src/diagnostics.rs
index c7d9b3e624fc..c42f1095ab10 100644
--- a/rust/pin-init/internal/src/diagnostics.rs
+++ b/rust/pin-init/internal/src/diagnostics.rs
@@ -3,7 +3,7 @@
 use std::fmt::Display;
 
 use proc_macro2::TokenStream;
-use quote::quote_spanned;
+use quote::{quote, quote_spanned};
 use syn::{spanned::Spanned, Error};
 
 pub(crate) struct DiagCtxt(TokenStream);
@@ -29,16 +29,55 @@ const fn warn() {}
         ));
     }
 
-    pub(crate) fn with(
-        fun: impl FnOnce(&mut DiagCtxt) -> Result<TokenStream, ErrorGuaranteed>,
+    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
+        match f(&mut dcx) {
+            Ok(stream) => {
+                if dcx.0.is_empty() {
+                    stream
+                } else {
+                    merge_diag(stream, dcx.0)
+                }
             }
-            Err(ErrorGuaranteed(())) => dcx.0,
+            Err(ErrorGuaranteed(())) => convert_diag(dcx.0),
         }
     }
+
+    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
+                })
+            },
+        )
+    }
 }
diff --git a/rust/pin-init/internal/src/lib.rs b/rust/pin-init/internal/src/lib.rs
index c488019d6250..0410024ba1e6 100644
--- a/rust/pin-init/internal/src/lib.rs
+++ b/rust/pin-init/internal/src/lib.rs
@@ -25,31 +25,31 @@
 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(args, 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(args, 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(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(input, dcx)).into()
 }
 #[proc_macro]
 pub fn init(input: TokenStream) -> TokenStream {
     let input = parse_macro_input!(input);
-    DiagCtxt::with(|dcx| {
+    DiagCtxt::for_expr(|dcx| {
         init::expand_with_cfg(input, Some("::core::convert::Infallible"), false, dcx)
     })
     .into()
@@ -58,7 +58,7 @@ pub fn init(input: TokenStream) -> TokenStream {
 #[proc_macro]
 pub fn pin_init(input: TokenStream) -> TokenStream {
     let input = parse_macro_input!(input);
-    DiagCtxt::with(|dcx| {
+    DiagCtxt::for_expr(|dcx| {
         init::expand_with_cfg(input, Some("::core::convert::Infallible"), true, dcx)
     })
     .into()

base-commit: dfb6a037fd586f1ffcba3b143dab58f5c88294d1
-- 
2.54.0