Cleanup bidirection proc-macro-srv protocol a bit

Author: Veykril

crates/load-cargo/src/lib.rs

@@ -724,6 +724,8 @@ impl ProcMacroExpander for Expander {
 
                 Ok(SubResponse::SpanParentResult { parent_span: None })
             }
+            // FIXME: implement this
+            SubRequest::SpanJoin { .. } => Ok(SubResponse::SpanJoinResult { span: None }),
         };
         match self.0.expand(
             subtree.view(),

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

@@ -12,8 +12,8 @@ use span::Span;
 use crate::{
     ProcMacro, ProcMacroKind, ServerError,
     bidirectional_protocol::msg::{
-        BidirectionalMessage, ExpandMacro, ExpandMacroData, ExpnGlobals, Request, Response,
-        SubRequest, SubResponse,
+        ApiVersionCheck, BidirectionalMessage, ExpandMacro, ExpandMacroData, ExpnGlobals,
+        ListMacros, Request, Response, SubRequest, SubResponse,
     },
     legacy_protocol::{
         SpanMode,
@@ -98,7 +98,7 @@ pub(crate) fn version_check(
     srv: &ProcMacroServerProcess,
     callback: SubCallback<'_>,
 ) -> Result<u32, ServerError> {
-    let request = BidirectionalMessage::Request(Request::ApiVersionCheck {});
+    let request = BidirectionalMessage::Request(Request::ApiVersionCheck(ApiVersionCheck {}));
 
     let response_payload = run_request(srv, request, callback)?;
 
@@ -135,9 +135,9 @@ pub(crate) fn find_proc_macros(
     dylib_path: &AbsPath,
     callback: SubCallback<'_>,
 ) -> Result<Result<Vec<(String, ProcMacroKind)>, String>, ServerError> {
-    let request = BidirectionalMessage::Request(Request::ListMacros {
+    let request = BidirectionalMessage::Request(Request::ListMacros(ListMacros {
         dylib_path: dylib_path.to_path_buf().into(),
-    });
+    }));
 
     let response_payload = run_request(srv, request, callback)?;
 
@@ -186,25 +186,12 @@ pub(crate) fn expand(
 
     match response_payload {
         BidirectionalMessage::Response(Response::ExpandMacro(it)) => Ok(it
-            .map(|tree| {
-                let mut expanded = FlatTree::to_subtree_resolved(tree, version, &span_data_table);
-                if proc_macro.needs_fixup_change() {
-                    proc_macro.change_fixup_to_match_old_server(&mut expanded);
-                }
-                expanded
-            })
-            .map_err(|msg| msg.0)),
-        BidirectionalMessage::Response(Response::ExpandMacroExtended(it)) => Ok(it
             .map(|resp| {
-                let mut expanded = FlatTree::to_subtree_resolved(
+                FlatTree::to_subtree_resolved(
                     resp.tree,
                     version,
                     &deserialize_span_data_index_map(&resp.span_data_table),
-                );
-                if proc_macro.needs_fixup_change() {
-                    proc_macro.change_fixup_to_match_old_server(&mut expanded);
-                }
-                expanded
+                )
             })
             .map_err(|msg| msg.0)),
         _ => Err(ServerError { message: "unexpected response".to_owned(), io: None }),

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

@@ -1,6 +1,8 @@
 //! Bidirectional protocol messages
+#![expect(clippy::disallowed_types)]
 
 use std::{
+    collections::{HashMap, HashSet},
     io::{self, BufRead, Write},
     ops::Range,
 };
@@ -16,13 +18,54 @@ use crate::{
 
 #[derive(Debug, Serialize, Deserialize)]
 pub enum SubRequest {
-    FilePath { file_id: u32 },
-    SourceText { file_id: u32, ast_id: u32, start: u32, end: u32 },
-    LocalFilePath { file_id: u32 },
-    LineColumn { file_id: u32, ast_id: u32, offset: u32 },
-    ByteRange { file_id: u32, ast_id: u32, start: u32, end: u32 },
-    SpanSource { file_id: u32, ast_id: u32, start: u32, end: u32, ctx: u32 },
-    SpanParent { file_id: u32, ast_id: u32, start: u32, end: u32, ctx: u32 },
+    FilePath {
+        file_id: u32,
+    },
+    SourceText {
+        file_id: u32,
+        ast_id: u32,
+        start: u32,
+        end: u32,
+    },
+    LocalFilePath {
+        file_id: u32,
+    },
+    LineColumn {
+        file_id: u32,
+        ast_id: u32,
+        offset: u32,
+    },
+    ByteRange {
+        file_id: u32,
+        ast_id: u32,
+        start: u32,
+        end: u32,
+    },
+    SpanSource {
+        file_id: u32,
+        ast_id: u32,
+        start: u32,
+        end: u32,
+        ctx: u32,
+    },
+    SpanParent {
+        file_id: u32,
+        ast_id: u32,
+        start: u32,
+        end: u32,
+        ctx: u32,
+    },
+    SpanJoin {
+        file_id: u32,
+        ast_id_first: u32,
+        start_first: u32,
+        end_first: u32,
+        ctx_first: u32,
+        ast_id_second: u32,
+        start_second: u32,
+        end_second: u32,
+        ctx_second: u32,
+    },
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -54,6 +97,9 @@ pub enum SubResponse {
     SpanParentResult {
         parent_span: Option<ParentSpan>,
     },
+    SpanJoinResult {
+        span: Option<SpanJoin>,
+    },
     Cancel {
         reason: String,
     },
@@ -68,6 +114,15 @@ pub struct ParentSpan {
     pub ctx: u32,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct SpanJoin {
+    pub ast_id: u32,
+    pub start: u32,
+    pub end: u32,
+    pub ctx: u32,
+}
+
+#[expect(clippy::large_enum_variant)]
 #[derive(Debug, Serialize, Deserialize)]
 pub enum BidirectionalMessage {
     Request(Request),
@@ -78,21 +133,29 @@ pub enum BidirectionalMessage {
 
 #[derive(Debug, Serialize, Deserialize)]
 pub enum Request {
-    ListMacros { dylib_path: Utf8PathBuf },
+    ListMacros(ListMacros),
     ExpandMacro(Box<ExpandMacro>),
-    ApiVersionCheck {},
+    ApiVersionCheck(ApiVersionCheck),
     SetConfig(ServerConfig),
 }
 
+#[expect(clippy::large_enum_variant)]
 #[derive(Debug, Serialize, Deserialize)]
 pub enum Response {
     ListMacros(Result<Vec<(String, ProcMacroKind)>, String>),
-    ExpandMacro(Result<FlatTree, PanicMessage>),
     ApiVersionCheck(u32),
     SetConfig(ServerConfig),
-    ExpandMacroExtended(Result<ExpandMacroExtended, PanicMessage>),
+    ExpandMacro(Result<ExpandMacroResponse, PanicMessage>),
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ListMacros {
+    pub dylib_path: Utf8PathBuf,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ApiVersionCheck {}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ExpandMacro {
     pub lib: Utf8PathBuf,
@@ -102,9 +165,11 @@ pub struct ExpandMacro {
 }
 
 #[derive(Debug, Serialize, Deserialize)]
-pub struct ExpandMacroExtended {
+pub struct ExpandMacroResponse {
     pub tree: FlatTree,
     pub span_data_table: Vec<u32>,
+    pub tracked_env_vars: HashMap<Box<str>, Option<Box<str>>>,
+    pub tracked_paths: HashSet<Box<str>>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]

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

@@ -240,8 +240,9 @@ impl ProcMacroServerProcess {
     /// Enable support for rust-analyzer span mode if the server supports it.
     pub(crate) fn rust_analyzer_spans(&self) -> bool {
         match self.protocol {
-            Protocol::LegacyJson { mode } => mode == SpanMode::RustAnalyzer,
-            Protocol::BidirectionalPostcardPrototype { mode } => mode == SpanMode::RustAnalyzer,
+            Protocol::LegacyJson { mode } | Protocol::BidirectionalPostcardPrototype { mode } => {
+                mode == SpanMode::RustAnalyzer
+            }
         }
     }
 

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

@@ -1,11 +1,11 @@
-//! Library interface for `proc-macro-srv-cli`.
-//!
-//! This module exposes the server main loop and protocol format for integration testing.
-
-#![cfg_attr(feature = "in-rust-tree", feature(rustc_private))]
-
-#[cfg(feature = "in-rust-tree")]
-extern crate rustc_driver as _;
-
-#[cfg(feature = "sysroot-abi")]
-pub mod main_loop;
+//! Library interface for `proc-macro-srv-cli`.
+//!
+//! This module exposes the server main loop and protocol format for integration testing.
+
+#![cfg_attr(feature = "in-rust-tree", feature(rustc_private))]
+
+#[cfg(feature = "in-rust-tree")]
+extern crate rustc_driver as _;
+
+#[cfg(feature = "sysroot-abi")]
+pub mod main_loop;

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

@@ -1,4 +1,5 @@
 //! The main loop of the proc-macro server.
+use proc_macro_api::bidirectional_protocol::msg::{ApiVersionCheck, ListMacros};
 use proc_macro_api::{
     ProtocolFormat, bidirectional_protocol::msg as bidirectional, legacy_protocol::msg as legacy,
     version::CURRENT_API_VERSION,
@@ -72,7 +73,7 @@ fn run_new(
 
         match req {
             bidirectional::BidirectionalMessage::Request(request) => match request {
-                bidirectional::Request::ListMacros { dylib_path } => {
+                bidirectional::Request::ListMacros(ListMacros { dylib_path }) => {
                     let res = srv.list_macros(&dylib_path).map(|macros| {
                         macros
                             .into_iter()
@@ -83,7 +84,7 @@ fn run_new(
                     send_response(stdout, bidirectional::Response::ListMacros(res))?;
                 }
 
-                bidirectional::Request::ApiVersionCheck {} => {
+                bidirectional::Request::ApiVersionCheck(ApiVersionCheck {}) => {
                     send_response(
                         stdout,
                         bidirectional::Response::ApiVersionCheck(CURRENT_API_VERSION),
@@ -142,6 +143,7 @@ fn handle_expand_id(
     let attributes = attributes
         .map(|it| it.to_tokenstream_unresolved::<SpanTrans>(CURRENT_API_VERSION, |_, b| b));
 
+    let mut tracked_env = Default::default();
     let res = srv
         .expand(
             lib,
@@ -153,11 +155,18 @@ fn handle_expand_id(
             def_site,
             call_site,
             mixed_site,
+            &mut tracked_env,
             None,
         )
         .map(|it| {
             legacy::FlatTree::from_tokenstream_raw::<SpanTrans>(it, call_site, CURRENT_API_VERSION)
         })
+        .map(|tree| bidirectional::ExpandMacroResponse {
+            tree,
+            span_data_table: vec![],
+            tracked_env_vars: tracked_env.env_vars,
+            tracked_paths: tracked_env.paths,
+        })
         .map_err(|e| legacy::PanicMessage(e.into_string().unwrap_or_default()));
 
     send_response(stdout, bidirectional::Response::ExpandMacro(res))
@@ -343,6 +352,46 @@ impl proc_macro_srv::ProcMacroClientInterface for ProcMacroClientHandle<'_> {
             other => handle_failure(other),
         }
     }
+
+    fn span_join(
+        &mut self,
+        first: proc_macro_srv::span::Span,
+        second: proc_macro_srv::span::Span,
+    ) -> Option<proc_macro_srv::span::Span> {
+        assert_eq!(first.anchor.file_id, second.anchor.file_id);
+        let response = self.roundtrip(bidirectional::SubRequest::SpanJoin {
+            file_id: first.anchor.file_id.as_u32(),
+            ast_id_first: first.anchor.ast_id.into_raw(),
+            start_first: first.range.start().into(),
+            end_first: first.range.end().into(),
+            ctx_first: first.ctx.into_u32(),
+            ast_id_second: second.anchor.ast_id.into_raw(),
+            start_second: second.range.start().into(),
+            end_second: second.range.end().into(),
+            ctx_second: second.ctx.into_u32(),
+        });
+
+        match response {
+            Ok(bidirectional::SubResponse::SpanJoinResult { span }) => {
+                span.map(|bidirectional::SpanJoin { ast_id, start, end, ctx }| {
+                    proc_macro_srv::span::Span {
+                        range: proc_macro_srv::span::TextRange::new(
+                            proc_macro_srv::span::TextSize::new(start),
+                            proc_macro_srv::span::TextSize::new(end),
+                        ),
+                        anchor: proc_macro_srv::span::SpanAnchor {
+                            file_id: first.anchor.file_id,
+                            ast_id: proc_macro_srv::span::ErasedFileAstId::from_raw(ast_id),
+                        },
+                        // SAFETY: spans originate from the server. If the protocol is violated,
+                        // undefined behavior is the caller’s responsibility.
+                        ctx: unsafe { proc_macro_srv::span::SyntaxContext::from_u32(ctx) },
+                    }
+                })
+            }
+            other => handle_failure(other),
+        }
+    }
 }
 
 fn handle_expand_ra(
@@ -383,6 +432,8 @@ fn handle_expand_ra(
         })
     });
 
+    let mut tracked_env = Default::default();
+
     let res = srv
         .expand(
             lib,
@@ -394,6 +445,7 @@ fn handle_expand_ra(
             def_site,
             call_site,
             mixed_site,
+            &mut tracked_env,
             Some(&mut ProcMacroClientHandle { stdin, stdout, buf }),
         )
         .map(|it| {
@@ -407,10 +459,15 @@ fn handle_expand_ra(
                 legacy::serialize_span_data_index_map(&span_data_table),
             )
         })
-        .map(|(tree, span_data_table)| bidirectional::ExpandMacroExtended { tree, span_data_table })
+        .map(|(tree, span_data_table)| bidirectional::ExpandMacroResponse {
+            tree,
+            span_data_table,
+            tracked_env_vars: tracked_env.env_vars,
+            tracked_paths: tracked_env.paths,
+        })
         .map_err(|e| legacy::PanicMessage(e.into_string().unwrap_or_default()));
 
-    send_response(stdout, bidirectional::Response::ExpandMacroExtended(res))
+    send_response(stdout, bidirectional::Response::ExpandMacro(res))
 }
 
 fn run_old(
@@ -480,6 +537,7 @@ fn run_old(
                             def_site,
                             call_site,
                             mixed_site,
+                            &mut Default::default(),
                             None,
                         )
                         .map(|it| {
@@ -522,6 +580,7 @@ fn run_old(
                             def_site,
                             call_site,
                             mixed_site,
+                            &mut Default::default(),
                             None,
                         )
                         .map(|it| {