Decouple proc-macro server protocol from the server implementation

Author: Veykril

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",

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 {

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 {

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;

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

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,

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

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(())
+}

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