Faster env snapshotting in proc-macro-srv

6 files changed, 147 insertions, 117 deletions

diff --git a/src/tools/rust-analyzer/crates/proc-macro-api/src/lib.rs b/src/tools/rust-analyzer/crates/proc-macro-api/src/lib.rs
index 039509d87c2..9ba6a09cb73 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-api/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-api/src/lib.rs
@@ -156,23 +156,26 @@ impl ProcMacro {
         let call_site = span_data_table.insert_full(call_site).0;
         let mixed_site = span_data_table.insert_full(mixed_site).0;
         let task = ExpandMacro {
-            macro_body: FlatTree::new(subtree, version, &mut span_data_table),
-            macro_name: self.name.to_string(),
-            attributes: attr.map(|subtree| FlatTree::new(subtree, version, &mut span_data_table)),
+            data: msg::ExpandMacroData {
+                macro_body: FlatTree::new(subtree, version, &mut span_data_table),
+                macro_name: self.name.to_string(),
+                attributes: attr
+                    .map(|subtree| FlatTree::new(subtree, version, &mut span_data_table)),
+                has_global_spans: ExpnGlobals {
+                    serialize: version >= HAS_GLOBAL_SPANS,
+                    def_site,
+                    call_site,
+                    mixed_site,
+                },
+                span_data_table: if version >= RUST_ANALYZER_SPAN_SUPPORT {
+                    serialize_span_data_index_map(&span_data_table)
+                } else {
+                    Vec::new()
+                },
+            },
             lib: self.dylib_path.to_path_buf().into(),
             env: env.into(),
             current_dir,
-            has_global_spans: ExpnGlobals {
-                serialize: version >= HAS_GLOBAL_SPANS,
-                def_site,
-                call_site,
-                mixed_site,
-            },
-            span_data_table: if version >= RUST_ANALYZER_SPAN_SUPPORT {
-                serialize_span_data_index_map(&span_data_table)
-            } else {
-                Vec::new()
-            },
         };
 
         let response = self.process.send_task(msg::Request::ExpandMacro(Box::new(task)))?;
diff --git a/src/tools/rust-analyzer/crates/proc-macro-api/src/msg.rs b/src/tools/rust-analyzer/crates/proc-macro-api/src/msg.rs
index ad0e1f187b6..49ebedba7dd 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-api/src/msg.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-api/src/msg.rs
@@ -72,6 +72,16 @@ pub struct PanicMessage(pub String);
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ExpandMacro {
+    pub lib: Utf8PathBuf,
+    /// Environment variables to set during macro expansion.
+    pub env: Vec<(String, String)>,
+    pub current_dir: Option<String>,
+    #[serde(flatten)]
+    pub data: ExpandMacroData,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ExpandMacroData {
     /// Argument of macro call.
     ///
     /// In custom derive this will be a struct or enum; in attribute-like macro - underlying
@@ -86,13 +96,6 @@ pub struct ExpandMacro {
 
     /// Possible attributes for the attribute-like macros.
     pub attributes: Option<FlatTree>,
-
-    pub lib: Utf8PathBuf,
-
-    /// Environment variables to set during macro expansion.
-    pub env: Vec<(String, String)>,
-
-    pub current_dir: Option<String>,
     /// marker for serde skip stuff
     #[serde(skip_serializing_if = "ExpnGlobals::skip_serializing_if")]
     #[serde(default)]
@@ -268,25 +271,30 @@ mod tests {
         let tt = fixture_token_tree();
         let mut span_data_table = Default::default();
         let task = ExpandMacro {
-            macro_body: FlatTree::new(&tt, CURRENT_API_VERSION, &mut span_data_table),
-            macro_name: Default::default(),
-            attributes: None,
+            data: ExpandMacroData {
+                macro_body: FlatTree::new(&tt, CURRENT_API_VERSION, &mut span_data_table),
+                macro_name: Default::default(),
+                attributes: None,
+                has_global_spans: ExpnGlobals {
+                    serialize: true,
+                    def_site: 0,
+                    call_site: 0,
+                    mixed_site: 0,
+                },
+                span_data_table: Vec::new(),
+            },
             lib: Utf8PathBuf::from_path_buf(std::env::current_dir().unwrap()).unwrap(),
             env: Default::default(),
             current_dir: Default::default(),
-            has_global_spans: ExpnGlobals {
-                serialize: true,
-                def_site: 0,
-                call_site: 0,
-                mixed_site: 0,
-            },
-            span_data_table: Vec::new(),
         };
 
         let json = serde_json::to_string(&task).unwrap();
         // println!("{}", json);
         let back: ExpandMacro = serde_json::from_str(&json).unwrap();
 
-        assert_eq!(tt, back.macro_body.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table));
+        assert_eq!(
+            tt,
+            back.data.macro_body.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table)
+        );
     }
 }
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv-cli/src/main.rs b/src/tools/rust-analyzer/crates/proc-macro-srv-cli/src/main.rs
index df0ae3171f5..f9689712ad6 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv-cli/src/main.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv-cli/src/main.rs
@@ -33,12 +33,14 @@ fn run() -> io::Result<()> {
 #[cfg(any(feature = "sysroot-abi", rust_analyzer))]
 fn run() -> io::Result<()> {
     use proc_macro_api::msg::{self, Message};
+    use proc_macro_srv::EnvSnapshot;
 
     let read_request = |buf: &mut String| msg::Request::read(&mut io::stdin().lock(), buf);
 
     let write_response = |msg: msg::Response| msg.write(&mut io::stdout().lock());
 
-    let mut srv = proc_macro_srv::ProcMacroSrv::default();
+    let env = EnvSnapshot::new();
+    let mut srv = proc_macro_srv::ProcMacroSrv::new(&env);
     let mut buf = String::new();
 
     while let Some(req) = read_request(&mut buf)? {
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/lib.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/lib.rs
index 3d309155d5d..34851ee0bee 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/lib.rs
@@ -32,7 +32,9 @@ use std::{
     collections::{hash_map::Entry, HashMap},
     env,
     ffi::OsString,
-    fs, thread,
+    fs,
+    path::{Path, PathBuf},
+    thread,
     time::SystemTime,
 };
 
@@ -50,15 +52,21 @@ use crate::server_impl::TokenStream;
 
 pub const RUSTC_VERSION_STRING: &str = env!("RUSTC_VERSION");
 
-#[derive(Default)]
-pub struct ProcMacroSrv {
+pub struct ProcMacroSrv<'env> {
     expanders: HashMap<(Utf8PathBuf, SystemTime), dylib::Expander>,
     span_mode: SpanMode,
+    env: &'env EnvSnapshot,
+}
+
+impl<'env> ProcMacroSrv<'env> {
+    pub fn new(env: &'env EnvSnapshot) -> Self {
+        Self { expanders: Default::default(), span_mode: Default::default(), env }
+    }
 }
 
 const EXPANDER_STACK_SIZE: usize = 8 * 1024 * 1024;
 
-impl ProcMacroSrv {
+impl<'env> ProcMacroSrv<'env> {
     pub fn set_span_mode(&mut self, span_mode: SpanMode) {
         self.span_mode = span_mode;
     }
@@ -69,52 +77,24 @@ impl ProcMacroSrv {
 
     pub fn expand(
         &mut self,
-        task: msg::ExpandMacro,
+        msg::ExpandMacro { lib, env, current_dir, data }: msg::ExpandMacro,
     ) -> Result<(msg::FlatTree, Vec<u32>), msg::PanicMessage> {
         let span_mode = self.span_mode;
-        let expander = self.expander(task.lib.as_ref()).map_err(|err| {
+        let snapped_env = self.env;
+        let expander = self.expander(lib.as_ref()).map_err(|err| {
             debug_assert!(false, "should list macros before asking to expand");
             msg::PanicMessage(format!("failed to load macro: {err}"))
         })?;
 
-        let prev_env = EnvSnapshot::new();
-        for (k, v) in &task.env {
-            env::set_var(k, v);
-        }
-        let prev_working_dir = match &task.current_dir {
-            Some(dir) => {
-                let prev_working_dir = std::env::current_dir().ok();
-                if let Err(err) = std::env::set_current_dir(dir) {
-                    eprintln!("Failed to set the current working dir to {dir}. Error: {err:?}")
-                }
-                prev_working_dir
-            }
-            None => None,
-        };
-
-        let ExpnGlobals { def_site, call_site, mixed_site, .. } = task.has_global_spans;
+        let prev_env = EnvChange::apply(snapped_env, env, current_dir.as_ref().map(<_>::as_ref));
 
         let result = match span_mode {
-            SpanMode::Id => {
-                expand_id(task, expander, def_site, call_site, mixed_site).map(|it| (it, vec![]))
-            }
-            SpanMode::RustAnalyzer => {
-                expand_ra_span(task, expander, def_site, call_site, mixed_site)
-            }
+            SpanMode::Id => expand_id(data, expander).map(|it| (it, vec![])),
+            SpanMode::RustAnalyzer => expand_ra_span(data, expander),
         };
 
         prev_env.rollback();
 
-        if let Some(dir) = prev_working_dir {
-            if let Err(err) = std::env::set_current_dir(&dir) {
-                eprintln!(
-                    "Failed to set the current working dir to {}. Error: {:?}",
-                    dir.display(),
-                    err
-                )
-            }
-        }
-
         result.map_err(msg::PanicMessage)
     }
 
@@ -168,32 +148,28 @@ impl ProcMacroSrvSpan for Span {
 }
 
 fn expand_id(
-    task: msg::ExpandMacro,
+    msg::ExpandMacroData {
+        macro_body,
+        macro_name,
+        attributes,
+        has_global_spans: ExpnGlobals { serialize: _, def_site, call_site, mixed_site },
+        span_data_table: _,
+    }: msg::ExpandMacroData,
     expander: &dylib::Expander,
-    def_site: usize,
-    call_site: usize,
-    mixed_site: usize,
 ) -> Result<msg::FlatTree, String> {
     let def_site = TokenId(def_site as u32);
     let call_site = TokenId(call_site as u32);
     let mixed_site = TokenId(mixed_site as u32);
 
-    let macro_body = task.macro_body.to_subtree_unresolved(CURRENT_API_VERSION);
-    let attributes = task.attributes.map(|it| it.to_subtree_unresolved(CURRENT_API_VERSION));
+    let macro_body = macro_body.to_subtree_unresolved(CURRENT_API_VERSION);
+    let attributes = attributes.map(|it| it.to_subtree_unresolved(CURRENT_API_VERSION));
     let result = thread::scope(|s| {
         let thread = thread::Builder::new()
             .stack_size(EXPANDER_STACK_SIZE)
-            .name(task.macro_name.clone())
+            .name(macro_name.clone())
             .spawn_scoped(s, || {
                 expander
-                    .expand(
-                        &task.macro_name,
-                        macro_body,
-                        attributes,
-                        def_site,
-                        call_site,
-                        mixed_site,
-                    )
+                    .expand(&macro_name, macro_body, attributes, def_site, call_site, mixed_site)
                     .map(|it| msg::FlatTree::new_raw(&it, CURRENT_API_VERSION))
             });
         let res = match thread {
@@ -210,35 +186,31 @@ fn expand_id(
 }
 
 fn expand_ra_span(
-    task: msg::ExpandMacro,
+    msg::ExpandMacroData {
+        macro_body,
+        macro_name,
+        attributes,
+        has_global_spans: ExpnGlobals { serialize: _, def_site, call_site, mixed_site },
+        span_data_table,
+    }: msg::ExpandMacroData,
     expander: &dylib::Expander,
-    def_site: usize,
-    call_site: usize,
-    mixed_site: usize,
 ) -> Result<(msg::FlatTree, Vec<u32>), String> {
-    let mut span_data_table = deserialize_span_data_index_map(&task.span_data_table);
+    let mut span_data_table = deserialize_span_data_index_map(&span_data_table);
 
     let def_site = span_data_table[def_site];
     let call_site = span_data_table[call_site];
     let mixed_site = span_data_table[mixed_site];
 
-    let macro_body = task.macro_body.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table);
+    let macro_body = macro_body.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table);
     let attributes =
-        task.attributes.map(|it| it.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table));
+        attributes.map(|it| it.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table));
     let result = thread::scope(|s| {
         let thread = thread::Builder::new()
             .stack_size(EXPANDER_STACK_SIZE)
-            .name(task.macro_name.clone())
+            .name(macro_name.clone())
             .spawn_scoped(s, || {
                 expander
-                    .expand(
-                        &task.macro_name,
-                        macro_body,
-                        attributes,
-                        def_site,
-                        call_site,
-                        mixed_site,
-                    )
+                    .expand(&macro_name, macro_body, attributes, def_site, call_site, mixed_site)
                     .map(|it| {
                         (
                             msg::FlatTree::new(&it, CURRENT_API_VERSION, &mut span_data_table),
@@ -269,31 +241,74 @@ impl PanicMessage {
     }
 }
 
-struct EnvSnapshot {
+pub struct EnvSnapshot {
     vars: HashMap<OsString, OsString>,
 }
 
 impl EnvSnapshot {
-    fn new() -> EnvSnapshot {
+    pub fn new() -> EnvSnapshot {
         EnvSnapshot { vars: env::vars_os().collect() }
     }
+}
+
+struct EnvChange<'snap> {
+    changed_vars: Vec<String>,
+    prev_working_dir: Option<PathBuf>,
+    snap: &'snap EnvSnapshot,
+}
+
+impl<'snap> EnvChange<'snap> {
+    fn apply(
+        snap: &'snap EnvSnapshot,
+        new_vars: Vec<(String, String)>,
+        current_dir: Option<&Path>,
+    ) -> EnvChange<'snap> {
+        let prev_working_dir = match current_dir {
+            Some(dir) => {
+                let prev_working_dir = std::env::current_dir().ok();
+                if let Err(err) = std::env::set_current_dir(dir) {
+                    eprintln!(
+                        "Failed to set the current working dir to {}. Error: {err:?}",
+                        dir.display()
+                    )
+                }
+                prev_working_dir
+            }
+            None => None,
+        };
+        EnvChange {
+            snap,
+            changed_vars: new_vars
+                .into_iter()
+                .map(|(k, v)| {
+                    env::set_var(&k, v);
+                    k
+                })
+                .collect(),
+            prev_working_dir,
+        }
+    }
 
     fn rollback(self) {}
 }
 
-impl Drop for EnvSnapshot {
+impl Drop for EnvChange<'_> {
     fn drop(&mut self) {
-        for (name, value) in env::vars_os() {
-            let old_value = self.vars.remove(&name);
-            if old_value != Some(value) {
-                match old_value {
-                    None => env::remove_var(name),
-                    Some(old_value) => env::set_var(name, old_value),
-                }
+        for name in self.changed_vars.drain(..) {
+            match self.snap.vars.get::<std::ffi::OsStr>(name.as_ref()) {
+                Some(prev_val) => env::set_var(name, prev_val),
+                None => env::remove_var(name),
             }
         }
-        for (name, old_value) in self.vars.drain() {
-            env::set_var(name, old_value)
+
+        if let Some(dir) = &self.prev_working_dir {
+            if let Err(err) = std::env::set_current_dir(&dir) {
+                eprintln!(
+                    "Failed to set the current working dir to {}. Error: {:?}",
+                    dir.display(),
+                    err
+                )
+            }
         }
     }
 }
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/proc_macros.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/proc_macros.rs
index c5c3dff9db8..d48c5b30dee 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/proc_macros.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/proc_macros.rs
@@ -1,9 +1,10 @@
 //! Proc macro ABI
 
-use libloading::Library;
 use proc_macro::bridge;
 use proc_macro_api::ProcMacroKind;
 
+use libloading::Library;
+
 use crate::{dylib::LoadProcMacroDylibError, ProcMacroSrvSpan};
 
 pub(crate) struct ProcMacros {
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/tests/utils.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/tests/utils.rs
index 0579c426d09..03b1117a5bd 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/tests/utils.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/tests/utils.rs
@@ -5,7 +5,7 @@ use proc_macro_api::msg::TokenId;
 use span::{ErasedFileAstId, FileId, Span, SpanAnchor, SyntaxContextId};
 use tt::TextRange;
 
-use crate::{dylib, proc_macro_test_dylib_path, ProcMacroSrv};
+use crate::{dylib, proc_macro_test_dylib_path, EnvSnapshot, ProcMacroSrv};
 
 fn parse_string(call_site: TokenId, src: &str) -> crate::server_impl::TokenStream<TokenId> {
     crate::server_impl::TokenStream::with_subtree(
@@ -96,7 +96,8 @@ fn assert_expand_impl(
 
 pub(crate) fn list() -> Vec<String> {
     let dylib_path = proc_macro_test_dylib_path();
-    let mut srv = ProcMacroSrv::default();
+    let env = EnvSnapshot::new();
+    let mut srv = ProcMacroSrv::new(&env);
     let res = srv.list_macros(&dylib_path).unwrap();
     res.into_iter().map(|(name, kind)| format!("{name} [{kind:?}]")).collect()
 }