add postcard abstraction inside Message trait

Author: Shourya742

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,23 +183,23 @@ fn send_request_postcard(
     req: Request,
     buf: &mut Vec<u8>,
 ) -> Result<Option<Response>, 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)),
     })?;
 
     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,
             })?;

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<R: BufRead>(
+        from_proto: ProtocolReadPostcard<R>,
+        inp: &mut R,
+        buf: &mut Vec<u8>,
+    ) -> io::Result<Option<Self>> {
+        Ok(match from_proto(inp, buf)? {
+            None => None,
+            Some(buf) => Some(decode_cobs(buf)?),
+        })
+    }
+
+    fn write_postcard<W: Write>(
+        self,
+        to_proto: ProtocolWritePostcard<W>,
+        out: &mut W,
+    ) -> io::Result<()> {
+        let buf = encode_cobs(&self)?;
+        to_proto(out, &buf)
+    }
 }
 
 impl Message for Request {}
@@ -182,6 +205,15 @@ type ProtocolRead<R: BufRead> =
 #[allow(type_alias_bounds)]
 type ProtocolWrite<W: Write> = 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<R: BufRead> =
+    for<'i, 'buf> fn(inp: &'i mut R, buf: &'buf mut Vec<u8>) -> io::Result<Option<&'buf mut [u8]>>;
+/// Type alias for a function that writes protocol postcard messages to an output stream.
+#[allow(type_alias_bounds)]
+type ProtocolWritePostcard<W: Write> =
+    for<'o, 'msg> fn(out: &'o mut W, msg: &'msg [u8]) -> io::Result<()>;
+
 #[cfg(test)]
 mod tests {
     use intern::{Symbol, sym};

…api/src/legacy_protocol/postcard_wire.rs …acro-api/src/legacy_protocol/postcard.rscrates/proc-macro-api/src/legacy_protocol/postcard_wire.rs renamed to crates/proc-macro-api/src/legacy_protocol/postcard.rs crates/proc-macro-api/src/legacy_protocol/postcard_wire.rs renamed to 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<T: serde::Serialize>(value: &T) -> Result<Vec<u8>, postcard::Error> {
-    postcard::to_allocvec_cobs(value)
+pub fn encode_cobs<T: serde::Serialize>(value: &T) -> io::Result<Vec<u8>> {
+    postcard::to_allocvec_cobs(value).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
 }
 
-pub fn decode_cobs<T: serde::de::DeserializeOwned>(bytes: &mut [u8]) -> Result<T, postcard::Error> {
-    postcard::from_bytes_cobs(bytes)
+pub fn decode_cobs<T: serde::de::DeserializeOwned>(bytes: &mut [u8]) -> io::Result<T> {
+    postcard::from_bytes_cobs(bytes).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
 }

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