Auto merge of #127617 - lnicola:sync-from-ra, r=lnicola

Subtree update of `rust-analyzer`

r? `@ghost`

15 files changed, 366 insertions, 183 deletions

diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/Cargo.toml b/src/tools/rust-analyzer/crates/proc-macro-srv/Cargo.toml
index f8db1c6a30b..735f781c439 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/Cargo.toml
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/Cargo.toml
@@ -13,8 +13,9 @@ doctest = false
 
 [dependencies]
 object.workspace = true
-libloading = "0.8.0"
-memmap2 = "0.5.4"
+libloading.workspace = true
+memmap2.workspace = true
+snap.workspace = true
 
 stdx.workspace = true
 tt.workspace = true
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/build.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/build.rs
index 874d1c6cd38..9a17cfc9f36 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/build.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/build.rs
@@ -1,27 +1,15 @@
 //! Determine rustc version `proc-macro-srv` (and thus the sysroot ABI) is
 //! build with and make it accessible at runtime for ABI selection.
 
-use std::{env, fs::File, io::Write, path::PathBuf, process::Command};
+use std::{env, process::Command};
 
 fn main() {
-    println!("cargo:rustc-check-cfg=cfg(rust_analyzer)");
-
-    let mut path = PathBuf::from(env::var_os("OUT_DIR").unwrap());
-    path.push("rustc_version.rs");
-    let mut f = File::create(&path).unwrap();
+    println!("cargo::rustc-check-cfg=cfg(rust_analyzer)");
 
     let rustc = env::var("RUSTC").expect("proc-macro-srv's build script expects RUSTC to be set");
     let output = Command::new(rustc).arg("--version").output().expect("rustc --version must run");
     let version_string = std::str::from_utf8(&output.stdout[..])
         .expect("rustc --version output must be UTF-8")
         .trim();
-
-    write!(
-        f,
-        "
-    #[allow(dead_code)]
-    pub(crate) const RUSTC_VERSION_STRING: &str = {version_string:?};
-    "
-    )
-    .unwrap();
+    println!("cargo::rustc-env=RUSTC_VERSION={}", version_string);
 }
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/proc-macro-test/build.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/proc-macro-test/build.rs
index 6a0ae362d88..ff2f5d18639 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/proc-macro-test/build.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/proc-macro-test/build.rs
@@ -8,7 +8,7 @@
 //! 1.58) and future ABIs (stage1, nightly)
 
 use std::{
-    env, fs,
+    env,
     path::{Path, PathBuf},
     process::Command,
 };
@@ -30,8 +30,7 @@ fn main() {
 
     if !has_features {
         println!("proc-macro-test testing only works on nightly toolchains");
-        let info_path = out_dir.join("proc_macro_test_location.txt");
-        fs::File::create(info_path).unwrap();
+        println!("cargo::rustc-env=PROC_MACRO_TEST_LOCATION=\"\"");
         return;
     }
 
@@ -121,6 +120,5 @@ fn main() {
     // This file is under `target_dir` and is already under `OUT_DIR`.
     let artifact_path = artifact_path.expect("no dylib for proc-macro-test-impl found");
 
-    let info_path = out_dir.join("proc_macro_test_location.txt");
-    fs::write(info_path, artifact_path.to_str().unwrap()).unwrap();
+    println!("cargo::rustc-env=PROC_MACRO_TEST_LOCATION={}", artifact_path.display());
 }
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/proc-macro-test/imp/src/lib.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/proc-macro-test/imp/src/lib.rs
index 5f8530d08c4..a1707364f3c 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/proc-macro-test/imp/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/proc-macro-test/imp/src/lib.rs
@@ -1,6 +1,6 @@
 //! Exports a few trivial procedural macros for testing.
 
-#![warn(rust_2018_idioms, unused_lifetimes)]
+
 #![feature(proc_macro_span, proc_macro_def_site)]
 #![allow(clippy::all)]
 
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/proc-macro-test/src/lib.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/proc-macro-test/src/lib.rs
index 739c6ec6f44..6464adb2ca7 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/proc-macro-test/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/proc-macro-test/src/lib.rs
@@ -1,6 +1,3 @@
 //! Exports a few trivial procedural macros for testing.
 
-#![warn(rust_2018_idioms, unused_lifetimes)]
-
-pub static PROC_MACRO_TEST_LOCATION: &str =
-    include_str!(concat!(env!("OUT_DIR"), "/proc_macro_test_location.txt"));
+pub static PROC_MACRO_TEST_LOCATION: &str = env!("PROC_MACRO_TEST_LOCATION");
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib.rs
index 22c34ff1678..78ae4574c40 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib.rs
@@ -1,13 +1,15 @@
 //! Handles dynamic library loading for proc macro
 
+mod version;
+
+use proc_macro::bridge;
 use std::{fmt, fs::File, io};
 
 use libloading::Library;
 use memmap2::Mmap;
 use object::Object;
 use paths::{AbsPath, Utf8Path, Utf8PathBuf};
-use proc_macro::bridge;
-use proc_macro_api::{read_dylib_info, ProcMacroKind};
+use proc_macro_api::ProcMacroKind;
 
 use crate::ProcMacroSrvSpan;
 
@@ -119,33 +121,45 @@ impl ProcMacroLibraryLibloading {
         let abs_file: &AbsPath = file
             .try_into()
             .map_err(|_| invalid_data_err(format!("expected an absolute path, got {file}")))?;
-        let version_info = read_dylib_info(abs_file)?;
+        let version_info = version::read_dylib_info(abs_file)?;
 
         let lib = load_library(file).map_err(invalid_data_err)?;
-        let proc_macros =
-            crate::proc_macros::ProcMacros::from_lib(&lib, symbol_name, version_info)?;
+        let proc_macros = crate::proc_macros::ProcMacros::from_lib(
+            &lib,
+            symbol_name,
+            &version_info.version_string,
+        )?;
         Ok(ProcMacroLibraryLibloading { _lib: lib, proc_macros })
     }
 }
 
