Pass crate environment to proc macros

6 files changed, 42 insertions, 9 deletions

diff --git a/crates/base_db/src/input.rs b/crates/base_db/src/input.rs
index a693e7f80..9567bcc42 100644
--- a/crates/base_db/src/input.rs
+++ b/crates/base_db/src/input.rs
@@ -151,8 +151,12 @@ pub enum ProcMacroKind {
 }
 
 pub trait ProcMacroExpander: fmt::Debug + Send + Sync + RefUnwindSafe {
-    fn expand(&self, subtree: &Subtree, attrs: Option<&Subtree>)
-        -> Result<Subtree, ExpansionError>;
+    fn expand(
+        &self,
+        subtree: &Subtree,
+        attrs: Option<&Subtree>,
+        env: &Env,
+    ) -> Result<Subtree, ExpansionError>;
 }
 
 #[derive(Debug, Clone)]
@@ -418,6 +422,10 @@ impl Env {
     pub fn get(&self, env: &str) -> Option<String> {
         self.entries.get(env).cloned()
     }
+
+    pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
+        self.entries.iter().map(|(k, v)| (k.as_str(), v.as_str()))
+    }
 }
 
 #[derive(Debug)]
diff --git a/crates/hir_expand/src/db.rs b/crates/hir_expand/src/db.rs
index 077de3727..06f0a3ed9 100644
--- a/crates/hir_expand/src/db.rs
+++ b/crates/hir_expand/src/db.rs
@@ -271,7 +271,7 @@ fn expand_proc_macro(
         _ => unreachable!(),
     };
 
-    expander.expand(db, lazy_id, &macro_arg.0)
+    expander.expand(db, loc.krate, &macro_arg.0)
 }
 
 fn parse_or_expand(db: &dyn AstDatabase, file_id: HirFileId) -> Option<SyntaxNode> {
diff --git a/crates/hir_expand/src/proc_macro.rs b/crates/hir_expand/src/proc_macro.rs
index 38882d2b6..7c77f6ce0 100644
--- a/crates/hir_expand/src/proc_macro.rs
+++ b/crates/hir_expand/src/proc_macro.rs
@@ -1,6 +1,6 @@
 //! Proc Macro Expander stub
 
-use crate::{db::AstDatabase, LazyMacroId};
+use crate::db::AstDatabase;
 use base_db::{CrateId, ProcMacroId};
 use tt::buffer::{Cursor, TokenBuffer};
 
@@ -32,7 +32,7 @@ impl ProcMacroExpander {
     pub fn expand(
         self,
         db: &dyn AstDatabase,
-        _id: LazyMacroId,
+        calling_crate: CrateId,
         tt: &tt::Subtree,
     ) -> Result<tt::Subtree, mbe::ExpandError> {
         match self.proc_macro_id {
@@ -47,7 +47,10 @@ impl ProcMacroExpander {
                 let tt = remove_derive_attrs(tt)
                     .ok_or_else(|| err!("Fail to remove derive for custom derive"))?;
 
-                proc_macro.expander.expand(&tt, None).map_err(mbe::ExpandError::from)
+                // Proc macros have access to the environment variables of the invoking crate.
+                let env = &krate_graph[calling_crate].env;
+
+                proc_macro.expander.expand(&tt, None, &env).map_err(mbe::ExpandError::from)
             }
             None => Err(mbe::ExpandError::UnresolvedProcMacro),
         }
diff --git a/crates/proc_macro_api/src/lib.rs b/crates/proc_macro_api/src/lib.rs
index a61afbbd6..2ea456fb0 100644
--- a/crates/proc_macro_api/src/lib.rs
+++ b/crates/proc_macro_api/src/lib.rs
@@ -16,7 +16,7 @@ use std::{
     sync::Arc,
 };
 
-use base_db::ProcMacro;
+use base_db::{Env, ProcMacro};
 use tt::{SmolStr, Subtree};
 
 use crate::process::{ProcMacroProcessSrv, ProcMacroProcessThread};
@@ -44,12 +44,14 @@ impl base_db::ProcMacroExpander for ProcMacroProcessExpander {
         &self,
         subtree: &Subtree,
         attr: Option<&Subtree>,
+        env: &Env,
     ) -> Result<Subtree, tt::ExpansionError> {
         let task = ExpansionTask {
             macro_body: subtree.clone(),
             macro_name: self.name.to_string(),
             attributes: attr.cloned(),
             lib: self.dylib_path.to_path_buf(),
+            env: env.iter().map(|(k, v)| (k.to_string(), v.to_string())).collect(),
         };
 
         let result: ExpansionResult = self.process.send_task(msg::Request::ExpansionMacro(task))?;
diff --git a/crates/proc_macro_api/src/rpc.rs b/crates/proc_macro_api/src/rpc.rs
index b85f92eea..cf830b59f 100644
--- a/crates/proc_macro_api/src/rpc.rs
+++ b/crates/proc_macro_api/src/rpc.rs
@@ -51,6 +51,9 @@ pub struct ExpansionTask {
     pub attributes: Option<Subtree>,
 
     pub lib: PathBuf,
+
+    /// Environment variables to set during macro expansion.
+    pub env: Vec<(String, String)>,
 }
 
 #[derive(Clone, Eq, PartialEq, Debug, Default, Serialize, Deserialize)]
@@ -251,6 +254,7 @@ mod tests {
             macro_name: Default::default(),
             attributes: None,
             lib: Default::default(),
+            env: Default::default(),
         };
 
         let json = serde_json::to_string(&task).unwrap();
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>");