From e344bfa7eae548b3d8a0ea88ee2f228071cf1285 Mon Sep 17 00:00:00 2001 From: bit-aloo Date: Sat, 22 Nov 2025 21:06:51 +0530 Subject: [PATCH] add postcard abstraction inside Message trait --- crates/proc-macro-api/src/legacy_protocol.rs | 10 +++--- .../proc-macro-api/src/legacy_protocol/msg.rs | 34 ++++++++++++++++++- .../{postcard_wire.rs => postcard.rs} | 8 ++--- crates/proc-macro-srv-cli/src/main_loop.rs | 25 +++++--------- 4 files changed, 51 insertions(+), 26 deletions(-) rename crates/proc-macro-api/src/legacy_protocol/{postcard_wire.rs => postcard.rs} (64%) diff --git a/crates/proc-macro-api/src/legacy_protocol.rs b/crates/proc-macro-api/src/legacy_protocol.rs index 5d590520c8..cb869fddbc 100644 --- a/crates/proc-macro-api/src/legacy_protocol.rs +++ b/crates/proc-macro-api/src/legacy_protocol.rs @@ -2,7 +2,7 @@ pub mod json; pub mod msg; -pub mod postcard_wire; +pub mod postcard; use std::{ io::{BufRead, Write}, @@ -183,15 +183,15 @@ fn send_request_postcard( req: Request, buf: &mut Vec, ) -> Result, ServerError> { - let bytes = postcard_wire::encode_cobs(&req) + let bytes = postcard::encode_cobs(&req) .map_err(|_| ServerError { message: "failed to write request".into(), io: None })?; - postcard_wire::write_postcard(&mut writer, &bytes).map_err(|err| ServerError { + postcard::write_postcard(&mut writer, &bytes).map_err(|err| ServerError { message: "failed to write request".into(), io: Some(Arc::new(err)), })?; - let frame = postcard_wire::read_postcard(&mut reader, buf).map_err(|err| ServerError { + let frame = postcard::read_postcard(&mut reader, buf).map_err(|err| ServerError { message: "failed to read response".into(), io: Some(Arc::new(err)), })?; @@ -199,7 +199,7 @@ fn send_request_postcard( match frame { None => Ok(None), Some(bytes) => { - let resp: Response = postcard_wire::decode_cobs(bytes).map_err(|e| ServerError { + let resp: Response = postcard::decode_cobs(bytes).map_err(|e| ServerError { message: format!("failed to decode message: {e}"), io: None, })?; diff --git a/crates/proc-macro-api/src/legacy_protocol/msg.rs b/crates/proc-macro-api/src/legacy_protocol/msg.rs index 487f50b145..d3744dee0e 100644 --- a/crates/proc-macro-api/src/legacy_protocol/msg.rs +++ b/crates/proc-macro-api/src/legacy_protocol/msg.rs @@ -8,7 +8,10 @@ use paths::Utf8PathBuf; use serde::de::DeserializeOwned; use serde_derive::{Deserialize, Serialize}; -use crate::ProcMacroKind; +use crate::{ + ProcMacroKind, + legacy_protocol::postcard::{decode_cobs, encode_cobs}, +}; /// Represents requests sent from the client to the proc-macro-srv. #[derive(Debug, Serialize, Deserialize)] @@ -169,6 +172,26 @@ pub trait Message: serde::Serialize + DeserializeOwned { let text = serde_json::to_string(&self)?; to_proto(out, &text) } + + fn read_postcard( + from_proto: ProtocolReadPostcard, + inp: &mut R, + buf: &mut Vec, + ) -> io::Result> { + Ok(match from_proto(inp, buf)? { + None => None, + Some(buf) => Some(decode_cobs(buf)?), + }) + } + + fn write_postcard( + self, + to_proto: ProtocolWritePostcard, + out: &mut W, + ) -> io::Result<()> { + let buf = encode_cobs(&self)?; + to_proto(out, &buf) + } } impl Message for Request {} @@ -182,6 +205,15 @@ type ProtocolRead = #[allow(type_alias_bounds)] type ProtocolWrite = for<'o, 'msg> fn(out: &'o mut W, msg: &'msg str) -> io::Result<()>; +/// Type alias for a function that reads protocol postcard messages from a buffered input stream. +#[allow(type_alias_bounds)] +type ProtocolReadPostcard = + for<'i, 'buf> fn(inp: &'i mut R, buf: &'buf mut Vec) -> io::Result>; +/// Type alias for a function that writes protocol postcard messages to an output stream. +#[allow(type_alias_bounds)] +type ProtocolWritePostcard = + for<'o, 'msg> fn(out: &'o mut W, msg: &'msg [u8]) -> io::Result<()>; + #[cfg(test)] mod tests { use intern::{Symbol, sym}; diff --git a/crates/proc-macro-api/src/legacy_protocol/postcard_wire.rs b/crates/proc-macro-api/src/legacy_protocol/postcard.rs similarity index 64% rename from crates/proc-macro-api/src/legacy_protocol/postcard_wire.rs rename to crates/proc-macro-api/src/legacy_protocol/postcard.rs index 2c8bb5a0c9..eab26439db 100644 --- a/crates/proc-macro-api/src/legacy_protocol/postcard_wire.rs +++ b/crates/proc-macro-api/src/legacy_protocol/postcard.rs @@ -18,10 +18,10 @@ pub fn write_postcard(out: &mut impl Write, msg: &[u8]) -> io::Result<()> { out.flush() } -pub fn encode_cobs(value: &T) -> Result, postcard::Error> { - postcard::to_allocvec_cobs(value) +pub fn encode_cobs(value: &T) -> io::Result> { + postcard::to_allocvec_cobs(value).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e)) } -pub fn decode_cobs(bytes: &mut [u8]) -> Result { - postcard::from_bytes_cobs(bytes) +pub fn decode_cobs(bytes: &mut [u8]) -> io::Result { + postcard::from_bytes_cobs(bytes).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e)) } diff --git a/crates/proc-macro-srv-cli/src/main_loop.rs b/crates/proc-macro-srv-cli/src/main_loop.rs index 017efb7222..b0e7108d20 100644 --- a/crates/proc-macro-srv-cli/src/main_loop.rs +++ b/crates/proc-macro-srv-cli/src/main_loop.rs @@ -1,5 +1,5 @@ //! The main loop of the proc-macro server. -use std::{io, thread}; +use std::io; use proc_macro_api::{ legacy_protocol::{ @@ -8,14 +8,13 @@ use proc_macro_api::{ self, ExpandMacroData, ExpnGlobals, Message, SpanMode, SpanTransformer, deserialize_span_data_index_map, serialize_span_data_index_map, }, + postcard::{read_postcard, write_postcard}, }, version::CURRENT_API_VERSION, }; use proc_macro_srv::{EnvSnapshot, SpanId}; use crate::ProtocolFormat; -use std::io::BufReader; - struct SpanTrans; impl SpanTransformer for SpanTrans { @@ -183,23 +182,18 @@ fn run_postcard() -> io::Result<()> { } } - let stdin = io::stdin(); - let stdout = io::stdout(); - let mut reader = BufReader::new(stdin.lock()); - let mut writer = stdout.lock(); - let mut buf = vec![0; 1024]; + let mut buf = Vec::new(); + let mut read_request = + || msg::Request::read_postcard(read_postcard, &mut io::stdin().lock(), &mut buf); + let write_response = + |msg: msg::Response| msg.write_postcard(write_postcard, &mut io::stdout().lock()); let env = proc_macro_srv::EnvSnapshot::default(); let srv = proc_macro_srv::ProcMacroSrv::new(&env); let mut span_mode = msg::SpanMode::Id; - use proc_macro_api::legacy_protocol::postcard_wire; - while let Some(req) = postcard_wire::read_postcard(&mut reader, &mut buf)? { - let Ok(req) = postcard_wire::decode_cobs(req) else { - thread::sleep(std::time::Duration::from_secs(1)); - continue; - }; + while let Some(req) = read_request()? { let res = match req { msg::Request::ListMacros { dylib_path } => { msg::Response::ListMacros(srv.list_macros(&dylib_path).map(|macros| { @@ -302,8 +296,7 @@ fn run_postcard() -> io::Result<()> { } }; - let res = postcard_wire::encode_cobs(&res).unwrap(); - postcard_wire::write_postcard(&mut writer, &res)?; + write_response(res)?; } Ok(())