-pub struct Expander {
+pub(crate) struct Expander {
     inner: ProcMacroLibraryLibloading,
+    path: Utf8PathBuf,
+}
+
+impl Drop for Expander {
+    fn drop(&mut self) {
+        #[cfg(windows)]
+        std::fs::remove_file(&self.path).ok();
+        _ = self.path;
+    }
 }
 
 impl Expander {
-    pub fn new(lib: &Utf8Path) -> Result<Expander, LoadProcMacroDylibError> {
+    pub(crate) fn new(lib: &Utf8Path) -> Result<Expander, LoadProcMacroDylibError> {
         // Some libraries for dynamic loading require canonicalized path even when it is
         // already absolute
         let lib = lib.canonicalize_utf8()?;
 
-        let lib = ensure_file_with_lock_free_access(&lib)?;
+        let path = ensure_file_with_lock_free_access(&lib)?;
 
-        let library = ProcMacroLibraryLibloading::open(lib.as_ref())?;
+        let library = ProcMacroLibraryLibloading::open(path.as_ref())?;
 
-        Ok(Expander { inner: library })
+        Ok(Expander { inner: library, path })
     }
 
-    pub fn expand<S: ProcMacroSrvSpan>(
+    pub(crate) fn expand<S: ProcMacroSrvSpan>(
         &self,
         macro_name: &str,
         macro_body: tt::Subtree<S>,
@@ -164,7 +178,7 @@ impl Expander {
         result.map_err(|e| e.into_string().unwrap_or_default())
     }
 
-    pub fn list_macros(&self) -> Vec<(String, ProcMacroKind)> {
+    pub(crate) fn list_macros(&self) -> Vec<(String, ProcMacroKind)> {
         self.inner.proc_macros.list_macros()
     }
 }
@@ -193,7 +207,7 @@ fn ensure_file_with_lock_free_access(path: &Utf8Path) -> io::Result<Utf8PathBuf>
     unique_name.push_str(file_name);
 
     to.push(unique_name);
-    std::fs::copy(path, &to).unwrap();
+    std::fs::copy(path, &to)?;
     Ok(to)
 }
 
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib/version.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib/version.rs
new file mode 100644
index 00000000000..1f7ef7914ba
--- /dev/null
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib/version.rs
@@ -0,0 +1,180 @@
+//! Reading proc-macro rustc version information from binary data
+
+use std::{
+    fs::File,
+    io::{self, Read},
+};
+
+use memmap2::Mmap;
+use object::read::{File as BinaryFile, Object, ObjectSection};
+use paths::AbsPath;
+use snap::read::FrameDecoder as SnapDecoder;
+
+#[derive(Debug)]
+#[allow(dead_code)]
+pub struct RustCInfo {
+    pub version: (usize, usize, usize),
+    pub channel: String,
+    pub commit: Option<String>,
+    pub date: Option<String>,
+    // something like "rustc 1.58.1 (db9d1b20b 2022-01-20)"
+    pub version_string: String,
+}
+
+/// Read rustc dylib information
+pub fn read_dylib_info(dylib_path: &AbsPath) -> io::Result<RustCInfo> {
+    macro_rules! err {
+        ($e:literal) => {
+            io::Error::new(io::ErrorKind::InvalidData, $e)
+        };
+    }
+
+    let ver_str = read_version(dylib_path)?;
+    let mut items = ver_str.split_whitespace();
+    let tag = items.next().ok_or_else(|| err!("version format error"))?;
+    if tag != "rustc" {
+        return Err(err!("version format error (No rustc tag)"));
+    }
+
+    let version_part = items.next().ok_or_else(|| err!("no version string"))?;
+    let mut version_parts = version_part.split('-');
+    let version = version_parts.next().ok_or_else(|| err!("no version"))?;
+    let channel = version_parts.next().unwrap_or_default().to_owned();
+
+    let commit = match items.next() {
+        Some(commit) => {
+            match commit.len() {
+                0 => None,
+                _ => Some(commit[1..].to_string() /* remove ( */),
+            }
+        }
+        None => None,
+    };
+    let date = match items.next() {
+        Some(date) => {
+            match date.len() {
+                0 => None,
+                _ => Some(date[0..date.len() - 2].to_string() /* remove ) */),
+            }
+        }
+        None => None,
+    };
+
+    let version_numbers = version
+        .split('.')
+        .map(|it| it.parse::<usize>())
+        .collect::<Result<Vec<_>, _>>()
+        .map_err(|_| err!("version number error"))?;
+
+    if version_numbers.len() != 3 {
+        return Err(err!("version number format error"));
+    }
+    let version = (version_numbers[0], version_numbers[1], version_numbers[2]);
+
+    Ok(RustCInfo { version, channel, commit, date, version_string: ver_str })
+}
+
+/// This is used inside read_version() to locate the ".rustc" section
+/// from a proc macro crate's binary file.
+fn read_section<'a>(dylib_binary: &'a [u8], section_name: &str) -> io::Result<&'a [u8]> {
+    BinaryFile::parse(dylib_binary)
+        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
+        .section_by_name(section_name)
+        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "section read error"))?
+        .data()
+        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
+}
+
+/// Check the version of rustc that was used to compile a proc macro crate's
+/// binary file.
+///
+/// A proc macro crate binary's ".rustc" section has following byte layout:
+/// * [b'r',b'u',b's',b't',0,0,0,5] is the first 8 bytes
+/// * ff060000 734e6150 is followed, it's the snappy format magic bytes,
+///   means bytes from here(including this sequence) are compressed in
+///   snappy compression format. Version info is inside here, so decompress
+///   this.
+///
+/// The bytes you get after decompressing the snappy format portion has
+/// following layout:
+/// * [b'r',b'u',b's',b't',0,0,0,5] is the first 8 bytes(again)
+/// * [crate root bytes] next 8 bytes (4 in old versions) is to store
+///   crate root position, according to rustc's source code comment
+/// * [length byte] next 1 byte tells us how many bytes we should read next
+///   for the version string's utf8 bytes
+/// * [version string bytes encoded in utf8] <- GET THIS BOI
+/// * [some more bytes that we don't really care but about still there] :-)
+///
+/// Check this issue for more about the bytes layout:
+/// <https://github.com/rust-lang/rust-analyzer/issues/6174>
+pub fn read_version(dylib_path: &AbsPath) -> io::Result<String> {
+    let dylib_file = File::open(dylib_path)?;
+    let dylib_mmapped = unsafe { Mmap::map(&dylib_file) }?;
+
+    let dot_rustc = read_section(&dylib_mmapped, ".rustc")?;
+
+    // check if magic is valid
+    if &dot_rustc[0..4] != b"rust" {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            format!("unknown metadata magic, expected `rust`, found `{:?}`", &dot_rustc[0..4]),
+        ));
+    }
+    let version = u32::from_be_bytes([dot_rustc[4], dot_rustc[5], dot_rustc[6], dot_rustc[7]]);
+    // Last supported version is:
+    // https://github.com/rust-lang/rust/commit/b94cfefc860715fb2adf72a6955423d384c69318
+    let (snappy_portion, bytes_before_version) = match version {
+        5 | 6 => (&dot_rustc[8..], 13),
+        7 | 8 => {
+            let len_bytes = &dot_rustc[8..12];
+            let data_len = u32::from_be_bytes(len_bytes.try_into().unwrap()) as usize;
+            (&dot_rustc[12..data_len + 12], 13)
+        }
+        9 => {
+            let len_bytes = &dot_rustc[8..16];
+            let data_len = u64::from_le_bytes(len_bytes.try_into().unwrap()) as usize;
+            (&dot_rustc[16..data_len + 12], 17)
+        }
+        _ => {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("unsupported metadata version {version}"),
+            ));
+        }
+    };
+
+    let mut uncompressed: Box<dyn Read> = if &snappy_portion[0..4] == b"rust" {
+        // Not compressed.
+        Box::new(snappy_portion)
+    } else {
+        Box::new(SnapDecoder::new(snappy_portion))
+    };
+
+    // We're going to skip over the bytes before the version string, so basically:
+    // 8 bytes for [b'r',b'u',b's',b't',0,0,0,5]
+    // 4 or 8 bytes for [crate root bytes]
+    // 1 byte for length of version string
+    // so 13 or 17 bytes in total, and we should check the last of those bytes
+    // to know the length
+    let mut bytes = [0u8; 17];
+    uncompressed.read_exact(&mut bytes[..bytes_before_version])?;
+    let length = bytes[bytes_before_version - 1];
+
+    let mut version_string_utf8 = vec![0u8; length as usize];
+    uncompressed.read_exact(&mut version_string_utf8)?;
+    let version_string = String::from_utf8(version_string_utf8);
+    version_string.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
+}
+
+#[test]
+fn test_version_check() {
+    let path = paths::AbsPathBuf::assert(crate::proc_macro_test_dylib_path());
+    let info = read_dylib_info(&path).unwrap();
+    assert_eq!(
+        info.version_string,
+        crate::RUSTC_VERSION_STRING,
+        "sysroot ABI mismatch: dylib rustc version (read from .rustc section): {:?} != proc-macro-srv version (read from 'rustc --version'): {:?}",
+        info.version_string,
+        crate::RUSTC_VERSION_STRING,
+    );
+}
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 2472c1e3119..e6281035e1a 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
@@ -13,7 +13,6 @@
 #![cfg(any(feature = "sysroot-abi", rust_analyzer))]
 #![cfg_attr(feature = "in-rust-tree", feature(rustc_private))]
 #![feature(proc_macro_internals, proc_macro_diagnostic, proc_macro_span)]
