add postcard related methods to proc-macro-apo

Cargo.lock

@@ -560,6 +560,18 @@ version = "1.15.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
 
+[[package]]
+name = "embedded-io"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ef1a6892d9eef45c8fa6b9e0086428a2cca8491aca8f787c534a3d6d0bcb3ced"
+
+[[package]]
+name = "embedded-io"
+version = "0.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "edd0f118536f44f5ccd48bcb8b111bdc3de888b58c74639dfb034a357d0f206d"
+
 [[package]]
 name = "ena"
 version = "0.14.3"
@@ -1785,6 +1797,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6764c3b5dd454e283a30e6dfe78e9b31096d9e32036b5d1eaac7a6119ccb9a24"
 dependencies = [
  "cobs",
+ "embedded-io 0.4.0",
+ "embedded-io 0.6.1",
  "heapless",
  "serde",
 ]
@@ -1820,6 +1834,7 @@ dependencies = [
  "indexmap",
  "intern",
  "paths",
+ "postcard",
  "proc-macro-srv",
  "rustc-hash 2.1.1",
  "serde",

Cargo.toml

@@ -127,6 +127,7 @@ object = { version = "0.36.7", default-features = false, features = [
     "macho",
     "pe",
 ] }
+postcard = {version = "1.1.3", features = ["alloc"]}
 process-wrap = { version = "8.2.1", features = ["std"] }
 pulldown-cmark-to-cmark = "10.0.4"
 pulldown-cmark = { version = "0.9.6", default-features = false }

crates/proc-macro-api/Cargo.toml

@@ -29,6 +29,7 @@ proc-macro-srv = {workspace = true, optional = true}
 span = { path = "../span", version = "0.0.0", default-features = false}
 
 intern.workspace = true
+postcard.workspace = true
 
 [features]
 sysroot-abi = ["proc-macro-srv", "proc-macro-srv/sysroot-abi"]

crates/proc-macro-api/src/legacy_protocol.rs

@@ -2,6 +2,7 @@
 
 pub mod json;
 pub mod msg;
+pub mod postcard_wire;
 
 use std::{
     io::{BufRead, Write},
@@ -151,7 +152,11 @@ fn send_task(srv: &ProcMacroServerProcess, req: Request) -> Result<Response, Ser
         return Err(server_error.clone());
     }
 
-    srv.send_task(send_request, req)
+    if srv.use_postcard() {
+        srv.send_task_bin(send_request_postcard, req)
+    } else {
+        srv.send_task(send_request, req)
+    }
 }
 
 /// Sends a request to the server and reads the response.
@@ -171,3 +176,34 @@ fn send_request(
     })?;
     Ok(res)
 }
+
+fn send_request_postcard(
+    mut writer: &mut dyn Write,
+    mut reader: &mut dyn BufRead,
+    req: Request,
+    buf: &mut Vec<u8>,
+) -> Result<Option<Response>, ServerError> {
+    let bytes = postcard_wire::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 {
+        message: "failed to write request".into(),
+        io: Some(Arc::new(err)),
+    })?;
+
+    let frame = postcard_wire::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 {
+                message: format!("failed to decode message: {e}"),
+                io: None,
+            })?;
+            Ok(Some(resp))
+        }
+    }
+}

crates/proc-macro-api/src/legacy_protocol/postcard_wire.rs

@@ -0,0 +1,27 @@
+//! Postcard encode and decode implementations.
+
+use std::io::{self, BufRead, Write};
+
+pub fn read_postcard<'a>(
+    input: &mut impl BufRead,
+    buf: &'a mut Vec<u8>,
+) -> io::Result<Option<&'a mut [u8]>> {
+    buf.clear();
+    let n = input.read_until(0, buf)?;
+    if n == 0 {
+        return Ok(None);
+    }
+    Ok(Some(&mut buf[..]))
+}
+pub fn write_postcard(out: &mut impl Write, msg: &[u8]) -> io::Result<()> {
+    out.write_all(msg)?;
+    out.flush()
+}
+
+pub fn encode_cobs<T: serde::Serialize>(value: &T) -> Result<Vec<u8>, postcard::Error> {
+    postcard::to_allocvec_cobs(value)
+}
+
+pub fn decode_cobs<T: serde::de::DeserializeOwned>(bytes: &mut [u8]) -> Result<T, postcard::Error> {
+    postcard::from_bytes_cobs(bytes)
+}

crates/proc-macro-api/src/lib.rs

@@ -31,9 +31,10 @@ pub mod version {
     /// Whether literals encode their kind as an additional u32 field and idents their rawness as a u32 field.
     pub const EXTENDED_LEAF_DATA: u32 = 5;
     pub const HASHED_AST_ID: u32 = 6;
+    pub const POSTCARD_WIRE: u32 = 7;
 
     /// Current API version of the proc-macro protocol.
-    pub const CURRENT_API_VERSION: u32 = HASHED_AST_ID;
+    pub const CURRENT_API_VERSION: u32 = POSTCARD_WIRE;
 }
 
 /// Represents different kinds of procedural macros that can be expanded by the external server.

crates/proc-macro-api/src/process.rs

