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

4 files changed, 94 insertions, 19 deletions

diff --git a/crates/proc-macro-api/src/lib.rs b/crates/proc-macro-api/src/lib.rs
index a5d4cfeb275..be68f6eff66 100644
--- a/crates/proc-macro-api/src/lib.rs
+++ b/crates/proc-macro-api/src/lib.rs
@@ -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 }),
         }
     }
 }
diff --git a/crates/proc-macro-api/src/msg.rs b/crates/proc-macro-api/src/msg.rs
index 18fd9ed7284..7ef6d5b0556 100644
--- a/crates/proc-macro-api/src/msg.rs
+++ b/crates/proc-macro-api/src/msg.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();
diff --git a/crates/proc-macro-api/src/msg/flat.rs b/crates/proc-macro-api/src/msg/flat.rs
index a12581ac138..8dfaba52625 100644
--- a/crates/proc-macro-api/src/msg/flat.rs
+++ b/crates/proc-macro-api/src/msg/flat.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 {
diff --git a/crates/proc-macro-api/src/process.rs b/crates/proc-macro-api/src/process.rs
index 9a20fa63ed7..ed55ca5bd53 100644
--- a/crates/proc-macro-api/src/process.rs
+++ b/crates/proc-macro-api/src/process.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 }),
         }
     }