Merge pull request #21340 from Shourya742/2025-12-26-better-bm-ergonomics

Veykril · web-flow · commit f9e29cfb70ec · 2025-12-27T08:59:21.000Z
Refactor: Bidirectional messages
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/proc-macro-srv-cli/Cargo.toml b/crates/proc-macro-srv-cli/Cargo.toml
@@ -14,7 +14,6 @@ publish = false
 proc-macro-srv.workspace = true
 proc-macro-api.workspace = true
 postcard.workspace = true
-crossbeam-channel.workspace = true
 clap = {version = "4.5.42", default-features = false, features = ["std"]}
 
 [features]
diff --git a/crates/proc-macro-srv-cli/src/main_loop.rs b/crates/proc-macro-srv-cli/src/main_loop.rs
@@ -1,13 +1,12 @@
 //! The main loop of the proc-macro server.
-use std::io;
-
 use proc_macro_api::{
     Codec,
     bidirectional_protocol::msg as bidirectional,
     legacy_protocol::msg as legacy,
     transport::codec::{json::JsonProtocol, postcard::PostcardProtocol},
     version::CURRENT_API_VERSION,
 };
+use std::{io, sync::mpsc};
 
 use legacy::Message;
 
@@ -170,6 +169,21 @@ fn handle_expand_id<W: std::io::Write, C: Codec>(
     send_response::<_, C>(stdout, bidirectional::Response::ExpandMacro(res))
 }
 