@@ -31,6 +31,7 @@ pub(crate) struct ProcMacroServerProcess {
 #[derive(Debug)]
 enum Protocol {
     LegacyJson { mode: SpanMode },
+    Postcard { mode: SpanMode },
 }
 
 /// Maintains the state of the proc-macro server process.
@@ -82,7 +83,11 @@ impl ProcMacroServerProcess {
                 if srv.version >= version::RUST_ANALYZER_SPAN_SUPPORT
                     && let Ok(mode) = srv.enable_rust_analyzer_spans()
                 {
-                    srv.protocol = Protocol::LegacyJson { mode };
+                    if srv.version >= version::POSTCARD_WIRE {
+                        srv.protocol = Protocol::Postcard { mode };
+                    } else {
+                        srv.protocol = Protocol::LegacyJson { mode };
+                    }
                 }
                 tracing::info!("Proc-macro server protocol: {:?}", srv.protocol);
                 Ok(srv)
@@ -99,6 +104,10 @@ impl ProcMacroServerProcess {
         self.exited.get().map(|it| &it.0)
     }
 
+    pub(crate) fn use_postcard(&self) -> bool {
+        matches!(self.protocol, Protocol::Postcard { .. })
+    }
+
     /// Retrieves the API version of the proc-macro server.
     pub(crate) fn version(&self) -> u32 {
         self.version
@@ -108,6 +117,7 @@ impl ProcMacroServerProcess {
     pub(crate) fn rust_analyzer_spans(&self) -> bool {
         match self.protocol {
             Protocol::LegacyJson { mode } => mode == SpanMode::RustAnalyzer,
+            Protocol::Postcard { mode } => mode == SpanMode::RustAnalyzer,
         }
     }
 
@@ -115,6 +125,7 @@ impl ProcMacroServerProcess {
     fn version_check(&self) -> Result<u32, ServerError> {
         match self.protocol {
             Protocol::LegacyJson { .. } => legacy_protocol::version_check(self),
+            Protocol::Postcard { .. } => legacy_protocol::version_check(self),
         }
     }
 
@@ -122,6 +133,7 @@ impl ProcMacroServerProcess {
     fn enable_rust_analyzer_spans(&self) -> Result<SpanMode, ServerError> {
         match self.protocol {
             Protocol::LegacyJson { .. } => legacy_protocol::enable_rust_analyzer_spans(self),
+            Protocol::Postcard { .. } => legacy_protocol::enable_rust_analyzer_spans(self),
         }
     }
 
@@ -132,6 +144,7 @@ impl ProcMacroServerProcess {
     ) -> Result<Result<Vec<(String, ProcMacroKind)>, String>, ServerError> {
         match self.protocol {
             Protocol::LegacyJson { .. } => legacy_protocol::find_proc_macros(self, dylib_path),
+            Protocol::Postcard { .. } => legacy_protocol::find_proc_macros(self, dylib_path),
         }
     }
 
@@ -188,6 +201,55 @@ impl ProcMacroServerProcess {
                 }
             })
     }
+
+    pub(crate) fn send_task_bin<Request, Response>(
+        &self,
+        serialize_req: impl FnOnce(
+            &mut dyn Write,
+            &mut dyn BufRead,
+            Request,
+            &mut Vec<u8>,
+        ) -> Result<Option<Response>, ServerError>,
+        req: Request,
+    ) -> Result<Response, ServerError> {
+        let state = &mut *self.state.lock().unwrap();
+        let mut buf = Vec::<u8>::new();
+        serialize_req(&mut state.stdin, &mut state.stdout, req, &mut buf)
+            .and_then(|res| {
+                res.ok_or_else(|| ServerError {
+                    message: "proc-macro server did not respond with data".to_owned(),
+                    io: Some(Arc::new(io::Error::new(
+                        io::ErrorKind::BrokenPipe,
+                        "proc-macro server did not respond with data",
+                    ))),
+                })
+            })
+            .map_err(|e| {
+                if e.io.as_ref().map(|it| it.kind()) == Some(io::ErrorKind::BrokenPipe) {
+                    match state.process.child.try_wait() {
+                        Ok(None) | Err(_) => e,
+                        Ok(Some(status)) => {
+                            let mut msg = String::new();
+                            if !status.success()
+                                && let Some(stderr) = state.process.child.stderr.as_mut()
+                            {
+                                _ = stderr.read_to_string(&mut msg);
+                            }
+                            let server_error = ServerError {
+                                message: format!(
+                                    "proc-macro server exited with {status}{}{msg}",
+                                    if msg.is_empty() { "" } else { ": " }
+                                ),
+                                io: None,
+                            };
+                            self.exited.get_or_init(|| AssertUnwindSafe(server_error)).0.clone()
+                        }
+                    }
+                } else {
+                    e
+                }
+            })
+    }
 }
 
 /// Manages the execution of the proc-macro server process.

crates/proc-macro-srv-cli/src/main_loop.rs

@@ -200,8 +200,6 @@ fn run_postcard() -> io::Result<()> {
             thread::sleep(std::time::Duration::from_secs(1));
             continue;
         };
-        dbg!(&req);
-
         let res = match req {
             msg::Request::ListMacros { dylib_path } => {
                 msg::Response::ListMacros(srv.list_macros(&dylib_path).map(|macros| {
@@ -304,7 +302,6 @@ fn run_postcard() -> io::Result<()> {
             }
         };
 
-        dbg!(&res);
         let res = postcard_wire::encode_cobs(&res).unwrap();
         postcard_wire::write_postcard(&mut writer, &res)?;
     }