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.

10 files changed, 96 insertions, 102 deletions

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 4e1526fe48b..165936269d3 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
@@ -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 {
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 25c30b6db4a..516c7418bde 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
@@ -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 }),
         }
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 e3b0614546d..fcea75ef672 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
@@ -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 {
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv-cli/src/main_loop.rs b/src/tools/rust-analyzer/crates/proc-macro-srv-cli/src/main_loop.rs
index 25a49fbff9f..f54dff1f2d8 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv-cli/src/main_loop.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv-cli/src/main_loop.rs
@@ -26,7 +26,7 @@ pub(crate) fn run() -> io::Result<()> {
     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 srv = proc_macro_srv::ProcMacroSrv::new(&env);
 
     let mut span_mode = SpanMode::Id;
 
@@ -78,7 +78,7 @@ pub(crate) fn run() -> io::Result<()> {
                         })
                         .map_err(msg::PanicMessage)
                     }),
-                    SpanMode::RustAnalyzer { .. } => msg::Response::ExpandMacroExtended({
+                    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];
@@ -122,9 +122,6 @@ pub(crate) fn run() -> io::Result<()> {
             msg::Request::ApiVersionCheck {} => msg::Response::ApiVersionCheck(CURRENT_API_VERSION),
             msg::Request::SetConfig(config) => {
                 span_mode = config.span_mode;
-                if let SpanMode::RustAnalyzer { fixup_ast_id } = span_mode {
-                    srv.set_fixup_ast_id(fixup_ast_id);
-                }
                 msg::Response::SetConfig(config)
             }
         };
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib.rs
index c64a9833e3f..c49159df991 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/dylib.rs
@@ -3,7 +3,6 @@
 mod version;
 
 use proc_macro::bridge;
-use span::ErasedFileAstId;
 use std::{fmt, fs, io, time::SystemTime};
 
 use libloading::Library;
@@ -162,20 +161,14 @@ impl Expander {
         def_site: S,
         call_site: S,
         mixed_site: S,
-        fixup_ast_id: ErasedFileAstId,
     ) -> Result<TopSubtree<S>, String>
     where
         <S::Server as bridge::server::Types>::TokenStream: Default,
     {
-        let result = self.inner.proc_macros.expand(
-            macro_name,
-            macro_body,
-            attributes,
-            def_site,
-            call_site,
-            mixed_site,
-            fixup_ast_id,
-        );
+        let result = self
+            .inner
+            .proc_macros
+            .expand(macro_name, macro_body, attributes, def_site, call_site, mixed_site);
         result.map_err(|e| e.into_string().unwrap_or_default())
     }
 
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/lib.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/lib.rs
index 22afa018de0..223c5a54b70 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/lib.rs
@@ -41,7 +41,7 @@ use std::{
 };
 
 use paths::{Utf8Path, Utf8PathBuf};
-use span::{ErasedFileAstId, FIXUP_ERASED_FILE_AST_ID_MARKER, Span, TokenId};
+use span::{Span, TokenId};
 
 use crate::server_impl::TokenStream;
 
@@ -57,16 +57,11 @@ pub const RUSTC_VERSION_STRING: &str = env!("RUSTC_VERSION");
 pub struct ProcMacroSrv<'env> {
     expanders: Mutex<HashMap<Utf8PathBuf, Arc<dylib::Expander>>>,
     env: &'env EnvSnapshot,
-    fixup_ast_id: ErasedFileAstId,
 }
 
 impl<'env> ProcMacroSrv<'env> {
     pub fn new(env: &'env EnvSnapshot) -> Self {
-        Self { expanders: Default::default(), env, fixup_ast_id: FIXUP_ERASED_FILE_AST_ID_MARKER }
-    }
-
-    pub fn set_fixup_ast_id(&mut self, fixup_ast_id: u32) {
-        self.fixup_ast_id = ErasedFileAstId::from_raw(fixup_ast_id);
+        Self { expanders: Default::default(), env }
     }
 }
 
@@ -106,7 +101,6 @@ impl ProcMacroSrv<'_> {
                             def_site,
                             call_site,
                             mixed_site,
-                            self.fixup_ast_id,
                         )
                         .map(|tt| tt.0)
                 });
