Pass the crate environment to proc macros

crates/base_db/src/input.rs

@@ -6,12 +6,12 @@
 //! actual IO. See `vfs` and `project_model` in the `rust-analyzer` crate for how
 //! actual IO is done and lowered to input.
 
-use std::{fmt, iter::FromIterator, ops, str::FromStr, sync::Arc};
+use std::{fmt, iter::FromIterator, ops, panic::RefUnwindSafe, str::FromStr, sync::Arc};
 
 use cfg::CfgOptions;
 use rustc_hash::{FxHashMap, FxHashSet};
 use syntax::SmolStr;
-use tt::TokenExpander;
+use tt::{ExpansionError, Subtree};
 use vfs::{file_set::FileSet, FileId, VfsPath};
 
 /// Files are grouped into source roots. A source root is a directory on the
@@ -150,11 +150,20 @@ pub enum ProcMacroKind {
     Attr,
 }
 
+pub trait ProcMacroExpander: fmt::Debug + Send + Sync + RefUnwindSafe {
+    fn expand(
+        &self,
+        subtree: &Subtree,
+        attrs: Option<&Subtree>,
+        env: &Env,
+    ) -> Result<Subtree, ExpansionError>;
+}
+
 #[derive(Debug, Clone)]
 pub struct ProcMacro {
     pub name: SmolStr,
     pub kind: ProcMacroKind,
-    pub expander: Arc<dyn TokenExpander>,
+    pub expander: Arc<dyn ProcMacroExpander>,
 }
 
 impl Eq for ProcMacro {}
@@ -413,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)]

crates/base_db/src/lib.rs

@@ -14,7 +14,7 @@ pub use crate::{
     change::Change,
     input::{
         CrateData, CrateDisplayName, CrateGraph, CrateId, CrateName, Dependency, Edition, Env,
-        ProcMacro, ProcMacroId, ProcMacroKind, SourceRoot, SourceRootId,
+        ProcMacro, ProcMacroExpander, ProcMacroId, ProcMacroKind, SourceRoot, SourceRootId,
     },
 };
 pub use salsa;

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> {

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),
         }

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};
@@ -39,17 +39,19 @@ impl PartialEq for ProcMacroProcessExpander {
     }
 }
 
-impl tt::TokenExpander for ProcMacroProcessExpander {
+impl base_db::ProcMacroExpander for ProcMacroProcessExpander {
     fn expand(
         &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))?;
@@ -90,7 +92,7 @@ impl ProcMacroClient {
                     ProcMacroKind::FuncLike => base_db::ProcMacroKind::FuncLike,
                     ProcMacroKind::Attr => base_db::ProcMacroKind::Attr,
                 };
-                let expander: Arc<dyn tt::TokenExpander> = Arc::new(ProcMacroProcessExpander {
+                let expander = Arc::new(ProcMacroProcessExpander {
                     process: self.process.clone(),
                     name: name.clone(),
                     dylib_path: dylib_path.into(),

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

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

crates/tt/src/lib.rs

@@ -1,7 +1,7 @@
 //! `tt` crate defines a `TokenTree` data structure: this is the interface (both
 //! input and output) of macros. It closely mirrors `proc_macro` crate's
 //! `TokenTree`.
-use std::{fmt, panic::RefUnwindSafe};
+use std::fmt;
 
 use stdx::impl_from;
 
@@ -247,8 +247,3 @@ impl fmt::Display for ExpansionError {
         }
     }
 }
-
-pub trait TokenExpander: fmt::Debug + Send + Sync + RefUnwindSafe {
-    fn expand(&self, subtree: &Subtree, attrs: Option<&Subtree>)
-        -> Result<Subtree, ExpansionError>;
-}