Add rust-analyzer-span server feature equivalent to the ID server

Author: Veykril

Cargo.lock

@@ -20,7 +20,10 @@ use triomphe::Arc;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    msg::{ExpandMacro, ExpnGlobals, FlatTree, PanicMessage, HAS_GLOBAL_SPANS},
+    msg::{
+        flat::serialize_span_data_index_map, ExpandMacro, ExpnGlobals, FlatTree, PanicMessage,
+        HAS_GLOBAL_SPANS, RUST_ANALYZER_SPAN_SUPPORT,
+    },
     process::ProcMacroProcessSrv,
 };
 
@@ -166,6 +169,11 @@ impl ProcMacro {
                 call_site,
                 mixed_site,
             },
+            span_data_table: if version >= RUST_ANALYZER_SPAN_SUPPORT {
+                serialize_span_data_index_map(&span_data_table)
+            } else {
+                Vec::new()
+            },
         };
 
         let response = self
@@ -178,9 +186,7 @@ impl ProcMacro {
             msg::Response::ExpandMacro(it) => {
                 Ok(it.map(|tree| FlatTree::to_subtree_resolved(tree, version, &span_data_table)))
             }
-            msg::Response::ListMacros(..) | msg::Response::ApiVersionCheck(..) => {
-                Err(ServerError { message: "unexpected response".to_string(), io: None })
-            }
+            _ => Err(ServerError { message: "unexpected response".to_string(), io: None }),
         }
     }
 }

crates/proc-macro-api/src/lib.rs

@@ -10,28 +10,42 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 use crate::ProcMacroKind;
 