@@ -162,41 +156,25 @@ impl ProcMacroSrv<'_> {
 
 pub trait ProcMacroSrvSpan: Copy + Send {
     type Server: proc_macro::bridge::server::Server<TokenStream = TokenStream<Self>>;
-    fn make_server(
-        call_site: Self,
-        def_site: Self,
-        mixed_site: Self,
-        fixup_ast_id: ErasedFileAstId,
-    ) -> Self::Server;
+    fn make_server(call_site: Self, def_site: Self, mixed_site: Self) -> Self::Server;
 }
 
 impl ProcMacroSrvSpan for TokenId {
     type Server = server_impl::token_id::TokenIdServer;
 
-    fn make_server(
-        call_site: Self,
-        def_site: Self,
-        mixed_site: Self,
-        _fixup_ast_id: ErasedFileAstId,
-    ) -> Self::Server {
+    fn make_server(call_site: Self, def_site: Self, mixed_site: Self) -> Self::Server {
         Self::Server { call_site, def_site, mixed_site }
     }
 }
 impl ProcMacroSrvSpan for Span {
     type Server = server_impl::rust_analyzer_span::RaSpanServer;
-    fn make_server(
-        call_site: Self,
-        def_site: Self,
-        mixed_site: Self,
-        fixup_ast_id: ErasedFileAstId,
-    ) -> Self::Server {
+    fn make_server(call_site: Self, def_site: Self, mixed_site: Self) -> Self::Server {
         Self::Server {
             call_site,
             def_site,
             mixed_site,
             tracked_env_vars: Default::default(),
             tracked_paths: Default::default(),
-            fixup_ast_id,
         }
     }
 }
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/proc_macros.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/proc_macros.rs
index 3b9c45894c6..18532706c4a 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/proc_macros.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/proc_macros.rs
@@ -1,7 +1,6 @@
 //! Proc macro ABI
 
 use proc_macro::bridge;
-use span::ErasedFileAstId;
 
 use crate::{ProcMacroKind, ProcMacroSrvSpan, server_impl::TopSubtree};
 
@@ -23,7 +22,6 @@ impl ProcMacros {
         def_site: S,
         call_site: S,
         mixed_site: S,
-        fixup_ast_id: ErasedFileAstId,
     ) -> Result<TopSubtree<S>, crate::PanicMessage> {
         let parsed_body = crate::server_impl::TokenStream::with_subtree(macro_body);
 
@@ -39,7 +37,7 @@ impl ProcMacros {
                 {
                     let res = client.run(
                         &bridge::server::SameThread,
-                        S::make_server(call_site, def_site, mixed_site, fixup_ast_id),
+                        S::make_server(call_site, def_site, mixed_site),
                         parsed_body,
                         cfg!(debug_assertions),
                     );
@@ -50,7 +48,7 @@ impl ProcMacros {
                 bridge::client::ProcMacro::Bang { name, client } if *name == macro_name => {
                     let res = client.run(
                         &bridge::server::SameThread,
-                        S::make_server(call_site, def_site, mixed_site, fixup_ast_id),
+                        S::make_server(call_site, def_site, mixed_site),
                         parsed_body,
                         cfg!(debug_assertions),
                     );
@@ -61,7 +59,7 @@ impl ProcMacros {
                 bridge::client::ProcMacro::Attr { name, client } if *name == macro_name => {
                     let res = client.run(
                         &bridge::server::SameThread,
-                        S::make_server(call_site, def_site, mixed_site, fixup_ast_id),
+                        S::make_server(call_site, def_site, mixed_site),
                         parsed_attributes,
                         parsed_body,
                         cfg!(debug_assertions),
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/rust_analyzer_span.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/rust_analyzer_span.rs
index b071ad80f5b..a1863efafbb 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/rust_analyzer_span.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/server_impl/rust_analyzer_span.rs
@@ -11,7 +11,7 @@ use std::{
 
 use intern::Symbol;
 use proc_macro::bridge::{self, server};
-use span::{ErasedFileAstId, Span};
+use span::{FIXUP_ERASED_FILE_AST_ID_MARKER, Span};
 use tt::{TextRange, TextSize};
 
 use crate::server_impl::{from_token_tree, literal_from_str, token_stream::TokenStreamBuilder};
@@ -28,7 +28,6 @@ pub struct RaSpanServer {
     pub call_site: Span,
     pub def_site: Span,
     pub mixed_site: Span,
-    pub fixup_ast_id: ErasedFileAstId,
 }
 
 impl server::Types for RaSpanServer {
@@ -182,10 +181,10 @@ impl server::Span for RaSpanServer {
     fn join(&mut self, first: Self::Span, second: Self::Span) -> Option<Self::Span> {
         // We can't modify the span range for fixup spans, those are meaningful to fixup, so just
         // prefer the non-fixup span.
-        if first.anchor.ast_id == self.fixup_ast_id {
+        if first.anchor.ast_id == FIXUP_ERASED_FILE_AST_ID_MARKER {
             return Some(second);
         }
-        if second.anchor.ast_id == self.fixup_ast_id {
+        if second.anchor.ast_id == FIXUP_ERASED_FILE_AST_ID_MARKER {
             return Some(first);
         }
         // FIXME: Once we can talk back to the client, implement a "long join" request for anchors
@@ -214,7 +213,7 @@ impl server::Span for RaSpanServer {
         end: Bound<usize>,
     ) -> Option<Self::Span> {
         // We can't modify the span range for fixup spans, those are meaningful to fixup.
-        if span.anchor.ast_id == self.fixup_ast_id {
+        if span.anchor.ast_id == FIXUP_ERASED_FILE_AST_ID_MARKER {
             return Some(span);
         }
         let length = span.range.len().into();
@@ -257,7 +256,7 @@ impl server::Span for RaSpanServer {
 
     fn end(&mut self, span: Self::Span) -> Self::Span {
         // We can't modify the span range for fixup spans, those are meaningful to fixup.
-        if span.anchor.ast_id == self.fixup_ast_id {
+        if span.anchor.ast_id == FIXUP_ERASED_FILE_AST_ID_MARKER {
             return span;
         }
         Span { range: TextRange::empty(span.range.end()), ..span }
@@ -265,7 +264,7 @@ impl server::Span for RaSpanServer {
 
     fn start(&mut self, span: Self::Span) -> Self::Span {
         // We can't modify the span range for fixup spans, those are meaningful to fixup.
-        if span.anchor.ast_id == self.fixup_ast_id {
+        if span.anchor.ast_id == FIXUP_ERASED_FILE_AST_ID_MARKER {
             return span;
         }
         Span { range: TextRange::empty(span.range.start()), ..span }
diff --git a/src/tools/rust-analyzer/crates/proc-macro-srv/src/tests/utils.rs b/src/tools/rust-analyzer/crates/proc-macro-srv/src/tests/utils.rs
index e0a69608dec..10af5662b5c 100644
--- a/src/tools/rust-analyzer/crates/proc-macro-srv/src/tests/utils.rs
+++ b/src/tools/rust-analyzer/crates/proc-macro-srv/src/tests/utils.rs
@@ -2,8 +2,7 @@
 
 use expect_test::Expect;
 use span::{
-    EditionedFileId, FIXUP_ERASED_FILE_AST_ID_MARKER, FileId, ROOT_ERASED_FILE_AST_ID, Span,
-    SpanAnchor, SyntaxContext, TokenId,
+    EditionedFileId, FileId, ROOT_ERASED_FILE_AST_ID, Span, SpanAnchor, SyntaxContext, TokenId,
 };
 use tt::TextRange;
 
@@ -68,17 +67,8 @@ fn assert_expand_impl(
     let input_ts_string = format!("{input_ts:?}");
     let attr_ts_string = attr_ts.as_ref().map(|it| format!("{it:?}"));
 
-    let res = expander
-        .expand(
-            macro_name,
-            input_ts,
-            attr_ts,
-            def_site,
-            call_site,
-            mixed_site,
-            FIXUP_ERASED_FILE_AST_ID_MARKER,
-        )
-        .unwrap();
+    let res =
+        expander.expand(macro_name, input_ts, attr_ts, def_site, call_site, mixed_site).unwrap();
     expect.assert_eq(&format!(
         "{input_ts_string}\n\n{}\n\n{res:?}",
         attr_ts_string.unwrap_or_default()
@@ -110,17 +100,7 @@ fn assert_expand_impl(
     let fixture_string = format!("{fixture:?}");
     let attr_string = attr.as_ref().map(|it| format!("{it:?}"));
 
-    let res = expander
-        .expand(
-            macro_name,
-            fixture,
-            attr,
-            def_site,
-            call_site,
-            mixed_site,
-            FIXUP_ERASED_FILE_AST_ID_MARKER,
-        )
-        .unwrap();
+    let res = expander.expand(macro_name, fixture, attr, def_site, call_site, mixed_site).unwrap();
     expect_spanned
         .assert_eq(&format!("{fixture_string}\n\n{}\n\n{res:#?}", attr_string.unwrap_or_default()));
 }
diff --git a/src/tools/rust-analyzer/crates/span/src/ast_id.rs b/src/tools/rust-analyzer/crates/span/src/ast_id.rs
index 51a693ca011..8e959711981 100644
--- a/src/tools/rust-analyzer/crates/span/src/ast_id.rs
+++ b/src/tools/rust-analyzer/crates/span/src/ast_id.rs
@@ -108,7 +108,8 @@ impl fmt::Debug for ErasedFileAstId {
 
 #[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
 enum ErasedFileAstIdKind {
-    Root,
+    /// This needs to not change because it's depended upon by the proc macro server.
+    Fixup,
     // The following are associated with `ErasedHasNameFileAstId`.
     Enum,
     Struct,
@@ -143,7 +144,7 @@ enum ErasedFileAstIdKind {
     /// Associated with [`BlockExprFileAstId`].
     BlockExpr,
     /// Keep this last.
-    Fixup,
+    Root,
 }
 
 // First hash, then index, then kind.
@@ -218,7 +219,7 @@ impl ErasedFileAstId {
     }
 
     #[inline]
-    pub fn from_raw(v: u32) -> Self {
+    pub const fn from_raw(v: u32) -> Self {
         Self(v)
     }
 }