-#![warn(rust_2018_idioms, unused_lifetimes)]
 #![allow(unreachable_pub, internal_features)]
 
 extern crate proc_macro;
@@ -27,13 +26,15 @@ extern crate rustc_lexer;
 
 mod dylib;
 mod proc_macros;
-mod server;
+mod server_impl;
 
 use std::{
     collections::{hash_map::Entry, HashMap},
     env,
     ffi::OsString,
-    fs, thread,
+    fs,
+    path::{Path, PathBuf},
+    thread,
     time::SystemTime,
 };
 
@@ -47,46 +48,25 @@ use proc_macro_api::{
 };
 use span::Span;
 
-use crate::server::TokenStream;
+use crate::server_impl::TokenStream;
 
-// see `build.rs`
-include!(concat!(env!("OUT_DIR"), "/rustc_version.rs"));
+pub const RUSTC_VERSION_STRING: &str = env!("RUSTC_VERSION");
 
-trait ProcMacroSrvSpan: Copy {
-    type Server: proc_macro::bridge::server::Server<TokenStream = TokenStream<Self>>;
-    fn make_server(call_site: Self, def_site: Self, mixed_site: Self) -> Self::Server;
+pub struct ProcMacroSrv<'env> {
+    expanders: HashMap<(Utf8PathBuf, SystemTime), dylib::Expander>,
+    span_mode: SpanMode,
+    env: &'env EnvSnapshot,
 }
 
