Decouple proc-macro server protocol from the server implementation

17 files changed, 243 insertions, 218 deletions

diff --git a/src/tools/rust-analyzer/Cargo.lock b/src/tools/rust-analyzer/Cargo.lock
index 2323fdf5333..10a95562cd9 100644
--- a/src/tools/rust-analyzer/Cargo.lock
+++ b/src/tools/rust-analyzer/Cargo.lock
@@ -1375,7 +1375,6 @@ dependencies = [
  "memmap2",
  "object 0.33.0",
  "paths",
- "proc-macro-api",
  "proc-macro-test",
  "ra-ap-rustc_lexer",
  "span",
diff --git a/src/tools/rust-analyzer/crates/load-cargo/src/lib.rs b/src/tools/rust-analyzer/crates/load-cargo/src/lib.rs
index 1b2162dad0f..98cc4967e78 100644
--- a/src/tools/rust-analyzer/crates/load-cargo/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/load-cargo/src/lib.rs
@@ -14,7 +14,7 @@ use ide_db::{
     prime_caches, ChangeWithProcMacros, FxHashMap, RootDatabase,
 };
 use itertools::Itertools;
-use proc_macro_api::{MacroDylib, ProcMacroServer};
+use proc_macro_api::{MacroDylib, ProcMacroClient};
 use project_model::{CargoConfig, PackageRoot, ProjectManifest, ProjectWorkspace};
 use span::Span;
 use vfs::{
@@ -42,7 +42,7 @@ pub fn load_workspace_at(
     cargo_config: &CargoConfig,
     load_config: &LoadCargoConfig,
     progress: &dyn Fn(String),
-) -> anyhow::Result<(RootDatabase, vfs::Vfs, Option<ProcMacroServer>)> {
+) -> anyhow::Result<(RootDatabase, vfs::Vfs, Option<ProcMacroClient>)> {
     let root = AbsPathBuf::assert_utf8(std::env::current_dir()?.join(root));
     let root = ProjectManifest::discover_single(&root)?;
     let mut workspace = ProjectWorkspace::load(root, cargo_config, progress)?;
@@ -59,7 +59,7 @@ pub fn load_workspace(
     ws: ProjectWorkspace,
     extra_env: &FxHashMap<String, String>,
     load_config: &LoadCargoConfig,
-) -> anyhow::Result<(RootDatabase, vfs::Vfs, Option<ProcMacroServer>)> {
+) -> anyhow::Result<(RootDatabase, vfs::Vfs, Option<ProcMacroClient>)> {
     let (sender, receiver) = unbounded();
     let mut vfs = vfs::Vfs::default();
     let mut loader = {
@@ -71,10 +71,10 @@ pub fn load_workspace(
     let proc_macro_server = match &load_config.with_proc_macro_server {
         ProcMacroServerChoice::Sysroot => ws
             .find_sysroot_proc_macro_srv()
-            .and_then(|it| ProcMacroServer::spawn(&it, extra_env).map_err(Into::into))
+            .and_then(|it| ProcMacroClient::spawn(&it, extra_env).map_err(Into::into))
             .map_err(|e| (e, true)),
         ProcMacroServerChoice::Explicit(path) => {
-            ProcMacroServer::spawn(path, extra_env).map_err(Into::into).map_err(|e| (e, true))
+            ProcMacroClient::spawn(path, extra_env).map_err(Into::into).map_err(|e| (e, true))
         }
         ProcMacroServerChoice::None => {
             Err((anyhow::format_err!("proc macro server disabled"), false))
@@ -362,7 +362,7 @@ impl SourceRootConfig {
 
 /// Load the proc-macros for the given lib path, disabling all expanders whose names are in `ignored_macros`.
 pub fn load_proc_macro(
-    server: &ProcMacroServer,
+    server: &ProcMacroClient,
     path: &AbsPath,
     ignored_macros: &[Box<str>],
 ) -> ProcMacroLoadResult {
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 e54d501b94c..bd33cbdae9d 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
@@ -13,18 +13,16 @@ use paths::{AbsPath, AbsPathBuf};
 use span::Span;
 use std::{fmt, io, sync::Arc};
 
-use serde::{Deserialize, Serialize};
-
 use crate::{
     msg::{
         deserialize_span_data_index_map, flat::serialize_span_data_index_map, ExpandMacro,
         ExpnGlobals, FlatTree, PanicMessage, SpanDataIndexMap, HAS_GLOBAL_SPANS,
         RUST_ANALYZER_SPAN_SUPPORT,
     },
-    process::ProcMacroProcessSrv,
+    process::ProcMacroServerProcess,
 };
 
-#[derive(Copy, Clone, Eq, PartialEq, Debug, Serialize, Deserialize)]
+#[derive(Copy, Clone, Eq, PartialEq, Debug, serde_derive::Serialize, serde_derive::Deserialize)]
 pub enum ProcMacroKind {
     CustomDerive,
     Attr,
@@ -37,12 +35,12 @@ pub enum ProcMacroKind {
 /// A handle to an external process which load dylibs with macros (.so or .dll)
 /// and runs actual macro expansion functions.
 #[derive(Debug)]
-pub struct ProcMacroServer {
+pub struct ProcMacroClient {
     /// Currently, the proc macro process expands all procedural macros sequentially.
     ///
     /// That means that concurrent salsa requests may block each other when expanding proc macros,
     /// which is unfortunate, but simple and good enough for the time being.
-    process: Arc<ProcMacroProcessSrv>,
+    process: Arc<ProcMacroServerProcess>,
     path: AbsPathBuf,
 }
 
@@ -62,7 +60,7 @@ impl MacroDylib {
 /// we share a single expander process for all macros.
 #[derive(Debug, Clone)]
 pub struct ProcMacro {
-    process: Arc<ProcMacroProcessSrv>,
+    process: Arc<ProcMacroServerProcess>,
     dylib_path: Arc<AbsPathBuf>,
     name: Box<str>,
     kind: ProcMacroKind,
@@ -95,15 +93,15 @@ impl fmt::Display for ServerError {
     }
 }
 
-impl ProcMacroServer {
+impl ProcMacroClient {
     /// Spawns an external process as the proc macro server and returns a client connected to it.
     pub fn spawn(
         process_path: &AbsPath,
         env: impl IntoIterator<Item = (impl AsRef<std::ffi::OsStr>, impl AsRef<std::ffi::OsStr>)>
             + Clone,
-    ) -> io::Result<ProcMacroServer> {
-        let process = ProcMacroProcessSrv::run(process_path, env)?;
-        Ok(ProcMacroServer { process: Arc::new(process), path: process_path.to_owned() })
+    ) -> io::Result<ProcMacroClient> {
+        let process = ProcMacroServerProcess::run(process_path, env)?;
+        Ok(ProcMacroClient { process: Arc::new(process), path: process_path.to_owned() })
     }
 
     pub fn path(&self) -> &AbsPath {
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 bbd9f582df9..2405a5bfb28 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
@@ -11,8 +11,8 @@ use crate::ProcMacroKind;
 
 pub use crate::msg::flat::{
     deserialize_span_data_index_map, serialize_span_data_index_map, FlatTree, SpanDataIndexMap,
-    TokenId,
 };
+pub use span::TokenId;
 
 // The versions of the server protocol
 pub const NO_VERSION_CHECK_VERSION: u32 = 0;
diff --git a/src/tools/rust-analyzer/crates/proc-macro-api/src/msg/flat.rs b/src/tools/rust-analyzer/crates/proc-macro-api/src/msg/flat.rs
index ce4b060fca5..39d3d9c4e5f 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-api/src/msg/flat.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-api/src/msg/flat.rs
@@ -40,7 +40,9 @@ use std::collections::VecDeque;
 use intern::Symbol;
 use rustc_hash::FxHashMap;
 use serde_derive::{Deserialize, Serialize};
-use span::{EditionedFileId, ErasedFileAstId, Span, SpanAnchor, SyntaxContextId, TextRange};
+use span::{
+    EditionedFileId, ErasedFileAstId, Span, SpanAnchor, SyntaxContextId, TextRange, TokenId,
+};
 
 use crate::msg::{ENCODE_CLOSE_SPAN_VERSION, EXTENDED_LEAF_DATA};
 
@@ -78,15 +80,6 @@ pub fn deserialize_span_data_index_map(map: &[u32]) -> SpanDataIndexMap {
         .collect()
 }
 
-#[derive(Clone, Copy, PartialEq, Eq, Hash)]
-pub struct TokenId(pub u32);
-
-impl std::fmt::Debug for TokenId {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        self.0.fmt(f)
-    }
-}
-
 #[derive(Serialize, Deserialize, Debug)]
 pub struct FlatTree {
     subtree: Vec<u32>,
diff --git a/src/tools/rust-analyzer/crates/proc-macro-api/src/process.rs b/src/tools/rust-analyzer/crates/proc-macro-api/src/process.rs
index 4d62efdd6b1..ed1547d5781 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-api/src/process.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-api/src/process.rs
@@ -17,7 +17,7 @@ use crate::{
 };
 
 #[derive(Debug)]
-pub(crate) struct ProcMacroProcessSrv {
+pub(crate) struct ProcMacroServerProcess {
     /// The state of the proc-macro server process, the protocol is currently strictly sequential
     /// hence the lock on the state.
     state: Mutex<ProcessSrvState>,
@@ -34,17 +34,17 @@ struct ProcessSrvState {
     stdout: BufReader<ChildStdout>,
 }
 
-impl ProcMacroProcessSrv {
+impl ProcMacroServerProcess {
     pub(crate) fn run(
         process_path: &AbsPath,
         env: impl IntoIterator<Item = (impl AsRef<std::ffi::OsStr>, impl AsRef<std::ffi::OsStr>)>
             + Clone,
-    ) -> io::Result<ProcMacroProcessSrv> {
+    ) -> io::Result<ProcMacroServerProcess> {
         let create_srv = |null_stderr| {
             let mut process = Process::run(process_path, env.clone(), null_stderr)?;
             let (stdin, stdout) = process.stdio().expect("couldn't access child stdio");
 
-            io::Result::Ok(ProcMacroProcessSrv {
+            io::Result::Ok(ProcMacroServerProcess {
                 state: Mutex::new(ProcessSrvState { process, stdin, stdout }),
                 version: 0,
                 mode: SpanMode::Id,
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 137efd5e7a0..de59e88aac4 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
@@ -6,7 +6,10 @@
 #[cfg(feature = "in-rust-tree")]
 extern crate rustc_driver as _;
 
-use std::io;
+#[cfg(any(feature = "sysroot-abi", rust_analyzer))]
+mod main_loop;
+#[cfg(any(feature = "sysroot-abi", rust_analyzer))]
+use main_loop::run;
 
 fn main() -> std::io::Result<()> {
     let v = std::env::var("RUST_ANALYZER_INTERNALS_DO_NOT_USE");
@@ -22,57 +25,10 @@ fn main() -> std::io::Result<()> {
 }
 
 #[cfg(not(any(feature = "sysroot-abi", rust_analyzer)))]
-fn run() -> io::Result<()> {
-    Err(io::Error::new(
-        io::ErrorKind::Unsupported,
+fn run() -> std::io::Result<()> {
+    Err(std::io::Error::new(
+        std::io::ErrorKind::Unsupported,
         "proc-macro-srv-cli needs to be compiled with the `sysroot-abi` feature to function"
             .to_owned(),
     ))
 }
-
-#[cfg(any(feature = "sysroot-abi", rust_analyzer))]
-fn run() -> io::Result<()> {
-    use proc_macro_api::{
-        json::{read_json, write_json},
-        msg::{self, Message},
-    };
-    use proc_macro_srv::EnvSnapshot;
-
-    let read_request =
-        |buf: &mut String| msg::Request::read(read_json, &mut io::stdin().lock(), buf);
-
-    let write_response = |msg: msg::Response| msg.write(write_json, &mut io::stdout().lock());
-
-    let env = EnvSnapshot::default();
-    let mut srv = proc_macro_srv::ProcMacroSrv::new(&env);
-    let mut buf = String::new();
-
-    while let Some(req) = read_request(&mut buf)? {
-        let res = match req {
-            msg::Request::ListMacros { dylib_path } => {
-                msg::Response::ListMacros(srv.list_macros(&dylib_path))
-            }
-            msg::Request::ExpandMacro(task) => match srv.span_mode() {
-                msg::SpanMode::Id => {
-                    msg::Response::ExpandMacro(srv.expand(*task).map(|(it, _)| it))
-                }
-                msg::SpanMode::RustAnalyzer => msg::Response::ExpandMacroExtended(
-                    srv.expand(*task).map(|(tree, span_data_table)| msg::ExpandMacroExtended {
-                        tree,
-                        span_data_table,
-                    }),
-                ),
-            },
-            msg::Request::ApiVersionCheck {} => {
-                msg::Response::ApiVersionCheck(proc_macro_api::msg::CURRENT_API_VERSION)
-            }
-            msg::Request::SetConfig(config) => {
-                srv.set_span_mode(config.span_mode);
-                msg::Response::SetConfig(config)
-            }
-        };
-        write_response(res)?
-    }
-
-    Ok(())
-}
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv-cli/src/main_loop.rs b/src/tools/rust-analyzer/crates/proc-macro-srv-cli/src/main_loop.rs
new file mode 100644
index 00000000000..74590f94c76
--- /dev/null
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv-cli/src/main_loop.rs
@@ -0,0 +1,128 @@
+//! The main loop of the proc-macro server.
+use std::io;
+
+use proc_macro_api::{
+    json::{read_json, write_json},
+    msg::{
+        self, deserialize_span_data_index_map, serialize_span_data_index_map, ExpandMacroData,
+        ExpnGlobals, Message, SpanMode, TokenId, CURRENT_API_VERSION,
+    },
+};
+use proc_macro_srv::EnvSnapshot;
+
+pub(crate) fn run() -> io::Result<()> {
+    fn macro_kind_to_api(kind: proc_macro_srv::ProcMacroKind) -> proc_macro_api::ProcMacroKind {
+        match kind {
+            proc_macro_srv::ProcMacroKind::CustomDerive => {
+                proc_macro_api::ProcMacroKind::CustomDerive
+            }
+            proc_macro_srv::ProcMacroKind::Bang => proc_macro_api::ProcMacroKind::Bang,
+            proc_macro_srv::ProcMacroKind::Attr => proc_macro_api::ProcMacroKind::Attr,
+        }
+    }
+
+    let read_request =
+        |buf: &mut String| msg::Request::read(read_json, &mut io::stdin().lock(), buf);
+
+    let write_response = |msg: msg::Response| msg.write(write_json, &mut io::stdout().lock());
+
+    let env = EnvSnapshot::default();
+    let mut srv = proc_macro_srv::ProcMacroSrv::new(&env);
+    let mut buf = String::new();
+
+    let mut span_mode = SpanMode::Id;
+
+    while let Some(req) = read_request(&mut buf)? {
+        let res = match req {
+            msg::Request::ListMacros { dylib_path } => {
+                msg::Response::ListMacros(srv.list_macros(&dylib_path).map(|macros| {
+                    macros.into_iter().map(|(name, kind)| (name, macro_kind_to_api(kind))).collect()
+                }))
+            }
+            msg::Request::ExpandMacro(task) => {
+                let msg::ExpandMacro {
+                    lib,
+                    env,
+                    current_dir,
+                    data:
+                        ExpandMacroData {
+                            macro_body,
+                            macro_name,
+                            attributes,
+                            has_global_spans:
+                                ExpnGlobals { serialize: _, def_site, call_site, mixed_site },
+                            span_data_table,
+                        },
+                } = *task;
+                match span_mode {
+                    SpanMode::Id => msg::Response::ExpandMacro({
+                        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 = macro_body.to_subtree_unresolved(CURRENT_API_VERSION);
+                        let attributes =
+                            attributes.map(|it| it.to_subtree_unresolved(CURRENT_API_VERSION));
+
+                        srv.expand(
+                            lib,
+                            env,
+                            current_dir,
+                            macro_name,
+                            macro_body,
+                            attributes,
+                            def_site,
+                            call_site,
+                            mixed_site,
+                        )
+                        .map(|it| msg::FlatTree::new_raw(&it, CURRENT_API_VERSION))
+                        .map_err(msg::PanicMessage)
+                    }),
+                    SpanMode::RustAnalyzer => msg::Response::ExpandMacroExtended({
+                        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 =
+                            macro_body.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table);
+                        let attributes = attributes.map(|it| {
+                            it.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table)
+                        });
+                        srv.expand(
+                            lib,
+                            env,
+                            current_dir,
+                            macro_name,
+                            macro_body,
+                            attributes,
+                            def_site,
+                            call_site,
+                            mixed_site,
+                        )
+                        .map(|it| {
+                            (
+                                msg::FlatTree::new(&it, CURRENT_API_VERSION, &mut span_data_table),
+                                serialize_span_data_index_map(&span_data_table),
+                            )
+                        })
+                        .map(|(tree, span_data_table)| msg::ExpandMacroExtended {
+                            tree,
+                            span_data_table,
+                        })
+                        .map_err(msg::PanicMessage)
+                    }),
+                }
+            }
+            msg::Request::ApiVersionCheck {} => msg::Response::ApiVersionCheck(CURRENT_API_VERSION),
+            msg::Request::SetConfig(config) => {
+                span_mode = config.span_mode;
+                msg::Response::SetConfig(config)
+            }
+        };
+        write_response(res)?
+    }
+
+    Ok(())
+}
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 98385969459..00695c54737 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/Cargo.toml
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/Cargo.toml
@@ -23,7 +23,6 @@ syntax-bridge.workspace = true
 paths.workspace = true
 # span = {workspace = true, default-features = false} does not work
 span = { path = "../span", version = "0.0.0", default-features = false}
-proc-macro-api.workspace = true
 intern.workspace = true
 
 ra-ap-rustc_lexer.workspace = true
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 26f6af84aae..a5de61ded16 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
@@ -8,9 +8,8 @@ use std::{fmt, fs, io, time::SystemTime};
 use libloading::Library;
 use object::Object;
 use paths::{Utf8Path, Utf8PathBuf};
-use proc_macro_api::ProcMacroKind;
 
-use crate::ProcMacroSrvSpan;
+use crate::{proc_macros::ProcMacroKind, ProcMacroSrvSpan};
 
 const NEW_REGISTRAR_SYMBOL: &str = "_rustc_proc_macro_decls_";
 
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 c8d9e6cc299..2b7f53e66b4 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
@@ -14,6 +14,7 @@
 #![cfg_attr(feature = "in-rust-tree", feature(rustc_private))]
 #![feature(proc_macro_internals, proc_macro_diagnostic, proc_macro_span)]
 #![allow(unreachable_pub, internal_features, clippy::disallowed_types, clippy::print_stderr)]
+#![deny(deprecated_safe)]
 
 extern crate proc_macro;
 #[cfg(feature = "in-rust-tree")]
@@ -38,62 +39,75 @@ use std::{
 };
 
 use paths::{Utf8Path, Utf8PathBuf};
-use proc_macro_api::{
-    msg::{
-        self, deserialize_span_data_index_map, serialize_span_data_index_map, ExpnGlobals,
-        SpanMode, TokenId, CURRENT_API_VERSION,
-    },
-    ProcMacroKind,
-};
-use span::Span;
+use span::{Span, TokenId};
 
 use crate::server_impl::TokenStream;
 
+pub use crate::proc_macros::ProcMacroKind;
+
 pub const RUSTC_VERSION_STRING: &str = env!("RUSTC_VERSION");
 
 pub struct ProcMacroSrv<'env> {
     expanders: HashMap<Utf8PathBuf, 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 }
+        Self { expanders: Default::default(), env }
     }
 }
 
 const EXPANDER_STACK_SIZE: usize = 8 * 1024 * 1024;
 
 impl ProcMacroSrv<'_> {
-    pub fn set_span_mode(&mut self, span_mode: SpanMode) {
-        self.span_mode = span_mode;
-    }
-
-    pub fn span_mode(&self) -> SpanMode {
-        self.span_mode
-    }
-
-    pub fn expand(
+    pub fn expand<S: ProcMacroSrvSpan>(
         &mut self,
-        msg::ExpandMacro { lib, env, current_dir, data }: msg::ExpandMacro,
-    ) -> Result<(msg::FlatTree, Vec<u32>), msg::PanicMessage> {
-        let span_mode = self.span_mode;
+        lib: impl AsRef<Utf8Path>,
+        env: Vec<(String, String)>,
+        current_dir: Option<impl AsRef<Path>>,
+        macro_name: String,
+        macro_body: tt::Subtree<S>,
+        attribute: Option<tt::Subtree<S>>,
+        def_site: S,
+        call_site: S,
+        mixed_site: S,
+    ) -> Result<tt::Subtree<S>, String> {
         let snapped_env = self.env;
-        let expander = self
-            .expander(lib.as_ref())
-            .map_err(|err| msg::PanicMessage(format!("failed to load macro: {err}")))?;
+        let expander =
+            self.expander(lib.as_ref()).map_err(|err| format!("failed to load macro: {err}"))?;
 
         let prev_env = EnvChange::apply(snapped_env, env, current_dir.as_ref().map(<_>::as_ref));
 
-        let result = match span_mode {
-            SpanMode::Id => expand_id(data, expander).map(|it| (it, vec![])),
-            SpanMode::RustAnalyzer => expand_ra_span(data, expander),
-        };
-
+        // 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(macro_name.clone())
+                .spawn_scoped(s, move || {
+                    expander.expand(
+                        &macro_name,
+                        macro_body,
+                        attribute,
+                        def_site,
+                        call_site,
+                        mixed_site,
+                    )
+                });
+            let res = match thread {
+                Ok(handle) => handle.join(),
+                Err(e) => return Err(e.to_string()),
+            };
+
+            match res {
+                Ok(res) => res,
+                Err(e) => std::panic::resume_unwind(e),
+            }
+        });
         prev_env.rollback();
 
-        result.map_err(msg::PanicMessage)
+        result
     }
 
     pub fn list_macros(
@@ -123,7 +137,7 @@ impl ProcMacroSrv<'_> {
     }
 }
 
-trait ProcMacroSrvSpan: Copy {
+pub trait ProcMacroSrvSpan: Copy + Send {
     type Server: proc_macro::bridge::server::Server<TokenStream = TokenStream<Self>>;
     fn make_server(call_site: Self, def_site: Self, mixed_site: Self) -> Self::Server;
 }
@@ -147,93 +161,6 @@ impl ProcMacroSrvSpan for Span {
         }
     }
 }
-
-fn expand_id(
-    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,
-) -> 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 = 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(macro_name.clone())
-            .spawn_scoped(s, || {
-                expander
-                    .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 {
-            Ok(handle) => handle.join(),
-            Err(e) => return Err(e.to_string()),
-        };
-
-        match res {
-            Ok(res) => res,
-            Err(e) => std::panic::resume_unwind(e),
-        }
-    });
-    result
-}
-
-fn expand_ra_span(
-    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,
-) -> Result<(msg::FlatTree, Vec<u32>), String> {
-    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 = macro_body.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table);
-    let attributes =
-        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(macro_name.clone())
-            .spawn_scoped(s, || {
-                expander
-                    .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),
-                            serialize_span_data_index_map(&span_data_table),
-                        )
-                    })
-            });
-        let res = match thread {
-            Ok(handle) => handle.join(),
-            Err(e) => return Err(e.to_string()),
-        };
-
-        match res {
-            Ok(res) => res,
-            Err(e) => std::panic::resume_unwind(e),
-        }
-    });
-    result
-}
-
 pub struct PanicMessage {
     message: Option<String>,
 }
@@ -254,10 +181,13 @@ impl Default for EnvSnapshot {
     }
 }
 
+static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
+
 struct EnvChange<'snap> {
     changed_vars: Vec<String>,
     prev_working_dir: Option<PathBuf>,
     snap: &'snap EnvSnapshot,
+    _guard: std::sync::MutexGuard<'snap, ()>,
 }
 
 impl<'snap> EnvChange<'snap> {
@@ -266,6 +196,7 @@ impl<'snap> EnvChange<'snap> {
         new_vars: Vec<(String, String)>,
         current_dir: Option<&Path>,
     ) -> EnvChange<'snap> {
+        let guard = ENV_LOCK.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
         let prev_working_dir = match current_dir {
             Some(dir) => {
                 let prev_working_dir = std::env::current_dir().ok();
@@ -284,11 +215,13 @@ impl<'snap> EnvChange<'snap> {
             changed_vars: new_vars
                 .into_iter()
                 .map(|(k, v)| {
-                    env::set_var(&k, v);
+                    // SAFETY: We have acquired the environment lock
+                    unsafe { env::set_var(&k, v) };
                     k
                 })
                 .collect(),
             prev_working_dir,
+            _guard: guard,
         }
     }
 
@@ -298,9 +231,12 @@ impl<'snap> EnvChange<'snap> {
 impl Drop for EnvChange<'_> {
     fn drop(&mut self) {
         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),
+            // SAFETY: We have acquired the environment lock
+            unsafe {
+                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),
+                }
             }
         }
 
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 097b39a3f91..6dfc533a61f 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,12 +1,18 @@
 //! Proc macro ABI
 
 use proc_macro::bridge;
-use proc_macro_api::ProcMacroKind;
 
 use libloading::Library;
 
 use crate::{dylib::LoadProcMacroDylibError, ProcMacroSrvSpan};
 
+#[derive(Copy, Clone, Eq, PartialEq, Debug)]
+pub enum ProcMacroKind {
+    CustomDerive,
+    Attr,
+    Bang,
+}
+
 pub(crate) struct ProcMacros {
     exported_macros: Vec<bridge::client::ProcMacro>,
 }
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/token_id.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/token_id.rs
index e478b1c853b..081213c570c 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/token_id.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/token_id.rs
@@ -13,7 +13,7 @@ use crate::server_impl::{
     token_stream::TokenStreamBuilder,
 };
 mod tt {
-    pub use proc_macro_api::msg::TokenId;
+    pub use span::TokenId;
 
     pub use tt::*;
 
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 cc5d4a89131..4b8ea10ebc0 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
@@ -1,8 +1,7 @@
 //! utils used in proc-macro tests
 
 use expect_test::Expect;
-use proc_macro_api::msg::TokenId;
-use span::{EditionedFileId, ErasedFileAstId, FileId, Span, SpanAnchor, SyntaxContextId};
+use span::{EditionedFileId, ErasedFileAstId, FileId, Span, SpanAnchor, SyntaxContextId, TokenId};
 use tt::TextRange;
 
 use crate::{dylib, proc_macro_test_dylib_path, EnvSnapshot, ProcMacroSrv};
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/global_state.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/global_state.rs
index dd13bdba4cb..58b80797cf0 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/global_state.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/global_state.rs
@@ -17,7 +17,7 @@ use parking_lot::{
     MappedRwLockReadGuard, Mutex, RwLock, RwLockReadGuard, RwLockUpgradableReadGuard,
     RwLockWriteGuard,
 };
-use proc_macro_api::ProcMacroServer;
+use proc_macro_api::ProcMacroClient;
 use project_model::{ManifestPath, ProjectWorkspace, ProjectWorkspaceKind, WorkspaceBuildScripts};
 use rustc_hash::{FxHashMap, FxHashSet};
 use tracing::{span, trace, Level};
@@ -95,7 +95,7 @@ pub(crate) struct GlobalState {
     pub(crate) last_reported_status: lsp_ext::ServerStatusParams,
 
     // proc macros
-    pub(crate) proc_macro_clients: Arc<[anyhow::Result<ProcMacroServer>]>,
+    pub(crate) proc_macro_clients: Arc<[anyhow::Result<ProcMacroClient>]>,
     pub(crate) build_deps_changed: bool,
 
     // Flycheck
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/reload.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/reload.rs
index 3444773695b..1996c2b6421 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/reload.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/reload.rs
@@ -24,7 +24,7 @@ use ide_db::{
 use itertools::Itertools;
 use load_cargo::{load_proc_macro, ProjectFolders};
 use lsp_types::FileSystemWatcher;
-use proc_macro_api::ProcMacroServer;
+use proc_macro_api::ProcMacroClient;
 use project_model::{ManifestPath, ProjectWorkspace, ProjectWorkspaceKind, WorkspaceBuildScripts};
 use stdx::{format_to, thread::ThreadIntent};
 use triomphe::Arc;
@@ -650,7 +650,7 @@ impl GlobalState {
                 };
                 info!("Using proc-macro server at {path}");
 
-                ProcMacroServer::spawn(&path, &env).map_err(|err| {
+                ProcMacroClient::spawn(&path, &env).map_err(|err| {
                     tracing::error!(
                         "Failed to run proc-macro server from path {path}, error: {err:?}",
                     );
diff --git a/src/tools/rust-analyzer/crates/span/src/lib.rs b/src/tools/rust-analyzer/crates/span/src/lib.rs
index 20c3b087af5..8dc95735038 100644
--- a/src/tools/rust-analyzer/crates/span/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/span/src/lib.rs
@@ -358,6 +358,18 @@ impl HirFileId {
     }
 }
 
+/// Legacy span type, only defined here as it is still used by the proc-macro server.
+/// While rust-analyzer doesn't use this anymore at all, RustRover relies on the legacy type for
+/// proc-macro expansion.
+#[derive(Clone, Copy, PartialEq, Eq, Hash)]
+pub struct TokenId(pub u32);
+
+impl std::fmt::Debug for TokenId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
 #[cfg(not(feature = "ra-salsa"))]
 mod intern_id_proxy {
     use std::fmt;