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

Refactor: Bidirectional messages
This commit is contained in:
Lukas Wirth 2025-12-27 08:59:21 +00:00 committed by GitHub
commit f9e29cfb70
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 59 additions and 72 deletions

1
Cargo.lock generated
View File

@ -1877,7 +1877,6 @@ name = "proc-macro-srv-cli"
version = "0.0.0"
dependencies = [
"clap",
"crossbeam-channel",
"postcard",
"proc-macro-api",
"proc-macro-srv",

View File

@ -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]

View File

@ -1,6 +1,4 @@
//! The main loop of the proc-macro server.
use std::io;
use proc_macro_api::{
Codec,
bidirectional_protocol::msg as bidirectional,
@ -8,6 +6,7 @@ use proc_macro_api::{
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,

View File

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

View File

@ -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();

View File

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

View File

@ -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> {

View File

@ -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 {