Abtract away json protocol for proc-macro-srv

5 files changed, 311 insertions, 197 deletions

diff --git a/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol.rs b/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol.rs
new file mode 100644
index 00000000000..ee96b899fe5
--- /dev/null
+++ b/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol.rs
@@ -0,0 +1,172 @@
+//! The initial proc-macro-srv protocol, soon to be deprecated.
+
+pub mod json;
+pub mod msg;
+
+use std::{
+    io::{BufRead, Write},
+    sync::Arc,
+};
+
+use paths::AbsPath;
+use span::Span;
+
+use crate::{
+    ProcMacro, ProcMacroKind, ServerError,
+    legacy_protocol::{
+        json::{read_json, write_json},
+        msg::{
+            ExpandMacro, ExpandMacroData, ExpnGlobals, FlatTree, Message, Request, Response,
+            ServerConfig, SpanDataIndexMap, deserialize_span_data_index_map,
+            flat::serialize_span_data_index_map,
+        },
+    },
+    process::ProcMacroServerProcess,
+    version,
+};
+
+pub(crate) use crate::legacy_protocol::msg::SpanMode;
+
+/// 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 SpanId(pub u32);
+
+impl std::fmt::Debug for SpanId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+pub(crate) fn version_check(srv: &ProcMacroServerProcess) -> Result<u32, ServerError> {
+    let request = Request::ApiVersionCheck {};
+    let response = send_task(srv, request)?;
+
+    match response {
+        Response::ApiVersionCheck(version) => Ok(version),
+        _ => Err(ServerError { message: "unexpected response".to_owned(), io: None }),
+    }
+}
+
+/// Enable support for rust-analyzer span mode if the server supports it.
+pub(crate) fn enable_rust_analyzer_spans(
+    srv: &ProcMacroServerProcess,
+) -> Result<SpanMode, ServerError> {
+    let request = Request::SetConfig(ServerConfig { span_mode: SpanMode::RustAnalyzer });
+    let response = send_task(srv, request)?;
+
+    match response {
+        Response::SetConfig(ServerConfig { span_mode }) => Ok(span_mode),
+        _ => Err(ServerError { message: "unexpected response".to_owned(), io: None }),
+    }
+}
+
+/// Finds proc-macros in a given dynamic library.
+pub(crate) fn find_proc_macros(
+    srv: &ProcMacroServerProcess,
+    dylib_path: &AbsPath,
+) -> Result<Result<Vec<(String, ProcMacroKind)>, String>, ServerError> {
+    let request = Request::ListMacros { dylib_path: dylib_path.to_path_buf().into() };
+
+    let response = send_task(srv, request)?;
+
+    match response {
+        Response::ListMacros(it) => Ok(it),
+        _ => Err(ServerError { message: "unexpected response".to_owned(), io: None }),
+    }
+}
+
+pub(crate) fn expand(
+    proc_macro: &ProcMacro,
+    subtree: tt::SubtreeView<'_, Span>,
+    attr: Option<tt::SubtreeView<'_, Span>>,
+    env: Vec<(String, String)>,
+    def_site: Span,
+    call_site: Span,
+    mixed_site: Span,
+    current_dir: String,
+) -> Result<Result<tt::TopSubtree<span::SpanData<span::SyntaxContext>>, String>, crate::ServerError>
+{
+    let version = proc_macro.process.version();
+    let mut span_data_table = SpanDataIndexMap::default();
+    let def_site = span_data_table.insert_full(def_site).0;
+    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 {
+        data: ExpandMacroData {
+            macro_body: FlatTree::new(subtree, version, &mut span_data_table),
+            macro_name: proc_macro.name.to_string(),
+            attributes: attr.map(|subtree| FlatTree::new(subtree, version, &mut span_data_table)),
+            has_global_spans: ExpnGlobals {
+                serialize: version >= version::HAS_GLOBAL_SPANS,
+                def_site,
+                call_site,
+                mixed_site,
+            },
+            span_data_table: if proc_macro.process.rust_analyzer_spans() {
+                serialize_span_data_index_map(&span_data_table)
+            } else {
+                Vec::new()
+            },
+        },
+        lib: proc_macro.dylib_path.to_path_buf().into(),
+        env,
+        current_dir: Some(current_dir),
+    };
+
+    let response = send_task(&proc_macro.process, Request::ExpandMacro(Box::new(task)))?;
+
+    match response {
+        Response::ExpandMacro(it) => Ok(it
+            .map(|tree| {
+                let mut expanded = FlatTree::to_subtree_resolved(tree, version, &span_data_table);
+                if proc_macro.needs_fixup_change() {
+                    proc_macro.change_fixup_to_match_old_server(&mut expanded);
+                }
+                expanded
+            })
+            .map_err(|msg| msg.0)),
+        Response::ExpandMacroExtended(it) => Ok(it
+            .map(|resp| {
+                let mut expanded = FlatTree::to_subtree_resolved(
+                    resp.tree,
+                    version,
+                    &deserialize_span_data_index_map(&resp.span_data_table),
+                );
+                if proc_macro.needs_fixup_change() {
+                    proc_macro.change_fixup_to_match_old_server(&mut expanded);
+                }
+                expanded
+            })
+            .map_err(|msg| msg.0)),
+        _ => Err(ServerError { message: "unexpected response".to_owned(), io: None }),
+    }
+}
+
+/// Sends a request to the proc-macro server and waits for a response.
+fn send_task(srv: &ProcMacroServerProcess, req: Request) -> Result<Response, ServerError> {
+    if let Some(server_error) = srv.exited() {
+        return Err(server_error.clone());
+    }
+
+    srv.send_task(send_request, req)
+}
+
+/// Sends a request to the server and reads the response.
+fn send_request(
+    mut writer: &mut dyn Write,
+    mut reader: &mut dyn BufRead,
+    req: Request,
+    buf: &mut String,
+) -> Result<Option<Response>, ServerError> {
+    req.write(write_json, &mut writer).map_err(|err| ServerError {
+        message: "failed to write request".into(),
+        io: Some(Arc::new(err)),
+    })?;
+    let res = Response::read(read_json, &mut reader, buf).map_err(|err| ServerError {
+        message: "failed to read response".into(),
+        io: Some(Arc::new(err)),
+    })?;
+    Ok(res)
+}
diff --git a/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol/msg.rs b/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol/msg.rs
index 165936269d3..b795c455895 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol/msg.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol/msg.rs
@@ -1,5 +1,6 @@
 //! Defines messages for cross-process message passing based on `ndjson` wire protocol
 pub(crate) mod flat;