-impl ProcMacroSrvSpan for TokenId {
-    type Server = server::token_id::TokenIdServer;
-
-    fn make_server(call_site: Self, def_site: Self, mixed_site: Self) -> Self::Server {
-        Self::Server { interner: &server::SYMBOL_INTERNER, call_site, def_site, mixed_site }
+impl<'env> ProcMacroSrv<'env> {
+    pub fn new(env: &'env EnvSnapshot) -> Self {
+        Self { expanders: Default::default(), span_mode: Default::default(), env }
     }
 }
-impl ProcMacroSrvSpan for Span {
-    type Server = server::rust_analyzer_span::RaSpanServer;
-    fn make_server(call_site: Self, def_site: Self, mixed_site: Self) -> Self::Server {
-        Self::Server {
-            interner: &server::SYMBOL_INTERNER,
-            call_site,
-            def_site,
-            mixed_site,
-            tracked_env_vars: Default::default(),
-            tracked_paths: Default::default(),
-        }
-    }
-}
-
-#[derive(Default)]
-pub struct ProcMacroSrv {
-    expanders: HashMap<(Utf8PathBuf, SystemTime), dylib::Expander>,
-    span_mode: SpanMode,
-}
 
 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;
     }
@@ -97,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)
     }
 
@@ -169,33 +121,55 @@ impl ProcMacroSrv {
     }
 }
 
