Support spans with proc macro servers from before the ast id changes

The only thing changed is the value of the fixup ast id, so we just swap it.

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

@@ -55,7 +55,7 @@ pub enum SpanMode {
     Id,
 
     /// Rust Analyzer-specific span handling mode.
-    RustAnalyzer { fixup_ast_id: u32 },
+    RustAnalyzer,
 }
 
 /// Represents responses sent from the proc-macro-srv to the client.
@@ -308,7 +308,7 @@ mod tests {
     #[test]
     fn test_proc_macro_rpc_works() {
         let tt = fixture_token_tree();
-        for v in HASHED_AST_ID..=CURRENT_API_VERSION {
+        for v in RUST_ANALYZER_SPAN_SUPPORT..=CURRENT_API_VERSION {
             let mut span_data_table = Default::default();
             let task = ExpandMacro {
                 data: ExpandMacroData {

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

@@ -12,13 +12,13 @@ pub mod legacy_protocol {
 mod process;
 
 use paths::{AbsPath, AbsPathBuf};
-use span::Span;
+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, PanicMessage,
-        RUST_ANALYZER_SPAN_SUPPORT, Request, Response, SpanDataIndexMap,
+        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,
@@ -161,6 +161,38 @@ impl ProcMacro {
         self.kind
     }
 
+    fn needs_fixup_change(&self) -> bool {
+        let version = self.process.version();
+        (RUST_ANALYZER_SPAN_SUPPORT..HASHED_AST_ID).contains(&version)
+    }
+
+    /// On some server versions, the fixup ast id is different than ours. So change it to match.
+    fn change_fixup_to_match_old_server(&self, tt: &mut tt::TopSubtree<Span>) {
+        const OLD_FIXUP_AST_ID: ErasedFileAstId = ErasedFileAstId::from_raw(!0 - 1);
+        let change_ast_id = |ast_id: &mut ErasedFileAstId| {
+            if *ast_id == FIXUP_ERASED_FILE_AST_ID_MARKER {
+                *ast_id = OLD_FIXUP_AST_ID;
+            } else if *ast_id == OLD_FIXUP_AST_ID {
+                // Swap between them, that means no collision plus the change can be reversed by doing itself.
+                *ast_id = FIXUP_ERASED_FILE_AST_ID_MARKER;
+            }
+        };
+
+        for tt in &mut tt.0 {
+            match tt {
+                tt::TokenTree::Leaf(tt::Leaf::Ident(tt::Ident { span, .. }))
+                | tt::TokenTree::Leaf(tt::Leaf::Literal(tt::Literal { span, .. }))
+                | tt::TokenTree::Leaf(tt::Leaf::Punct(tt::Punct { span, .. })) => {
+                    change_ast_id(&mut span.anchor.ast_id);
+                }
+                tt::TokenTree::Subtree(subtree) => {
+                    change_ast_id(&mut subtree.delimiter.open.anchor.ast_id);
+                    change_ast_id(&mut subtree.delimiter.close.anchor.ast_id);
+                }
+            }
+        }
+    }
+
     /// Expands the procedural macro by sending an expansion request to the server.
     /// This includes span information and environmental context.
     pub fn expand(
@@ -173,6 +205,20 @@ impl ProcMacro {
         mixed_site: Span,
         current_dir: String,
     ) -> Result<Result<tt::TopSubtree<Span>, PanicMessage>, ServerError> {
+        let (mut subtree, mut attr) = (subtree, attr);
+        let (mut subtree_changed, mut attr_changed);
+        if self.needs_fixup_change() {
+            subtree_changed = tt::TopSubtree::from_subtree(subtree);
+            self.change_fixup_to_match_old_server(&mut subtree_changed);
+            subtree = subtree_changed.view();
+
+            if let Some(attr) = &mut attr {
+                attr_changed = tt::TopSubtree::from_subtree(*attr);
+                self.change_fixup_to_match_old_server(&mut attr_changed);
+                *attr = attr_changed.view();
+            }
+        }
+
         let version = self.process.version();
 
         let mut span_data_table = SpanDataIndexMap::default();
@@ -205,15 +251,23 @@ impl ProcMacro {
         let response = self.process.send_task(Request::ExpandMacro(Box::new(task)))?;
 
         match response {
-            Response::ExpandMacro(it) => {
-                Ok(it.map(|tree| FlatTree::to_subtree_resolved(tree, version, &span_data_table)))
-            }
+            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| {
-                FlatTree::to_subtree_resolved(
+                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 }),
         }

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

@@ -8,7 +8,6 @@ use std::{
 };
 
 use paths::AbsPath;
-use span::FIXUP_ERASED_FILE_AST_ID_MARKER;
 use stdx::JodChild;
 
 use crate::{
@@ -16,7 +15,8 @@ use crate::{
     legacy_protocol::{
         json::{read_json, write_json},
         msg::{
-            CURRENT_API_VERSION, HASHED_AST_ID, Message, Request, Response, ServerConfig, SpanMode,
+            CURRENT_API_VERSION, Message, RUST_ANALYZER_SPAN_SUPPORT, Request, Response,
+            ServerConfig, SpanMode,
         },
     },
 };
@@ -71,9 +71,7 @@ impl ProcMacroServerProcess {
             Ok(v) => {
                 tracing::info!("Proc-macro server version: {v}");
                 srv.version = v;
-                if srv.version >= HASHED_AST_ID {
-                    // We don't enable spans on versions prior to `HASHED_AST_ID`, because their ast id layout
-                    // is different.
+                if srv.version >= RUST_ANALYZER_SPAN_SUPPORT {
                     if let Ok(mode) = srv.enable_rust_analyzer_spans() {
                         srv.mode = mode;
                     }
@@ -113,11 +111,7 @@ impl ProcMacroServerProcess {
 
     /// 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 {
-                fixup_ast_id: FIXUP_ERASED_FILE_AST_ID_MARKER.into_raw(),
-            },
-        });
+        let request = Request::SetConfig(ServerConfig { span_mode: SpanMode::RustAnalyzer });
         let response = self.send_task(request)?;
 
         match response {