5 files changed, 62 insertions, 10 deletions

diff --git a/crates/proc_macro_srv/src/dylib.rs b/crates/proc_macro_srv/src/dylib.rs
index 2afb973cc..4e719f3aa 100644
--- a/crates/proc_macro_srv/src/dylib.rs
+++ b/crates/proc_macro_srv/src/dylib.rs
@@ -136,6 +136,7 @@ impl Expander {
                         &crate::proc_macro::bridge::server::SameThread,
                         crate::rustc_server::Rustc::default(),
                         parsed_body,
+                        false,
                     );
                     return res.map(|it| it.subtree);
                 }
@@ -144,6 +145,7 @@ impl Expander {
                         &crate::proc_macro::bridge::server::SameThread,
                         crate::rustc_server::Rustc::default(),
                         parsed_body,
+                        false,
                     );
                     return res.map(|it| it.subtree);
                 }
@@ -153,6 +155,7 @@ impl Expander {
                         crate::rustc_server::Rustc::default(),
                         parsed_attributes,
                         parsed_body,
+                        false,
                     );
                     return res.map(|it| it.subtree);
                 }
diff --git a/crates/proc_macro_srv/src/lib.rs b/crates/proc_macro_srv/src/lib.rs
index 9cca96994..d4f04ee06 100644
--- a/crates/proc_macro_srv/src/lib.rs
+++ b/crates/proc_macro_srv/src/lib.rs
@@ -24,7 +24,7 @@ use proc_macro::bridge::client::TokenStream;
 use proc_macro_api::{ExpansionResult, ExpansionTask, ListMacrosResult, ListMacrosTask};
 use std::{
     collections::{hash_map::Entry, HashMap},
-    fs,
+    env, fs,
     path::{Path, PathBuf},
     time::SystemTime,
 };