+trait ProcMacroSrvSpan: Copy {
+    type Server: proc_macro::bridge::server::Server<TokenStream = TokenStream<Self>>;
+    fn make_server(call_site: Self, def_site: Self, mixed_site: Self) -> Self::Server;
+}
+
+impl ProcMacroSrvSpan for TokenId {
+    type Server = server_impl::token_id::TokenIdServer;
+
+    fn make_server(call_site: Self, def_site: Self, mixed_site: Self) -> Self::Server {
+        Self::Server { interner: &server_impl::SYMBOL_INTERNER, call_site, def_site, mixed_site }
+    }
+}
+impl ProcMacroSrvSpan for Span {
+    type Server = server_impl::rust_analyzer_span::RaSpanServer;
+    fn make_server(call_site: Self, def_site: Self, mixed_site: Self) -> Self::Server {
+        Self::Server {
+            interner: &server_impl::SYMBOL_INTERNER,
+            call_site,
+            def_site,
+            mixed_site,
+            tracked_env_vars: Default::default(),
+            tracked_paths: Default::default(),
+        }
+    }
+}
+
 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 {
@@ -212,35 +186,33 @@ 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));
+    // Note, we spawn a new thread here so that thread locals allocation don't accumulate (this
+    // includes the proc-macro symbol interner)
     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),
@@ -271,31 +243,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 631fd84aa24..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,8 +1,9 @@
 //! Proc macro ABI
 
-use libloading::Library;
 use proc_macro::bridge;
-use proc_macro_api::{ProcMacroKind, RustCInfo};
+use proc_macro_api::ProcMacroKind;
+
+use libloading::Library;
 
 use crate::{dylib::LoadProcMacroDylibError, ProcMacroSrvSpan};
 
