[PATCH 3/3] rust: pin-init: internal: improve diagnostics robustness against panicking

[Gary Guo]

From: Gary Guo <gary@xxxxxxxxxxx>

Currently, if proc macro panicked, the diagnostics clean up is not
executed, and further invocation will cause the "DiagCtxt cannot be nested"
error. While we should aim to have no panics inside proc macros, producing
a sensible diagnostics message even when macro panicked is very useful for
developing.

Thus, catch proc macro panics and convert them to errors, and emit them
together with all diagnostics accumulated so far.

Ideally we would like panic location w/ line numbers being available as
well; however this is not currently implementable without overriding the
global panic hook.

Signed-off-by: Gary Guo <gary@xxxxxxxxxxx>
---
 rust/pin-init/internal/src/diagnostics.rs | 24 +++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/rust/pin-init/internal/src/diagnostics.rs b/rust/pin-init/internal/src/diagnostics.rs
index efdcf45f40cf..e88e520e326b 100644
--- a/rust/pin-init/internal/src/diagnostics.rs
+++ b/rust/pin-init/internal/src/diagnostics.rs
@@ -4,7 +4,7 @@
 use std::fmt::Display;
 use std::marker::PhantomData;
 
-use proc_macro2::TokenStream;
+use proc_macro2::{Span, TokenStream};
 use quote::{quote, quote_spanned};
 use syn::{spanned::Spanned, Error};
 
@@ -77,7 +77,27 @@ fn with(
             });
         });
 
-        let result = f(&mut DiagCtxt(PhantomData));
+        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}"),
+                ))
+            }
+        };
 
         let data = DIAGNOSTICS.with_borrow_mut(|data| data.take().unwrap());
 
-- 
2.54.0