+pub use self::flat::*;
 
 use std::io::{self, BufRead, Write};
 
@@ -9,24 +10,6 @@ use serde_derive::{Deserialize, Serialize};
 
 use crate::ProcMacroKind;
 
-pub use self::flat::{
-    FlatTree, SpanDataIndexMap, deserialize_span_data_index_map, serialize_span_data_index_map,
-};
-pub use span::TokenId;
-
-// The versions of the server protocol
-pub const NO_VERSION_CHECK_VERSION: u32 = 0;
-pub const VERSION_CHECK_VERSION: u32 = 1;
-pub const ENCODE_CLOSE_SPAN_VERSION: u32 = 2;
-pub const HAS_GLOBAL_SPANS: u32 = 3;
-pub const RUST_ANALYZER_SPAN_SUPPORT: u32 = 4;
-/// Whether literals encode their kind as an additional u32 field and idents their rawness as a u32 field.
-pub const EXTENDED_LEAF_DATA: u32 = 5;
-pub const HASHED_AST_ID: u32 = 6;
-
-/// Current API version of the proc-macro protocol.
-pub const CURRENT_API_VERSION: u32 = HASHED_AST_ID;
-
 /// Represents requests sent from the client to the proc-macro-srv.
 #[derive(Debug, Serialize, Deserialize)]
 pub enum Request {
@@ -48,7 +31,7 @@ pub enum Request {
 }
 
 /// Defines the mode used for handling span data.
-#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub enum SpanMode {
     /// Default mode, where spans are identified by an ID.
     #[default]
@@ -210,6 +193,8 @@ mod tests {
         TopSubtreeBuilder,
     };
 
+    use crate::version;
+
     use super::*;
 
     fn fixture_token_tree() -> TopSubtree<Span> {
@@ -308,7 +293,7 @@ mod tests {
     #[test]
     fn test_proc_macro_rpc_works() {
         let tt = fixture_token_tree();
-        for v in RUST_ANALYZER_SPAN_SUPPORT..=CURRENT_API_VERSION {
+        for v in version::RUST_ANALYZER_SPAN_SUPPORT..=version::CURRENT_API_VERSION {
             let mut span_data_table = Default::default();
             let task = ExpandMacro {
                 data: ExpandMacroData {
diff --git a/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol/msg/flat.rs b/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol/msg/flat.rs
index 597ffa05d20..fb3542d24f4 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol/msg/flat.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-api/src/legacy_protocol/msg/flat.rs
@@ -40,9 +40,12 @@ use std::collections::VecDeque;
 use intern::Symbol;
 use rustc_hash::FxHashMap;
 use serde_derive::{Deserialize, Serialize};
-use span::{EditionedFileId, ErasedFileAstId, Span, SpanAnchor, SyntaxContext, TextRange, TokenId};
+use span::{EditionedFileId, ErasedFileAstId, Span, SpanAnchor, SyntaxContext, TextRange};
 
-use crate::legacy_protocol::msg::{ENCODE_CLOSE_SPAN_VERSION, EXTENDED_LEAF_DATA};
+use crate::{
+    legacy_protocol::SpanId,
+    version::{ENCODE_CLOSE_SPAN_VERSION, EXTENDED_LEAF_DATA},
+};
 
 pub type SpanDataIndexMap =
     indexmap::IndexSet<Span, std::hash::BuildHasherDefault<rustc_hash::FxHasher>>;
@@ -62,7 +65,7 @@ pub fn serialize_span_data_index_map(map: &SpanDataIndexMap) -> Vec<u32> {
 }
 
 pub fn deserialize_span_data_index_map(map: &[u32]) -> SpanDataIndexMap {
-    debug_assert!(map.len() % 5 == 0);
+    debug_assert!(map.len().is_multiple_of(5));
     map.chunks_exact(5)
         .map(|span| {
             let &[file_id, ast_id, start, end, e] = span else { unreachable!() };
@@ -91,27 +94,27 @@ pub struct FlatTree {
 }
 
 struct SubtreeRepr {
-    open: TokenId,
-    close: TokenId,
+    open: SpanId,
+    close: SpanId,
     kind: tt::DelimiterKind,
     tt: [u32; 2],
 }
 
 struct LiteralRepr {
-    id: TokenId,
+    id: SpanId,
     text: u32,
     suffix: u32,
     kind: u16,
 }
 
 struct PunctRepr {
-    id: TokenId,
+    id: SpanId,
     char: char,
     spacing: tt::Spacing,
 }
 
 struct IdentRepr {
-    id: TokenId,
+    id: SpanId,
     text: u32,
     is_raw: bool,
 }
@@ -122,7 +125,7 @@ impl FlatTree {
         version: u32,
         span_data_table: &mut SpanDataIndexMap,
     ) -> FlatTree {
-        let mut w = Writer {
+        let mut w = Writer::<Span> {
             string_table: FxHashMap::default(),
             work: VecDeque::new(),
             span_data_table,
@@ -159,8 +162,11 @@ impl FlatTree {
         }
     }
 
-    pub fn new_raw(subtree: tt::SubtreeView<'_, TokenId>, version: u32) -> FlatTree {
-        let mut w = Writer {
+    pub fn new_raw<T: SpanTransformer<Table = ()>>(
+        subtree: tt::SubtreeView<'_, T::Span>,
+        version: u32,
+    ) -> FlatTree {
+        let mut w = Writer::<T> {
             string_table: FxHashMap::default(),
             work: VecDeque::new(),
             span_data_table: &mut (),
@@ -202,7 +208,7 @@ impl FlatTree {
         version: u32,
         span_data_table: &SpanDataIndexMap,
     ) -> tt::TopSubtree<Span> {
-        Reader {
+        Reader::<Span> {
             subtree: if version >= ENCODE_CLOSE_SPAN_VERSION {
                 read_vec(self.subtree, SubtreeRepr::read_with_close_span)
             } else {
@@ -227,8 +233,11 @@ impl FlatTree {
         .read()
     }
 
-    pub fn to_subtree_unresolved(self, version: u32) -> tt::TopSubtree<TokenId> {
-        Reader {
+    pub fn to_subtree_unresolved<T: SpanTransformer<Table = ()>>(
+        self,
+        version: u32,
+    ) -> tt::TopSubtree<T::Span> {
+        Reader::<T> {
             subtree: if version >= ENCODE_CLOSE_SPAN_VERSION {
                 read_vec(self.subtree, SubtreeRepr::read_with_close_span)
             } else {
@@ -283,7 +292,7 @@ impl SubtreeRepr {
             3 => tt::DelimiterKind::Bracket,
             other => panic!("bad kind {other}"),
         };
-        SubtreeRepr { open: TokenId(open), close: TokenId(!0), kind, tt: [lo, len] }
+        SubtreeRepr { open: SpanId(open), close: SpanId(!0), kind, tt: [lo, len] }
     }
     fn write_with_close_span(self) -> [u32; 5] {
         let kind = match self.kind {
@@ -302,7 +311,7 @@ impl SubtreeRepr {
             3 => tt::DelimiterKind::Bracket,
             other => panic!("bad kind {other}"),
         };
-        SubtreeRepr { open: TokenId(open), close: TokenId(close), kind, tt: [lo, len] }
+        SubtreeRepr { open: SpanId(open), close: SpanId(close), kind, tt: [lo, len] }
     }
 }
 
@@ -311,13 +320,13 @@ impl LiteralRepr {
         [self.id.0, self.text]
     }
     fn read([id, text]: [u32; 2]) -> LiteralRepr {
-        LiteralRepr { id: TokenId(id), text, kind: 0, suffix: !0 }
+        LiteralRepr { id: SpanId(id), text, kind: 0, suffix: !0 }
     }
     fn write_with_kind(self) -> [u32; 4] {
         [self.id.0, self.text, self.kind as u32, self.suffix]
     }
     fn read_with_kind([id, text, kind, suffix]: [u32; 4]) -> LiteralRepr {
-        LiteralRepr { id: TokenId(id), text, kind: kind as u16, suffix }
+        LiteralRepr { id: SpanId(id), text, kind: kind as u16, suffix }
     }
 }
 
@@ -335,7 +344,7 @@ impl PunctRepr {
             1 => tt::Spacing::Joint,
             other => panic!("bad spacing {other}"),
         };
-        PunctRepr { id: TokenId(id), char: char.try_into().unwrap(), spacing }
+        PunctRepr { id: SpanId(id), char: char.try_into().unwrap(), spacing }
     }
 }
 
@@ -344,44 +353,46 @@ impl IdentRepr {
         [self.id.0, self.text]
     }
     fn read(data: [u32; 2]) -> IdentRepr {
-        IdentRepr { id: TokenId(data[0]), text: data[1], is_raw: false }
+        IdentRepr { id: SpanId(data[0]), text: data[1], is_raw: false }
     }
     fn write_with_rawness(self) -> [u32; 3] {
         [self.id.0, self.text, self.is_raw as u32]
     }
     fn read_with_rawness([id, text, is_raw]: [u32; 3]) -> IdentRepr {
-        IdentRepr { id: TokenId(id), text, is_raw: is_raw == 1 }
+        IdentRepr { id: SpanId(id), text, is_raw: is_raw == 1 }
     }
 }
 
-trait InternableSpan: Copy {
+pub trait SpanTransformer {
     type Table;
-    fn token_id_of(table: &mut Self::Table, s: Self) -> TokenId;
-    fn span_for_token_id(table: &Self::Table, id: TokenId) -> Self;
+    type Span: Copy;
+    fn token_id_of(table: &mut Self::Table, s: Self::Span) -> SpanId;
+    fn span_for_token_id(table: &Self::Table, id: SpanId) -> Self::Span;
 }
-
-impl InternableSpan for TokenId {
+impl SpanTransformer for SpanId {
     type Table = ();
-    fn token_id_of((): &mut Self::Table, token_id: Self) -> TokenId {
+    type Span = Self;
+    fn token_id_of((): &mut Self::Table, token_id: Self::Span) -> SpanId {
         token_id
     }
 
-    fn span_for_token_id((): &Self::Table, id: TokenId) -> Self {
+    fn span_for_token_id((): &Self::Table, id: SpanId) -> Self::Span {
         id
     }
 }
-impl InternableSpan for Span {
+impl SpanTransformer for Span {
     type Table = SpanDataIndexMap;
-    fn token_id_of(table: &mut Self::Table, span: Self) -> TokenId {
-        TokenId(table.insert_full(span).0 as u32)
+    type Span = Self;
+    fn token_id_of(table: &mut Self::Table, span: Self::Span) -> SpanId {
+        SpanId(table.insert_full(span).0 as u32)
     }
-    fn span_for_token_id(table: &Self::Table, id: TokenId) -> Self {
+    fn span_for_token_id(table: &Self::Table, id: SpanId) -> Self::Span {
         *table.get_index(id.0 as usize).unwrap_or_else(|| &table[0])
     }
 }
 
-struct Writer<'a, 'span, S: InternableSpan> {
-    work: VecDeque<(usize, tt::iter::TtIter<'a, S>)>,
+struct Writer<'a, 'span, S: SpanTransformer> {
+    work: VecDeque<(usize, tt::iter::TtIter<'a, S::Span>)>,
     string_table: FxHashMap<std::borrow::Cow<'a, str>, u32>,
     span_data_table: &'span mut S::Table,
     version: u32,
@@ -394,8 +405,8 @@ struct Writer<'a, 'span, S: InternableSpan> {
     text: Vec<String>,
 }
 
-impl<'a, S: InternableSpan> Writer<'a, '_, S> {
-    fn write(&mut self, root: tt::SubtreeView<'a, S>) {
+impl<'a, T: SpanTransformer> Writer<'a, '_, T> {
+    fn write(&mut self, root: tt::SubtreeView<'a, T::Span>) {
         let subtree = root.top_subtree();
         self.enqueue(subtree, root.iter());
         while let Some((idx, subtree)) = self.work.pop_front() {
@@ -403,11 +414,11 @@ impl<'a, S: InternableSpan> Writer<'a, '_, S> {
         }
     }
 
-    fn token_id_of(&mut self, span: S) -> TokenId {
-        S::token_id_of(self.span_data_table, span)
+    fn token_id_of(&mut self, span: T::Span) -> SpanId {
+        T::token_id_of(self.span_data_table, span)
     }
 
-    fn subtree(&mut self, idx: usize, subtree: tt::iter::TtIter<'a, S>) {
+    fn subtree(&mut self, idx: usize, subtree: tt::iter::TtIter<'a, T::Span>) {
         let mut first_tt = self.token_tree.len();
         let n_tt = subtree.clone().count(); // FIXME: `count()` walks over the entire iterator.
         self.token_tree.resize(first_tt + n_tt, !0);
@@ -478,7 +489,11 @@ impl<'a, S: InternableSpan> Writer<'a, '_, S> {
         }
     }
 
-    fn enqueue(&mut self, subtree: &'a tt::Subtree<S>, contents: tt::iter::TtIter<'a, S>) -> u32 {
+    fn enqueue(
+        &mut self,
+        subtree: &'a tt::Subtree<T::Span>,
+        contents: tt::iter::TtIter<'a, T::Span>,
+    ) -> u32 {
         let idx = self.subtree.len();
         let open = self.token_id_of(subtree.delimiter.open);
         let close = self.token_id_of(subtree.delimiter.close);
@@ -507,7 +522,7 @@ impl<'a, S: InternableSpan> Writer<'a, '_, S> {
     }
 }
 
-struct Reader<'span, S: InternableSpan> {
+struct Reader<'span, S: SpanTransformer> {
     version: u32,
     subtree: Vec<SubtreeRepr>,
     literal: Vec<LiteralRepr>,
@@ -518,11 +533,11 @@ struct Reader<'span, S: InternableSpan> {
     span_data_table: &'span S::Table,
 }
 
-impl<S: InternableSpan> Reader<'_, S> {
-    pub(crate) fn read(self) -> tt::TopSubtree<S> {
-        let mut res: Vec<Option<(tt::Delimiter<S>, Vec<tt::TokenTree<S>>)>> =
+impl<T: SpanTransformer> Reader<'_, T> {
+    pub(crate) fn read(self) -> tt::TopSubtree<T::Span> {
+        let mut res: Vec<Option<(tt::Delimiter<T::Span>, Vec<tt::TokenTree<T::Span>>)>> =
             vec![None; self.subtree.len()];
-        let read_span = |id| S::span_for_token_id(self.span_data_table, id);
+        let read_span = |id| T::span_for_token_id(self.span_data_table, id);
         for i in (0..self.subtree.len()).rev() {
             let repr = &self.subtree[i];
             let token_trees = &self.token_tree[repr.tt[0] as usize..repr.tt[1] as usize];
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 516c7418bde..97919b85b51 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
@@ -5,24 +5,29 @@
 //! is used to provide basic infrastructure for communication between two
 //! processes: Client (RA itself), Server (the external program)
 
-pub mod legacy_protocol {
-    pub mod json;
-    pub mod msg;
-}
+pub mod legacy_protocol;
 mod process;
 
 use paths::{AbsPath, AbsPathBuf};
 use span::{ErasedFileAstId, FIXUP_ERASED_FILE_AST_ID_MARKER, Span};
 use std::{fmt, io, sync::Arc, time::SystemTime};
 
-use crate::{
-    legacy_protocol::msg::{
-        ExpandMacro, ExpandMacroData, ExpnGlobals, FlatTree, HAS_GLOBAL_SPANS, HASHED_AST_ID,
-        PanicMessage, RUST_ANALYZER_SPAN_SUPPORT, Request, Response, SpanDataIndexMap,
-        deserialize_span_data_index_map, flat::serialize_span_data_index_map,
-    },
-    process::ProcMacroServerProcess,
-};
+use crate::process::ProcMacroServerProcess;
+
+/// The versions of the server protocol
+pub mod version {
+    pub const NO_VERSION_CHECK_VERSION: u32 = 0;
+    pub const VERSION_CHECK_VERSION: u32 = 1;
+    pub const ENCODE_CLOSE_SPAN_VERSION: u32 = 2;
+    pub const HAS_GLOBAL_SPANS: u32 = 3;
+    pub const RUST_ANALYZER_SPAN_SUPPORT: u32 = 4;
+    /// Whether literals encode their kind as an additional u32 field and idents their rawness as a u32 field.
+    pub const EXTENDED_LEAF_DATA: u32 = 5;
+    pub const HASHED_AST_ID: u32 = 6;
+
+    /// Current API version of the proc-macro protocol.
+    pub const CURRENT_API_VERSION: u32 = HASHED_AST_ID;
+}
 
 /// Represents different kinds of procedural macros that can be expanded by the external server.
 #[derive(Copy, Clone, Eq, PartialEq, Debug, serde_derive::Serialize, serde_derive::Deserialize)]
@@ -163,7 +168,7 @@ impl ProcMacro {
 
     fn needs_fixup_change(&self) -> bool {
         let version = self.process.version();
-        (RUST_ANALYZER_SPAN_SUPPORT..HASHED_AST_ID).contains(&version)
+        (version::RUST_ANALYZER_SPAN_SUPPORT..version::HASHED_AST_ID).contains(&version)
     }
 
     /// On some server versions, the fixup ast id is different than ours. So change it to match.
@@ -204,7 +209,7 @@ impl ProcMacro {
         call_site: Span,
         mixed_site: Span,
         current_dir: String,
-    ) -> Result<Result<tt::TopSubtree<Span>, PanicMessage>, ServerError> {
+    ) -> Result<Result<tt::TopSubtree<Span>, String>, ServerError> {
         let (mut subtree, mut attr) = (subtree, attr);
         let (mut subtree_changed, mut attr_changed);
         if self.needs_fixup_change() {
@@ -219,57 +224,15 @@ impl ProcMacro {
             }
         }
 
-        let version = self.process.version();
-
-        let mut span_data_table = SpanDataIndexMap::default();
-        let def_site = span_data_table.insert_full(def_site).0;
-        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 {
-            data: 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(),
+        legacy_protocol::expand(
+            self,
+            subtree,
+            attr,
             env,
-            current_dir: Some(current_dir),
-        };
-
-        let response = self.process.send_task(Request::ExpandMacro(Box::new(task)))?;
-
-        match response {
-            Response::ExpandMacro(it) => Ok(it.map(|tree| {
-                let mut expanded = FlatTree::to_subtree_resolved(tree, version, &span_data_table);
-                if self.needs_fixup_change() {
-                    self.change_fixup_to_match_old_server(&mut expanded);
-                }
-                expanded
-            })),
-            Response::ExpandMacroExtended(it) => Ok(it.map(|resp| {
-                let mut expanded = FlatTree::to_subtree_resolved(
-                    resp.tree,
-                    version,
-                    &deserialize_span_data_index_map(&resp.span_data_table),
-                );
-                if self.needs_fixup_change() {
-                    self.change_fixup_to_match_old_server(&mut expanded);
-                }
-                expanded
-            })),
-            _ => Err(ServerError { message: "unexpected response".to_owned(), io: None }),
-        }
+            def_site,
+            call_site,
+            mixed_site,
+            current_dir,
+        )
     }
 }
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 fcea75ef672..278d9cbcda4 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
@@ -12,13 +12,8 @@ use stdx::JodChild;
 
 use crate::{
     ProcMacroKind, ServerError,
-    legacy_protocol::{
-        json::{read_json, write_json},
-        msg::{
-            CURRENT_API_VERSION, Message, RUST_ANALYZER_SPAN_SUPPORT, Request, Response,
-            ServerConfig, SpanMode,
-        },
-    },
+    legacy_protocol::{self, SpanMode},
+    version,
 };
 
 /// Represents a process handling proc-macro communication.
@@ -28,11 +23,16 @@ pub(crate) struct ProcMacroServerProcess {
     /// hence the lock on the state.
     state: Mutex<ProcessSrvState>,
     version: u32,
-    mode: SpanMode,
+    protocol: Protocol,
     /// Populated when the server exits.
     exited: OnceLock<AssertUnwindSafe<ServerError>>,
 }
 
+#[derive(Debug)]
+enum Protocol {
+    LegacyJson { mode: SpanMode },
+}
+
 /// Maintains the state of the proc-macro server process.
 #[derive(Debug)]
 struct ProcessSrvState {
@@ -56,27 +56,26 @@ impl ProcMacroServerProcess {
             io::Result::Ok(ProcMacroServerProcess {
                 state: Mutex::new(ProcessSrvState { process, stdin, stdout }),
                 version: 0,
-                mode: SpanMode::Id,
+                protocol: Protocol::LegacyJson { mode: SpanMode::Id },
                 exited: OnceLock::new(),
             })
         };
         let mut srv = create_srv()?;
         tracing::info!("sending proc-macro server version check");
         match srv.version_check() {
-            Ok(v) if v > CURRENT_API_VERSION => Err(io::Error::other(
-                format!( "The version of the proc-macro server ({v}) in your Rust toolchain is newer than the version supported by your rust-analyzer ({CURRENT_API_VERSION}).
-            This will prevent proc-macro expansion from working. Please consider updating your rust-analyzer to ensure compatibility with your current toolchain."
-                ),
-            )),
+            Ok(v) if v > version::CURRENT_API_VERSION => Err(io::Error::other(
+                format!( "The version of the proc-macro server ({v}) in your Rust toolchain is newer than the version supported by your rust-analyzer ({}).
+                        This will prevent proc-macro expansion from working. Please consider updating your rust-analyzer to ensure compatibility with your current toolchain."
+                        ,version::CURRENT_API_VERSION
+                    ),
+                )),
             Ok(v) => {
                 tracing::info!("Proc-macro server version: {v}");
                 srv.version = v;
-                if srv.version >= RUST_ANALYZER_SPAN_SUPPORT {
-                    if let Ok(mode) = srv.enable_rust_analyzer_spans() {
-                        srv.mode = mode;
-                    }
+                if srv.version >= version::RUST_ANALYZER_SPAN_SUPPORT && let Ok(mode) = srv.enable_rust_analyzer_spans() {
+                    srv.protocol = Protocol::LegacyJson { mode };
                 }
-                tracing::info!("Proc-macro server span mode: {:?}", srv.mode);
+                tracing::info!("Proc-macro server protocol: {:?}", srv.protocol);
                 Ok(srv)
             }
             Err(e) => {
@@ -98,25 +97,24 @@ impl ProcMacroServerProcess {
         self.version
     }
 
+    /// Enable support for rust-analyzer span mode if the server supports it.
+    pub(crate) fn rust_analyzer_spans(&self) -> bool {
+        match self.protocol {
+            Protocol::LegacyJson { mode } => mode == SpanMode::RustAnalyzer,
+        }
+    }
+
     /// Checks the API version of the running proc-macro server.
     fn version_check(&self) -> Result<u32, ServerError> {
-        let request = Request::ApiVersionCheck {};
-        let response = self.send_task(request)?;
-
-        match response {
-            Response::ApiVersionCheck(version) => Ok(version),
-            _ => Err(ServerError { message: "unexpected response".to_owned(), io: None }),
+        match self.protocol {
+            Protocol::LegacyJson { .. } => legacy_protocol::version_check(self),
         }
     }
 
     /// Enable support for rust-analyzer span mode if the server supports it.
     fn enable_rust_analyzer_spans(&self) -> Result<SpanMode, ServerError> {
-        let request = Request::SetConfig(ServerConfig { span_mode: SpanMode::RustAnalyzer });
-        let response = self.send_task(request)?;
-
-        match response {
-            Response::SetConfig(ServerConfig { span_mode }) => Ok(span_mode),
-            _ => Err(ServerError { message: "unexpected response".to_owned(), io: None }),
+        match self.protocol {
+            Protocol::LegacyJson { .. } => legacy_protocol::enable_rust_analyzer_spans(self),
         }
     }
 
@@ -125,25 +123,24 @@ impl ProcMacroServerProcess {
         &self,
         dylib_path: &AbsPath,
     ) -> Result<Result<Vec<(String, ProcMacroKind)>, String>, ServerError> {
-        let request = Request::ListMacros { dylib_path: dylib_path.to_path_buf().into() };
-
-        let response = self.send_task(request)?;
-
-        match response {
-            Response::ListMacros(it) => Ok(it),
-            _ => Err(ServerError { message: "unexpected response".to_owned(), io: None }),
+        match self.protocol {
+            Protocol::LegacyJson { .. } => legacy_protocol::find_proc_macros(self, dylib_path),
         }
     }
 
-    /// Sends a request to the proc-macro server and waits for a response.
-    pub(crate) fn send_task(&self, req: Request) -> Result<Response, ServerError> {
-        if let Some(server_error) = self.exited.get() {
-            return Err(server_error.0.clone());
-        }
-
+    pub(crate) fn send_task<Request, Response>(
+        &self,
+        serialize_req: impl FnOnce(
+            &mut dyn Write,
+            &mut dyn BufRead,
+            Request,
+            &mut String,
+        ) -> Result<Option<Response>, ServerError>,
+        req: Request,
+    ) -> Result<Response, ServerError> {
         let state = &mut *self.state.lock().unwrap();
         let mut buf = String::new();
-        send_request(&mut state.stdin, &mut state.stdout, req, &mut buf)
+        serialize_req(&mut state.stdin, &mut state.stdout, req, &mut buf)
             .and_then(|res| {
                 res.ok_or_else(|| {
                     let message = "proc-macro server did not respond with data".to_owned();
@@ -162,10 +159,10 @@ impl ProcMacroServerProcess {
                         Ok(None) | Err(_) => e,
                         Ok(Some(status)) => {
                             let mut msg = String::new();
-                            if !status.success() {
-                                if let Some(stderr) = state.process.child.stderr.as_mut() {
-                                    _ = stderr.read_to_string(&mut msg);
-                                }
+                            if !status.success()
+                                && let Some(stderr) = state.process.child.stderr.as_mut()
+                            {
+                                _ = stderr.read_to_string(&mut msg);
                             }
                             let server_error = ServerError {
                                 message: format!(
@@ -242,21 +239,3 @@ fn mk_child<'a>(
     }
     cmd.spawn()
 }
-
-/// Sends a request to the server and reads the response.
-fn send_request(
-    mut writer: &mut impl Write,
-    mut reader: &mut impl BufRead,
-    req: Request,
-    buf: &mut String,
-) -> Result<Option<Response>, ServerError> {
-    req.write(write_json, &mut writer).map_err(|err| ServerError {
-        message: "failed to write request".into(),
-        io: Some(Arc::new(err)),
-    })?;
-    let res = Response::read(read_json, &mut reader, buf).map_err(|err| ServerError {
-        message: "failed to read response".into(),
-        io: Some(Arc::new(err)),
-    })?;
-    Ok(res)
-}