@@ -29,15 +30,15 @@ impl ProcMacros {
     pub(crate) fn from_lib(
         lib: &Library,
         symbol_name: String,
-        info: RustCInfo,
+        version_string: &str,
     ) -> Result<ProcMacros, LoadProcMacroDylibError> {
-        if info.version_string == crate::RUSTC_VERSION_STRING {
+        if version_string == crate::RUSTC_VERSION_STRING {
             let macros =
                 unsafe { lib.get::<&&[bridge::client::ProcMacro]>(symbol_name.as_bytes()) }?;
 
             return Ok(Self { exported_macros: macros.to_vec() });
         }
-        Err(LoadProcMacroDylibError::AbiMismatch(info.version_string))
+        Err(LoadProcMacroDylibError::AbiMismatch(version_string.to_owned()))
     }
 
     pub(crate) fn expand<S: ProcMacroSrvSpan>(
@@ -49,11 +50,12 @@ impl ProcMacros {
         call_site: S,
         mixed_site: S,
     ) -> Result<tt::Subtree<S>, crate::PanicMessage> {
-        let parsed_body = crate::server::TokenStream::with_subtree(macro_body);
+        let parsed_body = crate::server_impl::TokenStream::with_subtree(macro_body);
 
-        let parsed_attributes = attributes.map_or_else(crate::server::TokenStream::new, |attr| {
-            crate::server::TokenStream::with_subtree(attr)
-        });
+        let parsed_attributes = attributes
+            .map_or_else(crate::server_impl::TokenStream::new, |attr| {
+                crate::server_impl::TokenStream::with_subtree(attr)
+            });
 
         for proc_macro in &self.exported_macros {
             match proc_macro {
@@ -117,16 +119,3 @@ impl ProcMacros {
             .collect()
     }
 }
-
-#[test]
-fn test_version_check() {
-    let path = paths::AbsPathBuf::assert(crate::proc_macro_test_dylib_path());
-    let info = proc_macro_api::read_dylib_info(&path).unwrap();
-    assert_eq!(
-        info.version_string,
-        crate::RUSTC_VERSION_STRING,
-        "sysroot ABI mismatch: dylib rustc version (read from .rustc section): {:?} != proc-macro-srv version (read from 'rustc --version'): {:?}",
-        info.version_string,
-        crate::RUSTC_VERSION_STRING,
-    );
-}
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl.rs
index e8b340a43d3..e8b340a43d3 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl.rs
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server/rust_analyzer_span.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/rust_analyzer_span.rs
index 0350bde4122..bb174ba1b22 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server/rust_analyzer_span.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/rust_analyzer_span.rs
@@ -14,7 +14,7 @@ use proc_macro::bridge::{self, server};
 use span::{Span, FIXUP_ERASED_FILE_AST_ID_MARKER};
 use tt::{TextRange, TextSize};
 
-use crate::server::{
+use crate::server_impl::{
     delim_to_external, delim_to_internal, literal_with_stringify_parts,
     token_stream::TokenStreamBuilder, Symbol, SymbolInternerRef, SYMBOL_INTERNER,
 };
@@ -29,7 +29,7 @@ mod tt {
     pub type Ident = ::tt::Ident<super::Span>;
 }
 
-type TokenStream = crate::server::TokenStream<Span>;
+type TokenStream = crate::server_impl::TokenStream<Span>;
 
 #[derive(Clone)]
 pub struct SourceFile;
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server/symbol.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/symbol.rs
index 540d06457f2..540d06457f2 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server/symbol.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/symbol.rs
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server/token_id.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/token_id.rs
index ad7bd954cf1..12edacbe39d 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server/token_id.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/token_id.rs
@@ -7,7 +7,7 @@ use std::{
 
 use proc_macro::bridge::{self, server};
 
-use crate::server::{
+use crate::server_impl::{
     delim_to_external, delim_to_internal, literal_with_stringify_parts,
     token_stream::TokenStreamBuilder, Symbol, SymbolInternerRef, SYMBOL_INTERNER,
 };
@@ -31,7 +31,7 @@ type Spacing = tt::Spacing;
 #[allow(unused)]
 type Literal = tt::Literal;
 type Span = tt::TokenId;
-type TokenStream = crate::server::TokenStream<Span>;
+type TokenStream = crate::server_impl::TokenStream<Span>;
 
 #[derive(Clone)]
 pub struct SourceFile;
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server/token_stream.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/token_stream.rs
index b1a448427c6..b1a448427c6 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server/token_stream.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/token_stream.rs
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 6050bc9e36e..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,10 +5,10 @@ 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::TokenStream<TokenId> {
-    crate::server::TokenStream::with_subtree(
+fn parse_string(call_site: TokenId, src: &str) -> crate::server_impl::TokenStream<TokenId> {
+    crate::server_impl::TokenStream::with_subtree(
         mbe::parse_to_token_tree_static_span(call_site, src).unwrap(),
     )
 }
@@ -17,8 +17,8 @@ fn parse_string_spanned(
     anchor: SpanAnchor,
     call_site: SyntaxContextId,
     src: &str,
-) -> crate::server::TokenStream<Span> {
-    crate::server::TokenStream::with_subtree(
+) -> crate::server_impl::TokenStream<Span> {
+    crate::server_impl::TokenStream::with_subtree(
         mbe::parse_to_token_tree(anchor, call_site, src).unwrap(),
     )
 }
@@ -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()
 }