@@ -37,7 +37,23 @@ pub(crate) struct ProcMacroSrv {
 impl ProcMacroSrv {
     pub fn expand(&mut self, task: &ExpansionTask) -> Result<ExpansionResult, String> {
         let expander = self.expander(&task.lib)?;
-        match expander.expand(&task.macro_name, &task.macro_body, task.attributes.as_ref()) {
+
+        let mut prev_env = HashMap::new();
+        for (k, v) in &task.env {
+            prev_env.insert(k.as_str(), env::var_os(k));
+            env::set_var(k, v);
+        }
+
+        let result = expander.expand(&task.macro_name, &task.macro_body, task.attributes.as_ref());
+
+        for (k, _) in &task.env {
+            match &prev_env[k.as_str()] {
+                Some(v) => env::set_var(k, v),
+                None => env::remove_var(k),
+            }
+        }
+
+        match result {
             Ok(expansion) => Ok(ExpansionResult { expansion }),
             Err(msg) => {
                 let msg = msg.as_str().unwrap_or("<unknown error>");
diff --git a/crates/proc_macro_srv/src/proc_macro/bridge/client.rs b/crates/proc_macro_srv/src/proc_macro/bridge/client.rs
index 55d6330cc..ca6749b9b 100644
--- a/crates/proc_macro_srv/src/proc_macro/bridge/client.rs
+++ b/crates/proc_macro_srv/src/proc_macro/bridge/client.rs
@@ -303,17 +303,21 @@ impl BridgeState<'_> {
 
 impl Bridge<'_> {
     fn enter<R>(self, f: impl FnOnce() -> R) -> R {
+        let force_show_panics = self.force_show_panics;
+
         // Hide the default panic output within `proc_macro` expansions.
         // NB. the server can't do this because it may use a different libstd.
         static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
         HIDE_PANICS_DURING_EXPANSION.call_once(|| {
             let prev = panic::take_hook();
             panic::set_hook(Box::new(move |info| {
-                let hide = BridgeState::with(|state| match state {
-                    BridgeState::NotConnected => false,
-                    BridgeState::Connected(_) | BridgeState::InUse => true,
+                let show = BridgeState::with(|state| match state {
+                    BridgeState::NotConnected => true,
+                    // Something weird is going on, so don't suppress any backtraces
+                    BridgeState::InUse => true,
+                    BridgeState::Connected(bridge) => force_show_panics,
                 });
-                if !hide {
+                if show {
                     prev(info)
                 }
             }));
diff --git a/crates/proc_macro_srv/src/proc_macro/bridge/mod.rs b/crates/proc_macro_srv/src/proc_macro/bridge/mod.rs
index b97886eb9..e4006a5ab 100644
--- a/crates/proc_macro_srv/src/proc_macro/bridge/mod.rs
+++ b/crates/proc_macro_srv/src/proc_macro/bridge/mod.rs
@@ -225,6 +225,9 @@ pub struct Bridge<'a> {
 
     /// Server-side function that the client uses to make requests.
     dispatch: closure::Closure<'a, Buffer<u8>, Buffer<u8>>,
+
+    /// If 'true', always invoke the default panic hook
+    force_show_panics: bool,
 }
 
 // impl<'a> !Sync for Bridge<'a> {}
diff --git a/crates/proc_macro_srv/src/proc_macro/bridge/server.rs b/crates/proc_macro_srv/src/proc_macro/bridge/server.rs
index 3acb239af..88fbdc078 100644
--- a/crates/proc_macro_srv/src/proc_macro/bridge/server.rs
+++ b/crates/proc_macro_srv/src/proc_macro/bridge/server.rs
@@ -138,6 +138,7 @@ pub trait ExecutionStrategy {
         input: Buffer<u8>,
         run_client: extern "C" fn(Bridge<'_>, D) -> Buffer<u8>,
         client_data: D,
+        force_show_panics: bool,
     ) -> Buffer<u8>;
 }
 
@@ -150,10 +151,14 @@ impl ExecutionStrategy for SameThread {
         input: Buffer<u8>,
         run_client: extern "C" fn(Bridge<'_>, D) -> Buffer<u8>,
         client_data: D,
+        force_show_panics: bool,
     ) -> Buffer<u8> {
         let mut dispatch = |b| dispatcher.dispatch(b);
 
-        run_client(Bridge { cached_buffer: input, dispatch: (&mut dispatch).into() }, client_data)
+        run_client(
+            Bridge { cached_buffer: input, dispatch: (&mut dispatch).into(), force_show_panics },
+            client_data,
+        )
     }
 }
 
@@ -169,6 +174,7 @@ impl ExecutionStrategy for CrossThread1 {
         input: Buffer<u8>,
         run_client: extern "C" fn(Bridge<'_>, D) -> Buffer<u8>,
         client_data: D,
+        force_show_panics: bool,
     ) -> Buffer<u8> {
         use std::sync::mpsc::channel;
 
@@ -182,7 +188,11 @@ impl ExecutionStrategy for CrossThread1 {
             };
 
             run_client(
-                Bridge { cached_buffer: input, dispatch: (&mut dispatch).into() },
+                Bridge {
+                    cached_buffer: input,
+                    dispatch: (&mut dispatch).into(),
+                    force_show_panics,
+                },
                 client_data,
             )
         });
@@ -204,6 +214,7 @@ impl ExecutionStrategy for CrossThread2 {
         input: Buffer<u8>,
         run_client: extern "C" fn(Bridge<'_>, D) -> Buffer<u8>,
         client_data: D,
+        force_show_panics: bool,
     ) -> Buffer<u8> {
         use std::sync::{Arc, Mutex};
 
@@ -229,7 +240,11 @@ impl ExecutionStrategy for CrossThread2 {
             };
 
             let r = run_client(
-                Bridge { cached_buffer: input, dispatch: (&mut dispatch).into() },
+                Bridge {
+                    cached_buffer: input,
+                    dispatch: (&mut dispatch).into(),
+                    force_show_panics,
+                },
                 client_data,
             );
 
@@ -268,6 +283,7 @@ fn run_server<
     input: I,
     run_client: extern "C" fn(Bridge<'_>, D) -> Buffer<u8>,
     client_data: D,
+    force_show_panics: bool,
 ) -> Result<O, PanicMessage> {
     let mut dispatcher =
         Dispatcher { handle_store: HandleStore::new(handle_counters), server: MarkedTypes(server) };
@@ -275,7 +291,13 @@ fn run_server<
     let mut b = Buffer::new();
     input.encode(&mut b, &mut dispatcher.handle_store);
 
-    b = strategy.run_bridge_and_client(&mut dispatcher, b, run_client, client_data);
+    b = strategy.run_bridge_and_client(
+        &mut dispatcher,
+        b,
+        run_client,
+        client_data,
+        force_show_panics,
+    );
 
     Result::decode(&mut &b[..], &mut dispatcher.handle_store)
 }
@@ -286,6 +308,7 @@ impl client::Client<fn(crate::TokenStream) -> crate::TokenStream> {
         strategy: &impl ExecutionStrategy,
         server: S,
         input: S::TokenStream,
+        force_show_panics: bool,
     ) -> Result<S::TokenStream, PanicMessage> {
         let client::Client { get_handle_counters, run, f } = *self;
         run_server(
@@ -295,6 +318,7 @@ impl client::Client<fn(crate::TokenStream) -> crate::TokenStream> {
             <MarkedTypes<S> as Types>::TokenStream::mark(input),
             run,
             f,
+            force_show_panics,
         )
         .map(<MarkedTypes<S> as Types>::TokenStream::unmark)
     }
@@ -307,6 +331,7 @@ impl client::Client<fn(crate::TokenStream, crate::TokenStream) -> crate::TokenSt
         server: S,
         input: S::TokenStream,
         input2: S::TokenStream,
+        force_show_panics: bool,
     ) -> Result<S::TokenStream, PanicMessage> {
         let client::Client { get_handle_counters, run, f } = *self;
         run_server(
@@ -319,6 +344,7 @@ impl client::Client<fn(crate::TokenStream, crate::TokenStream) -> crate::TokenSt
             ),
             run,
             f,
+            force_show_panics,
         )
         .map(<MarkedTypes<S> as Types>::TokenStream::unmark)
     }