diff --git a/crates/base_db/src/input.rs b/crates/base_db/src/input.rs index a693e7f80d..9567bcc424 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; + fn expand( + &self, + subtree: &Subtree, + attrs: Option<&Subtree>, + env: &Env, + ) -> Result; } #[derive(Debug, Clone)] @@ -418,6 +422,10 @@ impl Env { pub fn get(&self, env: &str) -> Option { self.entries.get(env).cloned() } + + pub fn iter(&self) -> impl Iterator { + 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 077de37270..06f0a3ed9c 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, ¯o_arg.0) + expander.expand(db, loc.krate, ¯o_arg.0) } fn parse_or_expand(db: &dyn AstDatabase, file_id: HirFileId) -> Option { diff --git a/crates/hir_expand/src/proc_macro.rs b/crates/hir_expand/src/proc_macro.rs index 38882d2b68..7c77f6ce0a 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 { 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 a61afbbd66..2ea456fb09 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 { 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 b85f92eead..cf830b59f1 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, 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 9cca969944..d4f04ee069 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 { 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("");