-pub use crate::msg::flat::{FlatTree, TokenId};
+pub use crate::msg::flat::{
+    deserialize_span_data_index_map, serialize_span_data_index_map, FlatTree, SpanDataIndexMap,
+    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;
 
-pub const CURRENT_API_VERSION: u32 = HAS_GLOBAL_SPANS;
+pub const CURRENT_API_VERSION: u32 = RUST_ANALYZER_SPAN_SUPPORT;
 
 #[derive(Debug, Serialize, Deserialize)]
 pub enum Request {
     ListMacros { dylib_path: PathBuf },
     ExpandMacro(ExpandMacro),
+    SetSpanMode(SpanMode),
     ApiVersionCheck {},
 }
 
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
+pub enum SpanMode {
+    #[default]
+    Id,
+    RustAnalyzer,
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub enum Response {
     ListMacros(Result<Vec<(String, ProcMacroKind)>, String>),
     ExpandMacro(Result<FlatTree, PanicMessage>),
+    ExpandMacroSpans(Result<(FlatTree, Vec<u32>), PanicMessage>),
     ApiVersionCheck(u32),
+    SetSpanMode(SpanMode),
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -64,9 +78,12 @@ pub struct ExpandMacro {
     #[serde(skip_serializing_if = "ExpnGlobals::skip_serializing_if")]
     #[serde(default)]
     pub has_global_spans: ExpnGlobals,
+    #[serde(skip_serializing_if = "Vec::is_empty")]
+    #[serde(default)]
+    pub span_data_table: Vec<u32>,
 }
 
-#[derive(Default, Debug, Serialize, Deserialize)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
 pub struct ExpnGlobals {
     #[serde(skip_serializing)]
     #[serde(default)]
@@ -241,6 +258,7 @@ mod tests {
                 call_site: 0,
                 mixed_site: 0,
             },
+            span_data_table: Vec::new(),
         };
 
         let json = serde_json::to_string(&task).unwrap();

crates/proc-macro-api/src/msg.rs

@@ -38,12 +38,45 @@
 use std::collections::{HashMap, VecDeque};
 
 use indexmap::IndexSet;
+use la_arena::RawIdx;
 use serde::{Deserialize, Serialize};
-use span::Span;
+use span::{ErasedFileAstId, FileId, Span, SpanAnchor, SyntaxContextId};
+use text_size::TextRange;
 
 use crate::msg::ENCODE_CLOSE_SPAN_VERSION;
 
-type SpanIndexMap = IndexSet<Span>;
+pub type SpanDataIndexMap = IndexSet<Span>;
+
+pub fn serialize_span_data_index_map(map: &SpanDataIndexMap) -> Vec<u32> {
+    map.iter()
+        .flat_map(|span| {
+            [
+                span.anchor.file_id.index(),
+                span.anchor.ast_id.into_raw().into_u32(),
+                span.range.start().into(),
+                span.range.end().into(),
+                span.ctx.into_u32(),
+            ]
+        })
+        .collect()
+}
+
+pub fn deserialize_span_data_index_map(map: &[u32]) -> SpanDataIndexMap {
+    debug_assert!(map.len() % 5 == 0);
+    map.chunks_exact(5)
+        .map(|span| {
+            let &[file_id, ast_id, start, end, e] = span else { unreachable!() };
+            Span {
+                anchor: SpanAnchor {
+                    file_id: FileId::from_raw(file_id),
+                    ast_id: ErasedFileAstId::from_raw(RawIdx::from_u32(ast_id)),
+                },
+                range: TextRange::new(start.into(), end.into()),
+                ctx: SyntaxContextId::from_u32(e),
+            }
+        })
+        .collect()
+}
 
 #[derive(Clone, Copy, PartialEq, Eq, Hash)]
 pub struct TokenId(pub u32);
@@ -93,7 +126,7 @@ impl FlatTree {
     pub fn new(
         subtree: &tt::Subtree<Span>,
         version: u32,
-        span_data_table: &mut SpanIndexMap,
+        span_data_table: &mut SpanDataIndexMap,
     ) -> FlatTree {
         let mut w = Writer {
             string_table: HashMap::new(),
@@ -155,7 +188,7 @@ impl FlatTree {
     pub fn to_subtree_resolved(
         self,
         version: u32,
-        span_data_table: &SpanIndexMap,
+        span_data_table: &SpanDataIndexMap,
     ) -> tt::Subtree<Span> {
         Reader {
             subtree: if version >= ENCODE_CLOSE_SPAN_VERSION {

crates/proc-macro-api/src/msg/flat.rs

@@ -9,7 +9,7 @@ use paths::{AbsPath, AbsPathBuf};
 use stdx::JodChild;
 
 use crate::{
-    msg::{Message, Request, Response, CURRENT_API_VERSION},
+    msg::{Message, Request, Response, SpanMode, CURRENT_API_VERSION, RUST_ANALYZER_SPAN_SUPPORT},
     ProcMacroKind, ServerError,
 };
 
@@ -19,6 +19,7 @@ pub(crate) struct ProcMacroProcessSrv {
     stdin: ChildStdin,
     stdout: BufReader<ChildStdout>,
     version: u32,
+    mode: SpanMode,
 }
 
 impl ProcMacroProcessSrv {
@@ -27,7 +28,13 @@ impl ProcMacroProcessSrv {
             let mut process = Process::run(process_path.clone(), null_stderr)?;
             let (stdin, stdout) = process.stdio().expect("couldn't access child stdio");
 
-            io::Result::Ok(ProcMacroProcessSrv { _process: process, stdin, stdout, version: 0 })
+            io::Result::Ok(ProcMacroProcessSrv {
+                _process: process,
+                stdin,
+                stdout,
+                version: 0,
+                mode: SpanMode::Id,
+            })
         };
         let mut srv = create_srv(true)?;
         tracing::info!("sending version check");
@@ -43,6 +50,11 @@ impl ProcMacroProcessSrv {
                 tracing::info!("got version {v}");
                 srv = create_srv(false)?;
                 srv.version = v;
+                if srv.version > RUST_ANALYZER_SPAN_SUPPORT {
+                    if let Ok(mode) = srv.enable_rust_analyzer_spans() {
+                        srv.mode = mode;
+                    }
+                }
                 Ok(srv)
             }
             Err(e) => {
@@ -62,9 +74,17 @@ impl ProcMacroProcessSrv {
 
         match response {
             Response::ApiVersionCheck(version) => Ok(version),
-            Response::ExpandMacro { .. } | Response::ListMacros { .. } => {
-                Err(ServerError { message: "unexpected response".to_string(), io: None })
-            }
+            _ => Err(ServerError { message: "unexpected response".to_string(), io: None }),
+        }
+    }
+
+    fn enable_rust_analyzer_spans(&mut self) -> Result<SpanMode, ServerError> {
+        let request = Request::SetSpanMode(crate::msg::SpanMode::RustAnalyzer);
+        let response = self.send_task(request)?;
+
+        match response {
+            Response::SetSpanMode(span_mode) => Ok(span_mode),
+            _ => Err(ServerError { message: "unexpected response".to_string(), io: None }),
         }
     }
 
@@ -78,9 +98,7 @@ impl ProcMacroProcessSrv {
 
         match response {
             Response::ListMacros(it) => Ok(it),
-            Response::ExpandMacro { .. } | Response::ApiVersionCheck { .. } => {
-                Err(ServerError { message: "unexpected response".to_string(), io: None })
-            }
+            _ => Err(ServerError { message: "unexpected response".to_string(), io: None }),
         }
     }
 

crates/proc-macro-api/src/process.rs

@@ -39,10 +39,17 @@ fn run() -> io::Result<()> {
             msg::Request::ListMacros { dylib_path } => {
                 msg::Response::ListMacros(srv.list_macros(&dylib_path))
             }
-            msg::Request::ExpandMacro(task) => msg::Response::ExpandMacro(srv.expand(task)),
+            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::ExpandMacroSpans(srv.expand(task)),
+            },
             msg::Request::ApiVersionCheck {} => {
                 msg::Response::ApiVersionCheck(proc_macro_api::msg::CURRENT_API_VERSION)
             }
+            msg::Request::SetSpanMode(span_mode) => {
+                srv.set_span_mode(span_mode);
+                msg::Response::SetSpanMode(span_mode)
+            }
         };
         write_response(res)?
     }

crates/proc-macro-srv-cli/src/main.rs

@@ -26,6 +26,8 @@ stdx.workspace = true
 tt.workspace = true
 mbe.workspace = true
 paths.workspace = true
+base-db.workspace = true
+span.workspace = true
 proc-macro-api.workspace = true
 
 [dev-dependencies]

crates/proc-macro-srv/Cargo.toml

@@ -11,7 +11,10 @@ use libloading::Library;
 use memmap2::Mmap;
 use object::Object;
 use paths::AbsPath;
-use proc_macro_api::{msg::TokenId, read_dylib_info, ProcMacroKind};
+use proc_macro::bridge;
+use proc_macro_api::{read_dylib_info, ProcMacroKind};
+
+use crate::ProcMacroSrvSpan;
 
 const NEW_REGISTRAR_SYMBOL: &str = "_rustc_proc_macro_decls_";
 
@@ -147,15 +150,18 @@ impl Expander {
         Ok(Expander { inner: library })
     }
 
-    pub fn expand(
+    pub fn expand<S: ProcMacroSrvSpan>(
         &self,
         macro_name: &str,
-        macro_body: &crate::tt::Subtree,
-        attributes: Option<&crate::tt::Subtree>,
-        def_site: TokenId,
-        call_site: TokenId,
-        mixed_site: TokenId,
-    ) -> Result<crate::tt::Subtree, String> {
+        macro_body: tt::Subtree<S>,
+        attributes: Option<tt::Subtree<S>>,
+        def_site: S,
+        call_site: S,
+        mixed_site: S,
+    ) -> Result<tt::Subtree<S>, String>
+    where
+        <S::Server as bridge::server::Types>::TokenStream: Default,
+    {
         let result = self
             .inner
             .proc_macros