Merge pull request #19985 from ChayimFriedman2/proc-macro-srv-ast-id

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

Author: Veykril

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,15 +8,15 @@ use std::{
 };
 
 use paths::AbsPath;
-use span::FIXUP_ERASED_FILE_AST_ID_MARKER;
 use stdx::JodChild;
 
 use crate::{
     ProcMacroKind, ServerError,
     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 {

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)
             }
         };

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())
     }
 

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,
         }
     }
 }

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),