add bidirectional handle in proc-macro-srv-cli to interact with client and srv

This commit is contained in:
bit-aloo
2025-12-27 09:07:08 +05:30
parent 8aa859476c
commit 231690657e
@@ -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 @@
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 BidirectionalProxy {
subreq_tx: mpsc::Sender<bidirectional::SubRequest>,
subresp_rx: mpsc::Receiver<bidirectional::SubResponse>,
}
impl proc_macro_srv::BidirectionalHandler for BidirectionalProxy {
fn source_text(&mut 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 = BidirectionalProxy { 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 macro_kind_to_api(kind: proc_macro_srv::ProcMacroKind) -> proc_macro_api::Pro
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,