+struct ProcMacroClientHandle {
+    subreq_tx: mpsc::Sender<bidirectional::SubRequest>,
+    subresp_rx: mpsc::Receiver<bidirectional::SubResponse>,
+}
+
+impl proc_macro_srv::ProcMacroClientInterface for ProcMacroClientHandle {
+    fn source_text(&self, file_id: u32, start: u32, end: u32) -> Option<String> {
+        self.subreq_tx.send(bidirectional::SubRequest::SourceText { file_id, start, end }).ok()?;
+
+        match self.subresp_rx.recv().ok()? {
+            bidirectional::SubResponse::SourceTextResult { text } => text,
+        }
+    }
+}
+
 fn handle_expand_ra<W: io::Write, R: io::BufRead, C: Codec>(
     srv: &proc_macro_srv::ProcMacroSrv<'_>,
     stdin: &mut R,
@@ -201,22 +215,20 @@ fn handle_expand_ra<W: io::Write, R: io::BufRead, C: Codec>(
         macro_body.to_tokenstream_resolved(CURRENT_API_VERSION, &span_data_table, |a, b| {
             srv.join_spans(a, b).unwrap_or(b)
         });
+
     let attributes = attributes.map(|it| {
         it.to_tokenstream_resolved(CURRENT_API_VERSION, &span_data_table, |a, b| {
             srv.join_spans(a, b).unwrap_or(b)
         })
     });
 
-    let (subreq_tx, subreq_rx) = crossbeam_channel::unbounded();
-    let (subresp_tx, subresp_rx) = crossbeam_channel::unbounded();
-    let (result_tx, result_rx) = crossbeam_channel::bounded(1);
+    let (subreq_tx, subreq_rx) = mpsc::channel();
+    let (subresp_tx, subresp_rx) = mpsc::channel();
+    let (result_tx, result_rx) = mpsc::channel();
 
-    std::thread::scope(|scope| {
-        scope.spawn(|| {
-            let callback = Box::new(move |req: proc_macro_srv::SubRequest| {
-                subreq_tx.send(req).unwrap();
-                subresp_rx.recv().unwrap()
-            });
+    std::thread::scope(|s| {
+        s.spawn(|| {
+            let callback = ProcMacroClientHandle { subreq_tx, subresp_rx };
 
             let res = srv
                 .expand(
@@ -229,7 +241,7 @@ fn handle_expand_ra<W: io::Write, R: io::BufRead, C: Codec>(
                     def_site,
                     call_site,
                     mixed_site,
-                    Some(callback),
+                    Some(Box::new(callback)),
                 )
                 .map(|it| {
                     (
@@ -253,27 +265,31 @@ fn handle_expand_ra<W: io::Write, R: io::BufRead, C: Codec>(
 
         loop {
             if let Ok(res) = result_rx.try_recv() {
-                send_response::<_, C>(stdout, bidirectional::Response::ExpandMacroExtended(res))
-                    .unwrap();
+                let _ = send_response::<_, C>(
+                    stdout,
+                    bidirectional::Response::ExpandMacroExtended(res),
+                );
                 break;
             }
 
-            let subreq = match subreq_rx.recv() {
+            let sub_req = match subreq_rx.recv() {
                 Ok(r) => r,
                 Err(_) => break,
             };
 
-            let api_req = from_srv_req(subreq);
-            bidirectional::BidirectionalMessage::SubRequest(api_req).write::<_, C>(stdout).unwrap();
-
-            let resp = bidirectional::BidirectionalMessage::read::<_, C>(stdin, buf)
-                .unwrap()
-                .expect("client closed connection");
-
+            if bidirectional::BidirectionalMessage::SubRequest(sub_req)
+                .write::<_, C>(stdout)
+                .is_err()
+            {
+                break;
+            }
+            let resp = match bidirectional::BidirectionalMessage::read::<_, C>(stdin, buf) {
+                Ok(Some(r)) => r,
+                _ => break,
+            };
             match resp {
-                bidirectional::BidirectionalMessage::SubResponse(api_resp) => {
-                    let srv_resp = from_client_res(api_resp);
-                    subresp_tx.send(srv_resp).unwrap();
+                bidirectional::BidirectionalMessage::SubResponse(resp) => {
+                    let _ = subresp_tx.send(resp);
                 }
                 other => panic!("expected SubResponse, got {other:?}"),
             }
@@ -425,22 +441,6 @@ fn run_<C: Codec>() -> io::Result<()> {
     Ok(())
 }
 
-fn from_srv_req(value: proc_macro_srv::SubRequest) -> bidirectional::SubRequest {
-    match value {
-        proc_macro_srv::SubRequest::SourceText { file_id, start, end } => {
-            bidirectional::SubRequest::SourceText { file_id: file_id.file_id().index(), start, end }
-        }
-    }
-}
-
-fn from_client_res(value: bidirectional::SubResponse) -> proc_macro_srv::SubResponse {
-    match value {
-        bidirectional::SubResponse::SourceTextResult { text } => {
-            proc_macro_srv::SubResponse::SourceTextResult { text }
-        }
-    }
-}
-
 fn send_response<W: std::io::Write, C: Codec>(
     stdout: &mut W,
     resp: bidirectional::Response,
diff --git a/crates/proc-macro-srv/src/dylib.rs b/crates/proc-macro-srv/src/dylib.rs
@@ -12,8 +12,8 @@ use object::Object;
 use paths::{Utf8Path, Utf8PathBuf};
 
 use crate::{
-    PanicMessage, ProcMacroKind, ProcMacroSrvSpan, SubCallback, dylib::proc_macros::ProcMacros,
-    token_stream::TokenStream,
+    PanicMessage, ProcMacroClientHandle, ProcMacroKind, ProcMacroSrvSpan,
+    dylib::proc_macros::ProcMacros, token_stream::TokenStream,
 };
 
 pub(crate) struct Expander {
@@ -45,7 +45,7 @@ impl Expander {
         def_site: S,
         call_site: S,
         mixed_site: S,
-        callback: Option<SubCallback>,
+        callback: Option<ProcMacroClientHandle>,
     ) -> Result<TokenStream<S>, PanicMessage>
     where
         <S::Server as bridge::server::Types>::TokenStream: Default,
diff --git a/crates/proc-macro-srv/src/dylib/proc_macros.rs b/crates/proc-macro-srv/src/dylib/proc_macros.rs
@@ -1,5 +1,5 @@
 //! Proc macro ABI
-use crate::{ProcMacroKind, ProcMacroSrvSpan, SubCallback, token_stream::TokenStream};
+use crate::{ProcMacroClientHandle, ProcMacroKind, ProcMacroSrvSpan, token_stream::TokenStream};
 use proc_macro::bridge;
 
 #[repr(transparent)]
@@ -20,7 +20,7 @@ impl ProcMacros {
         def_site: S,
         call_site: S,
         mixed_site: S,
-        callback: Option<SubCallback>,
+        callback: Option<ProcMacroClientHandle>,
     ) -> Result<TokenStream<S>, crate::PanicMessage> {
         let parsed_attributes = attribute.unwrap_or_default();
 
diff --git a/crates/proc-macro-srv/src/lib.rs b/crates/proc-macro-srv/src/lib.rs
@@ -47,7 +47,7 @@ use std::{
 };
 
 use paths::{Utf8Path, Utf8PathBuf};
-use span::{EditionedFileId, Span};
+use span::Span;
 use temp_dir::TempDir;
 
 pub use crate::server_impl::token_id::SpanId;
@@ -91,14 +91,10 @@ impl<'env> ProcMacroSrv<'env> {
     }
 }
 
-pub type SubCallback = Box<dyn Fn(SubRequest) -> SubResponse + Send + Sync + 'static>;
+pub type ProcMacroClientHandle = Box<dyn ProcMacroClientInterface + Send>;
 
-pub enum SubRequest {
-    SourceText { file_id: EditionedFileId, start: u32, end: u32 },
-}
-
-pub enum SubResponse {
-    SourceTextResult { text: Option<String> },
+pub trait ProcMacroClientInterface {
+    fn source_text(&self, file_id: u32, start: u32, end: u32) -> Option<String>;
 }
 
 const EXPANDER_STACK_SIZE: usize = 8 * 1024 * 1024;
@@ -115,7 +111,7 @@ impl ProcMacroSrv<'_> {
         def_site: S,
         call_site: S,
         mixed_site: S,
-        callback: Option<SubCallback>,
+        callback: Option<ProcMacroClientHandle>,
     ) -> Result<token_stream::TokenStream<S>, PanicMessage> {
         let snapped_env = self.env;
         let expander = self.expander(lib.as_ref()).map_err(|err| PanicMessage {
@@ -187,7 +183,7 @@ pub trait ProcMacroSrvSpan: Copy + Send + Sync {
         call_site: Self,
         def_site: Self,
         mixed_site: Self,
-        callback: Option<SubCallback>,
+        callback: Option<ProcMacroClientHandle>,
     ) -> Self::Server;
 }
 
@@ -198,7 +194,7 @@ impl ProcMacroSrvSpan for SpanId {
         call_site: Self,
         def_site: Self,
         mixed_site: Self,
-        callback: Option<SubCallback>,
+        callback: Option<ProcMacroClientHandle>,
     ) -> Self::Server {
         Self::Server {
             call_site,
@@ -217,7 +213,7 @@ impl ProcMacroSrvSpan for Span {
         call_site: Self,
         def_site: Self,
         mixed_site: Self,
-        callback: Option<SubCallback>,
+        callback: Option<ProcMacroClientHandle>,
     ) -> Self::Server {
         Self::Server {
             call_site,
diff --git a/crates/proc-macro-srv/src/server_impl/rust_analyzer_span.rs b/crates/proc-macro-srv/src/server_impl/rust_analyzer_span.rs
@@ -14,7 +14,7 @@ use proc_macro::bridge::server;
 use span::{FIXUP_ERASED_FILE_AST_ID_MARKER, Span, TextRange, TextSize};
 
 use crate::{
-    SubCallback, SubRequest, SubResponse,
+    ProcMacroClientHandle,
     bridge::{Diagnostic, ExpnGlobals, Literal, TokenTree},
     server_impl::literal_from_str,
 };
@@ -29,7 +29,7 @@ pub struct RaSpanServer {
     pub call_site: Span,
     pub def_site: Span,
     pub mixed_site: Span,
-    pub callback: Option<SubCallback>,
+    pub callback: Option<ProcMacroClientHandle>,
 }
 
 impl server::Types for RaSpanServer {
@@ -156,14 +156,7 @@ impl server::Span for RaSpanServer {
         let start: u32 = span.range.start().into();
         let end: u32 = span.range.end().into();
 
-        let req = SubRequest::SourceText { file_id, start, end };
-
-        let cb = self.callback.as_mut()?;
-        let response = cb(req);
-
-        match response {
-            SubResponse::SourceTextResult { text } => text,
-        }
+        self.callback.as_ref()?.source_text(file_id.file_id().index(), start, end)
     }
 
     fn parent(&mut self, _span: Self::Span) -> Option<Self::Span> {
diff --git a/crates/proc-macro-srv/src/server_impl/token_id.rs b/crates/proc-macro-srv/src/server_impl/token_id.rs
@@ -9,7 +9,7 @@ use intern::Symbol;
 use proc_macro::bridge::server;
 
 use crate::{
-    SubCallback,
+    ProcMacroClientHandle,
     bridge::{Diagnostic, ExpnGlobals, Literal, TokenTree},
     server_impl::literal_from_str,
 };
@@ -35,7 +35,7 @@ pub struct SpanIdServer {
     pub call_site: Span,
     pub def_site: Span,
     pub mixed_site: Span,
-    pub callback: Option<SubCallback>,
+    pub callback: Option<ProcMacroClientHandle>,
 }
 
 impl server::Types for SpanIdServer {
