Pass crate environment to proc macros

1 files changed, 18 insertions, 2 deletions

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>");