diff --git a/Cargo.lock b/Cargo.lock index 01de430925dc..98dba58ccaf4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -44,6 +44,15 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" +[[package]] +name = "atomic-polyfill" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8cf2bce30dfe09ef0bfaef228b9d414faaf7e563035494d7fe092dba54b300f4" +dependencies = [ + "critical-section", +] + [[package]] name = "autocfg" version = "1.4.0" @@ -119,6 +128,12 @@ version = "0.2.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "66bb12751a83493ef4b8da1120451a262554e216a247f14b48cb5e8fe7ed8bdf" +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "camino" version = "1.1.9" @@ -239,6 +254,12 @@ dependencies = [ "tracing", ] +[[package]] +name = "cobs" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67ba02a97a2bd10f4b59b25c7973101c79642302776489e030cd13cdab09ed15" + [[package]] name = "countme" version = "3.0.1" @@ -260,6 +281,12 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "critical-section" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "790eea4361631c5e7d22598ecd5723ff611904e3344ce8720784c93e3d83d40b" + [[package]] name = "crossbeam-channel" version = "0.5.15" @@ -416,6 +443,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" @@ -516,6 +555,15 @@ version = "0.31.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" +[[package]] +name = "hash32" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0c35f58762feb77d74ebe43bdbc3210f09be9fe6742234d573bacc26ed92b67" +dependencies = [ + "byteorder", +] + [[package]] name = "hashbrown" version = "0.14.5" @@ -542,6 +590,20 @@ dependencies = [ "hashbrown 0.15.2", ] +[[package]] +name = "heapless" +version = "0.7.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdc6457c0eb62c71aac4bc17216026d8410337c4126773b9c5daba343f17964f" +dependencies = [ + "atomic-polyfill", + "hash32", + "rustc_version", + "serde", + "spin", + "stable_deref_trait", +] + [[package]] name = "heck" version = "0.5.0" @@ -1510,6 +1572,19 @@ version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "350e9b48cbc6b0e028b0473b114454c6316e57336ee184ceab6e53f72c178b3e" +[[package]] +name = "postcard" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "170a2601f67cc9dba8edd8c4870b15f71a6a2dc196daec8c83f72b59dff628a8" +dependencies = [ + "cobs", + "embedded-io 0.4.0", + "embedded-io 0.6.1", + "heapless", + "serde", +] + [[package]] name = "powerfmt" version = "0.2.0" @@ -1523,6 +1598,7 @@ dependencies = [ "indexmap", "intern", "paths", + "postcard", "rustc-hash 2.1.1", "serde", "serde_derive", @@ -1919,6 +1995,15 @@ dependencies = [ "smallvec", ] +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + [[package]] name = "ryu" version = "1.0.20" @@ -2104,6 +2189,15 @@ dependencies = [ "vfs", ] +[[package]] +name = "spin" +version = "0.9.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" +dependencies = [ + "lock_api", +] + [[package]] name = "stable_deref_trait" version = "1.2.0" diff --git a/crates/proc-macro-api/Cargo.toml b/crates/proc-macro-api/Cargo.toml index f5ba40a994b1..52af72363c18 100644 --- a/crates/proc-macro-api/Cargo.toml +++ b/crates/proc-macro-api/Cargo.toml @@ -14,6 +14,7 @@ rust-version.workspace = true [dependencies] serde.workspace = true serde_derive.workspace = true +postcard = { version = "1.1.1", features = ["alloc"] } serde_json = { workspace = true, features = ["unbounded_depth"] } tracing.workspace = true rustc-hash.workspace = true @@ -24,7 +25,7 @@ paths = { workspace = true, features = ["serde1"] } tt.workspace = true stdx.workspace = true # span = {workspace = true, default-features = false} does not work -span = { path = "../span", version = "0.0.0", default-features = false} +span = { path = "../span", version = "0.0.0", default-features = false } intern.workspace = true diff --git a/crates/proc-macro-api/src/legacy_protocol/task_impl.rs b/crates/proc-macro-api/src/legacy_protocol/task_impl.rs new file mode 100644 index 000000000000..252de6104c9e --- /dev/null +++ b/crates/proc-macro-api/src/legacy_protocol/task_impl.rs @@ -0,0 +1,64 @@ +//! Implement how to send and receive legacy protocol message on stdio stream(protocol layer) +//! +//! The implementation follows: +//! 1. Send Request +//! 2. Receive Response + +use std::{ + io::{self, BufRead, Write}, + sync::Arc, +}; + +use crate::{ + ServerError, + legacy_protocol::{ + json::{read_json, write_json}, + msg::{Request, Response}, + }, + task::TaskClient, +}; + +/// Sends a request to the server and reads the response. +fn send_request( + mut writer: &mut dyn Write, + mut reader: &mut dyn BufRead, + req: Request, + buf: &mut String, +) -> Result, ServerError> { + use crate::legacy_protocol::msg::Message; + + req.write(write_json, &mut writer).map_err(|err| ServerError { + message: "failed to write request".into(), + io: Some(Arc::new(err)), + })?; + let res = Response::read(read_json, &mut reader, buf).map_err(|err| ServerError { + message: "failed to read response".into(), + io: Some(Arc::new(err)), + })?; + Ok(res) +} + +pub struct JsonTaskClient<'a> { + pub writer: &'a mut dyn Write, + pub reader: &'a mut dyn BufRead, + pub buf: &'a mut String, +} + +impl TaskClient for JsonTaskClient<'_> { + type Task = Request; + type TaskResult = Response; + + // Implement send_task for Json Legacy Protocol + // Basically what we have done in process.rs + fn send_task(&mut self, task: Self::Task) -> Result { + send_request(self.writer, self.reader, task, self.buf).and_then(|res| { + res.ok_or_else(|| { + let message = "proc-macro server did not respond with data".to_owned(); + ServerError { + io: Some(Arc::new(io::Error::new(io::ErrorKind::BrokenPipe, message.clone()))), + message, + } + }) + }) + } +} diff --git a/crates/proc-macro-api/src/lib.rs b/crates/proc-macro-api/src/lib.rs index 25c30b6db4a4..d289b8cae2f8 100644 --- a/crates/proc-macro-api/src/lib.rs +++ b/crates/proc-macro-api/src/lib.rs @@ -8,7 +8,19 @@ pub mod legacy_protocol { pub mod json; pub mod msg; + pub mod task_impl; } + +#[allow(dead_code)] +pub mod new_protocol { + pub mod msg; + pub mod postcard; + pub mod task_impl; +} + +#[allow(dead_code)] +pub mod task; + mod process; use paths::{AbsPath, AbsPathBuf}; diff --git a/crates/proc-macro-api/src/new_protocol/msg.rs b/crates/proc-macro-api/src/new_protocol/msg.rs new file mode 100644 index 000000000000..634a8628a39e --- /dev/null +++ b/crates/proc-macro-api/src/new_protocol/msg.rs @@ -0,0 +1,137 @@ +//! New Protocol Message Definitions covers the functionalities of `legacy_protocol` and with several changes: +//! +//! - Using [postcard](https://github.com/jamesmunns/postcard) as serde library, which provides binary, lightweight serialization. +//! - As we change to postcard, no need to use FlatTree, use `tt::TopSubtree` directly +//! +//! One possible communication may look like this: +//! client ------------------------- server +//! >-------Request---------> +//! <-------Query-----------< +//! >-------Reply-----------> +//! <-------Response-------< + +use paths::Utf8PathBuf; +use serde::{Deserialize, Serialize}; + +use crate::ProcMacroKind; +use crate::legacy_protocol::msg::FlatTree; + +/// Represents messages send from client to proc-macr-srv +#[derive(Debug, Serialize, Deserialize)] +pub enum C2SMsg { + Request(Request), + // NOTE: Reply is left empty as a placeholder + Reply, +} + +/// Represents messages send from client to proc-macr-srv +#[derive(Debug, Serialize, Deserialize)] +pub enum S2CMsg { + Response(Response), + // NOTE: Query is left empty as a place holder + Query, +} + +// NOTE: Directly copied from legacy_protocol +#[derive(Debug, Serialize, Deserialize)] +pub enum Request { + ListMacros { dylib_path: Utf8PathBuf }, + ExpandMacro(Box), + ApiVersionCheck {}, + SetConfig(ServerConfig), +} + +// NOTE: Directly copied from legacy_protocol +#[derive(Debug, Serialize, Deserialize)] +pub enum Response { + ListMacros(Result, String>), + ExpandMacro(Result), + ApiVersionCheck(u32), + SetConfig(ServerConfig), + ExpandMacroExtended(Result), +} + +// NOTE: Directly copied from legacy_protocol +#[derive(Debug, Serialize, Deserialize, Default)] +#[serde(default)] +pub struct ServerConfig { + pub span_mode: SpanMode, +} + +// NOTE: Directly copied from legacy_protocol, +// except the tree field +#[derive(Debug, Serialize, Deserialize)] +pub struct ExpandMacroExtended { + /// The expanded syntax tree. + pub tree: TreeWrapper, + /// Additional span data mappings. + pub span_data_table: Vec, +} + +// NOTE: Directly copied from legacy_protocol +#[derive(Debug, Serialize, Deserialize)] +pub struct PanicMessage(pub String); + +// NOTE: Directly copied from legacy_protocol +#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)] +pub enum SpanMode { + #[default] + Id, + RustAnalyzer, +} + +// NOTE: Directly copied from legacy_protocol +#[derive(Debug, Serialize, Deserialize)] +pub struct ExpandMacro { + pub lib: Utf8PathBuf, + + pub env: Vec<(String, String)>, + + pub current_dir: Option, + + #[serde(flatten)] + pub data: ExpandMacroData, +} + +// TODO: Maybe ideal we want to ser/de TokenTree directly +// +// Something like this should be ideal +// pub struct TreeWrapper(tt::TokenTree); +// currently, just use this wrapper to build the backbones +#[derive(Debug, Serialize, Deserialize)] +pub struct TreeWrapper(FlatTree); + +// NOTE: Directly copied from legacy_protocol, +// except the tree field +#[derive(Debug, Serialize, Deserialize)] +pub struct ExpandMacroData { + pub macro_body: TreeWrapper, + + pub macro_name: String, + + pub attributes: Option, + #[serde(skip_serializing_if = "ExpnGlobals::skip_serializing_if")] + #[serde(default)] + pub has_global_spans: ExpnGlobals, + + #[serde(skip_serializing_if = "Vec::is_empty")] + #[serde(default)] + pub span_data_table: Vec, +} + +// NOTE: Directly copied from legacy_protocol, +#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)] +pub struct ExpnGlobals { + #[serde(skip_serializing)] + #[serde(default)] + pub serialize: bool, + pub def_site: usize, + pub call_site: usize, + pub mixed_site: usize, +} + +impl ExpnGlobals { + fn skip_serializing_if(&self) -> bool { + !self.serialize + } +} diff --git a/crates/proc-macro-api/src/new_protocol/postcard.rs b/crates/proc-macro-api/src/new_protocol/postcard.rs new file mode 100644 index 000000000000..0f78a3b87642 --- /dev/null +++ b/crates/proc-macro-api/src/new_protocol/postcard.rs @@ -0,0 +1,42 @@ +//! Implement how to encoding and decoding new protocol message on stdio stream +//! +//! Current implementation encoded message to | message length | message content | + +use serde::de::DeserializeOwned; +use std::io::{self, BufRead, Read, Write}; + +use super::msg::{C2SMsg, S2CMsg}; + +fn read_usize_be(reader: &mut R) -> io::Result { + let mut buf = [0u8; std::mem::size_of::()]; + reader.read_exact(&mut buf)?; + Ok(usize::from_be_bytes(buf)) +} + +fn write_usize_be(writer: &mut W, value: usize) -> io::Result<()> { + writer.write_all(&value.to_be_bytes()) // Convert and write as Big-Endian +} + +pub trait ProtoPostcard: serde::Serialize + DeserializeOwned { + fn receive_proto(reader: &mut R) -> io::Result { + let msg_len = read_usize_be(reader)? as usize; + let mut buf = vec![0u8; msg_len]; + + reader.read_exact(&mut buf)?; + postcard::from_bytes::(&buf) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e)) + } + + fn send_proto(&self, writer: &mut W) -> io::Result<()> { + let bytes = postcard::to_allocvec(self) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + + write_usize_be(writer, bytes.len())?; + writer.write_all(&bytes)?; + writer.flush() + } +} + +// NOTE: With default implementation, C2SMsg and S2CMsg can both be send through stdio +impl ProtoPostcard for C2SMsg {} +impl ProtoPostcard for S2CMsg {} diff --git a/crates/proc-macro-api/src/new_protocol/task_impl.rs b/crates/proc-macro-api/src/new_protocol/task_impl.rs new file mode 100644 index 000000000000..f3e71adac92a --- /dev/null +++ b/crates/proc-macro-api/src/new_protocol/task_impl.rs @@ -0,0 +1,66 @@ +//! Implement how to send and receive new protocol message on stdio stream(protocol layer) +//! +//! Current implementation: +//! 1. Send Request +//! 2. Handle possible query and send reply +//! 3. Receive Response + +use std::io::{BufRead, Write}; + +use crate::{ + ServerError, + new_protocol::{ + msg::{C2SMsg, S2CMsg}, + postcard::ProtoPostcard, + }, + task::TaskClient, +}; + +pub struct PostcardTaskClient<'a> { + pub writer: &'a mut dyn Write, + pub reader: &'a mut dyn BufRead, +} + +impl TaskClient for PostcardTaskClient<'_> { + type Task = C2SMsg; + type TaskResult = S2CMsg; + + // Implement send_task for Postcard Protocol + // As this new Protocol will allow back and forth communication + // send_task will abstract these back and forth details + // Task will be C2SMsg::Request, which is basically legacy_protocol Request + // TaskResult will be S2CMsg::Response, which is basically legacy_protocol Response + fn send_task(&mut self, task: Self::Task) -> Result { + // Send the initial C2SMsg (client-to-server message) + task.send_proto(self.writer).map_err(|err| ServerError { + message: "failed to send request".into(), + io: Some(std::sync::Arc::new(err)), + })?; + + // Handle bidirectional communication loop + loop { + // Receive response from server + let server_msg = S2CMsg::receive_proto(self.reader).map_err(|err| ServerError { + message: "failed to read server message".into(), + io: Some(std::sync::Arc::new(err)), + })?; + + match server_msg { + S2CMsg::Response(_) => { + // Final response received, return it + return Ok(server_msg); + } + S2CMsg::Query => { + // Server sent a query, we need to send a reply + // For now, send an empty Reply as placeholder + let reply = C2SMsg::Reply; + reply.send_proto(self.writer).map_err(|err| ServerError { + message: "failed to send reply".into(), + io: Some(std::sync::Arc::new(err)), + })?; + // Continue the loop to wait for the next server message + } + } + } + } +} diff --git a/crates/proc-macro-api/src/process.rs b/crates/proc-macro-api/src/process.rs index fcea75ef672a..3d06ec604956 100644 --- a/crates/proc-macro-api/src/process.rs +++ b/crates/proc-macro-api/src/process.rs @@ -1,10 +1,10 @@ //! Handle process life-time and message passing for proc-macro client use std::{ - io::{self, BufRead, BufReader, Read, Write}, + io::{self, BufReader, Read}, panic::AssertUnwindSafe, process::{Child, ChildStdin, ChildStdout, Command, Stdio}, - sync::{Arc, Mutex, OnceLock}, + sync::{Mutex, OnceLock}, }; use paths::AbsPath; @@ -13,12 +13,13 @@ use stdx::JodChild; use crate::{ ProcMacroKind, ServerError, legacy_protocol::{ - json::{read_json, write_json}, msg::{ - CURRENT_API_VERSION, Message, RUST_ANALYZER_SPAN_SUPPORT, Request, Response, - ServerConfig, SpanMode, + CURRENT_API_VERSION, RUST_ANALYZER_SPAN_SUPPORT, Request, Response, ServerConfig, + SpanMode, }, + task_impl::JsonTaskClient, }, + task::TaskClient, }; /// Represents a process handling proc-macro communication. @@ -142,47 +143,61 @@ impl ProcMacroServerProcess { } let state = &mut *self.state.lock().unwrap(); - let mut buf = String::new(); - send_request(&mut state.stdin, &mut state.stdout, req, &mut buf) - .and_then(|res| { - res.ok_or_else(|| { - let message = "proc-macro server did not respond with data".to_owned(); - ServerError { - io: Some(Arc::new(io::Error::new( - io::ErrorKind::BrokenPipe, - message.clone(), - ))), - message, - } - }) - }) - .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() { - if let Some(stderr) = state.process.child.stderr.as_mut() { - _ = stderr.read_to_string(&mut msg); - } + // Check environment variable to determine which protocol to use + let protocol = std::env::var("RUST_ANALYZER_PROC_MACRO_PROTOCOL") + .unwrap_or_else(|_| "json".to_owned()); + + let result = match protocol.as_str() { + "postcard" => { + tracing::warn!("Postcard protocol requested but not fully implemented, using JSON"); + + let mut buf = String::new(); + let mut client = JsonTaskClient { + writer: &mut state.stdin, + reader: &mut state.stdout, + buf: &mut buf, + }; + client.send_task(req) + } + _ => { + // Default to JSON protocol + let mut buf = String::new(); + let mut client = JsonTaskClient { + writer: &mut state.stdin, + reader: &mut state.stdout, + buf: &mut buf, + }; + client.send_task(req) + } + }; + + result.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() { + if 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, - }; - // `AssertUnwindSafe` is fine here, we already correct initialized - // server_error at this point. - self.exited.get_or_init(|| AssertUnwindSafe(server_error)).0.clone() } + let server_error = ServerError { + message: format!( + "proc-macro server exited with {status}{}{msg}", + if msg.is_empty() { "" } else { ": " } + ), + io: None, + }; + // `AssertUnwindSafe` is fine here, we already correct initialized + // server_error at this point. + self.exited.get_or_init(|| AssertUnwindSafe(server_error)).0.clone() } - } else { - e } - }) + } else { + e + } + }) } } @@ -223,6 +238,26 @@ fn mk_child<'a>( ) -> io::Result { #[allow(clippy::disallowed_methods)] let mut cmd = Command::new(path); + + // Check for protocol selection environment variable + if let Ok(protocol) = std::env::var("RUST_ANALYZER_PROC_MACRO_PROTOCOL") { + match protocol.as_str() { + "postcard" => { + cmd.args(["--format", "postcard"]); + } + "json" => { + cmd.args(["--format", "json"]); + } + _ => { + tracing::warn!("Unknown protocol '{}', defaulting to json", protocol); + cmd.args(["--format", "json"]); + } + } + } else { + // Default to JSON protocol for backward compatibility + cmd.args(["--format", "json"]); + } + for env in extra_env { match env { (key, Some(val)) => cmd.env(key, val), @@ -242,21 +277,3 @@ fn mk_child<'a>( } cmd.spawn() } - -/// Sends a request to the server and reads the response. -fn send_request( - mut writer: &mut impl Write, - mut reader: &mut impl BufRead, - req: Request, - buf: &mut String, -) -> Result, ServerError> { - req.write(write_json, &mut writer).map_err(|err| ServerError { - message: "failed to write request".into(), - io: Some(Arc::new(err)), - })?; - let res = Response::read(read_json, &mut reader, buf).map_err(|err| ServerError { - message: "failed to read response".into(), - io: Some(Arc::new(err)), - })?; - Ok(res) -} diff --git a/crates/proc-macro-api/src/task.rs b/crates/proc-macro-api/src/task.rs new file mode 100644 index 000000000000..109a7b396c75 --- /dev/null +++ b/crates/proc-macro-api/src/task.rs @@ -0,0 +1,12 @@ +//! task.rs is used to abstract ProcMacroServerProcess::send_task in process.rs +//! The ultimate goal is to encapsulate legacy_protocol and don't directly use legacy_protocol in process.rs +//! + +use crate::ServerError; + +pub trait TaskClient { + type Task; + type TaskResult; + + fn send_task(&mut self, task: Self::Task) -> Result; +} diff --git a/crates/proc-macro-srv-cli/src/main.rs b/crates/proc-macro-srv-cli/src/main.rs index c47ed053254b..2a9d69b78184 100644 --- a/crates/proc-macro-srv-cli/src/main.rs +++ b/crates/proc-macro-srv-cli/src/main.rs @@ -9,6 +9,10 @@ extern crate rustc_driver as _; #[cfg(any(feature = "sysroot-abi", rust_analyzer))] mod main_loop; + +#[cfg(any(feature = "sysroot-abi", rust_analyzer))] +mod task_executor; + #[cfg(any(feature = "sysroot-abi", rust_analyzer))] use main_loop::run; diff --git a/crates/proc-macro-srv-cli/src/main_loop.rs b/crates/proc-macro-srv-cli/src/main_loop.rs index f54dff1f2d82..2be5c38fbb46 100644 --- a/crates/proc-macro-srv-cli/src/main_loop.rs +++ b/crates/proc-macro-srv-cli/src/main_loop.rs @@ -1,132 +1,39 @@ //! The main loop of the proc-macro server. use std::io; -use proc_macro_api::legacy_protocol::{ - json::{read_json, write_json}, - msg::{ - self, CURRENT_API_VERSION, ExpandMacroData, ExpnGlobals, Message, SpanMode, TokenId, - deserialize_span_data_index_map, serialize_span_data_index_map, - }, +use crate::task_executor::{ + TaskExecutor, json_task_executor::JsonTaskExecutor, + postcard_task_executor::PostcardTaskExecutor, }; -use proc_macro_srv::EnvSnapshot; pub(crate) fn run() -> io::Result<()> { - fn macro_kind_to_api(kind: proc_macro_srv::ProcMacroKind) -> proc_macro_api::ProcMacroKind { - match kind { - proc_macro_srv::ProcMacroKind::CustomDerive => { - proc_macro_api::ProcMacroKind::CustomDerive - } - proc_macro_srv::ProcMacroKind::Bang => proc_macro_api::ProcMacroKind::Bang, - proc_macro_srv::ProcMacroKind::Attr => proc_macro_api::ProcMacroKind::Attr, + // Check environment variable or command line args to determine protocol + let args: Vec = std::env::args().collect(); + let mut format = None; + + // Parse --format argument + for i in 1..args.len() { + if args[i] == "--format" && i + 1 < args.len() { + format = Some(args[i + 1].as_str()); + break; } } - let mut buf = String::new(); - let mut read_request = || msg::Request::read(read_json, &mut io::stdin().lock(), &mut buf); - let write_response = |msg: msg::Response| msg.write(write_json, &mut io::stdout().lock()); + // Default to JSON for backward compatibility + let protocol = format.unwrap_or("json"); - let env = EnvSnapshot::default(); - let srv = proc_macro_srv::ProcMacroSrv::new(&env); - - let mut span_mode = SpanMode::Id; - - 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| { - macros.into_iter().map(|(name, kind)| (name, macro_kind_to_api(kind))).collect() - })) - } - msg::Request::ExpandMacro(task) => { - let msg::ExpandMacro { - lib, - env, - current_dir, - data: - ExpandMacroData { - macro_body, - macro_name, - attributes, - has_global_spans: - ExpnGlobals { serialize: _, def_site, call_site, mixed_site }, - span_data_table, - }, - } = *task; - match span_mode { - SpanMode::Id => msg::Response::ExpandMacro({ - let def_site = TokenId(def_site as u32); - let call_site = TokenId(call_site as u32); - let mixed_site = TokenId(mixed_site as u32); - - let macro_body = macro_body.to_subtree_unresolved(CURRENT_API_VERSION); - let attributes = - attributes.map(|it| it.to_subtree_unresolved(CURRENT_API_VERSION)); - - srv.expand( - lib, - env, - current_dir, - macro_name, - macro_body, - attributes, - def_site, - call_site, - mixed_site, - ) - .map(|it| { - msg::FlatTree::new_raw(tt::SubtreeView::new(&it), CURRENT_API_VERSION) - }) - .map_err(msg::PanicMessage) - }), - SpanMode::RustAnalyzer => msg::Response::ExpandMacroExtended({ - let mut span_data_table = deserialize_span_data_index_map(&span_data_table); - - let def_site = span_data_table[def_site]; - let call_site = span_data_table[call_site]; - let mixed_site = span_data_table[mixed_site]; - - let macro_body = - macro_body.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table); - let attributes = attributes.map(|it| { - it.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table) - }); - srv.expand( - lib, - env, - current_dir, - macro_name, - macro_body, - attributes, - def_site, - call_site, - mixed_site, - ) - .map(|it| { - ( - msg::FlatTree::new( - tt::SubtreeView::new(&it), - CURRENT_API_VERSION, - &mut span_data_table, - ), - serialize_span_data_index_map(&span_data_table), - ) - }) - .map(|(tree, span_data_table)| msg::ExpandMacroExtended { - tree, - span_data_table, - }) - .map_err(msg::PanicMessage) - }), - } - } - msg::Request::ApiVersionCheck {} => msg::Response::ApiVersionCheck(CURRENT_API_VERSION), - msg::Request::SetConfig(config) => { - span_mode = config.span_mode; - msg::Response::SetConfig(config) - } - }; - write_response(res)? + match protocol { + "json" => { + let executor = JsonTaskExecutor; + executor.run() + } + "postcard" => { + let executor = PostcardTaskExecutor; + executor.run() + } + _ => Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("Unsupported protocol format: {}. Supported formats: json, postcard", protocol), + )), } - - Ok(()) } diff --git a/crates/proc-macro-srv-cli/src/task_executor.rs b/crates/proc-macro-srv-cli/src/task_executor.rs new file mode 100644 index 000000000000..ebf956a91ff3 --- /dev/null +++ b/crates/proc-macro-srv-cli/src/task_executor.rs @@ -0,0 +1,14 @@ +//! TaskExecutor trait and implementations for server-side protocol handling +//! This mirrors the client-side TaskClient pattern + +use std::io; + +pub(crate) mod json_task_executor; +#[allow(dead_code)] +pub(crate) mod postcard_task_executor; + +/// Generic trait for executing task processing loops on the server side +pub(crate) trait TaskExecutor { + /// Run the main server loop for processing client requests + fn run(&self) -> io::Result<()>; +} diff --git a/crates/proc-macro-srv-cli/src/task_executor/json_task_executor.rs b/crates/proc-macro-srv-cli/src/task_executor/json_task_executor.rs new file mode 100644 index 000000000000..02786de70c39 --- /dev/null +++ b/crates/proc-macro-srv-cli/src/task_executor/json_task_executor.rs @@ -0,0 +1,148 @@ +//! JSON protocol task executor for legacy protocol support + +use std::io; + +use proc_macro_api::legacy_protocol::{ + json::{read_json, write_json}, + msg::{ + self, CURRENT_API_VERSION, ExpandMacroData, ExpnGlobals, Message, SpanMode, TokenId, + deserialize_span_data_index_map, serialize_span_data_index_map, + }, +}; +use proc_macro_srv::EnvSnapshot; + +use crate::task_executor::TaskExecutor; + +pub(crate) struct JsonTaskExecutor; + +impl TaskExecutor for JsonTaskExecutor { + fn run(&self) -> io::Result<()> { + fn macro_kind_to_api(kind: proc_macro_srv::ProcMacroKind) -> proc_macro_api::ProcMacroKind { + match kind { + proc_macro_srv::ProcMacroKind::CustomDerive => { + proc_macro_api::ProcMacroKind::CustomDerive + } + proc_macro_srv::ProcMacroKind::Bang => proc_macro_api::ProcMacroKind::Bang, + proc_macro_srv::ProcMacroKind::Attr => proc_macro_api::ProcMacroKind::Attr, + } + } + + let mut buf = String::new(); + let mut read_request = || msg::Request::read(read_json, &mut io::stdin().lock(), &mut buf); + let write_response = |msg: msg::Response| msg.write(write_json, &mut io::stdout().lock()); + + let env = EnvSnapshot::default(); + let srv = proc_macro_srv::ProcMacroSrv::new(&env); + + let mut span_mode = SpanMode::Id; + + 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| { + macros + .into_iter() + .map(|(name, kind)| (name, macro_kind_to_api(kind))) + .collect() + })) + } + msg::Request::ExpandMacro(task) => { + let msg::ExpandMacro { + lib, + env, + current_dir, + data: + ExpandMacroData { + macro_body, + macro_name, + attributes, + has_global_spans: + ExpnGlobals { serialize: _, def_site, call_site, mixed_site }, + span_data_table, + }, + } = *task; + match span_mode { + SpanMode::Id => msg::Response::ExpandMacro({ + let def_site = TokenId(def_site as u32); + let call_site = TokenId(call_site as u32); + let mixed_site = TokenId(mixed_site as u32); + + let macro_body = macro_body.to_subtree_unresolved(CURRENT_API_VERSION); + let attributes = + attributes.map(|it| it.to_subtree_unresolved(CURRENT_API_VERSION)); + + srv.expand( + lib, + env, + current_dir, + macro_name, + macro_body, + attributes, + def_site, + call_site, + mixed_site, + ) + .map(|it| { + msg::FlatTree::new_raw( + tt::SubtreeView::new(&it), + CURRENT_API_VERSION, + ) + }) + .map_err(msg::PanicMessage) + }), + SpanMode::RustAnalyzer => msg::Response::ExpandMacroExtended({ + let mut span_data_table = + deserialize_span_data_index_map(&span_data_table); + + let def_site = span_data_table[def_site]; + let call_site = span_data_table[call_site]; + let mixed_site = span_data_table[mixed_site]; + + let macro_body = macro_body + .to_subtree_resolved(CURRENT_API_VERSION, &span_data_table); + let attributes = attributes.map(|it| { + it.to_subtree_resolved(CURRENT_API_VERSION, &span_data_table) + }); + srv.expand( + lib, + env, + current_dir, + macro_name, + macro_body, + attributes, + def_site, + call_site, + mixed_site, + ) + .map(|it| { + ( + msg::FlatTree::new( + tt::SubtreeView::new(&it), + CURRENT_API_VERSION, + &mut span_data_table, + ), + serialize_span_data_index_map(&span_data_table), + ) + }) + .map(|(tree, span_data_table)| msg::ExpandMacroExtended { + tree, + span_data_table, + }) + .map_err(msg::PanicMessage) + }), + } + } + msg::Request::ApiVersionCheck {} => { + msg::Response::ApiVersionCheck(CURRENT_API_VERSION) + } + msg::Request::SetConfig(config) => { + span_mode = config.span_mode; + msg::Response::SetConfig(config) + } + }; + write_response(res)? + } + + Ok(()) + } +} diff --git a/crates/proc-macro-srv-cli/src/task_executor/postcard_task_executor.rs b/crates/proc-macro-srv-cli/src/task_executor/postcard_task_executor.rs new file mode 100644 index 000000000000..0116e25d3561 --- /dev/null +++ b/crates/proc-macro-srv-cli/src/task_executor/postcard_task_executor.rs @@ -0,0 +1,83 @@ +//! Postcard binary protocol task executor for new protocol support + +use std::io; + +use proc_macro_api::new_protocol::{ + msg::{C2SMsg, PanicMessage, Request, Response, S2CMsg}, + postcard::ProtoPostcard, +}; +use proc_macro_srv::EnvSnapshot; + +use crate::task_executor::TaskExecutor; + +pub(crate) struct PostcardTaskExecutor; + +impl TaskExecutor for PostcardTaskExecutor { + fn run(&self) -> io::Result<()> { + fn macro_kind_to_api(kind: proc_macro_srv::ProcMacroKind) -> proc_macro_api::ProcMacroKind { + match kind { + proc_macro_srv::ProcMacroKind::CustomDerive => { + proc_macro_api::ProcMacroKind::CustomDerive + } + proc_macro_srv::ProcMacroKind::Bang => proc_macro_api::ProcMacroKind::Bang, + proc_macro_srv::ProcMacroKind::Attr => proc_macro_api::ProcMacroKind::Attr, + } + } + + let env = EnvSnapshot::default(); + let srv = proc_macro_srv::ProcMacroSrv::new(&env); + + let mut stdin = io::stdin().lock(); + let mut stdout = io::stdout().lock(); + + // Main communication loop + loop { + // Receive C2SMsg from client + let client_msg = match C2SMsg::receive_proto(&mut stdin) { + Ok(msg) => msg, + Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => { + // Client disconnected, exit gracefully + break; + } + Err(e) => return Err(e), + }; + + match client_msg { + C2SMsg::Request(request) => { + // Process the request and send response + let response = match request { + Request::ListMacros { dylib_path } => { + Response::ListMacros(srv.list_macros(&dylib_path).map(|macros| { + macros + .into_iter() + .map(|(name, kind)| (name, macro_kind_to_api(kind))) + .collect() + })) + } + Request::ExpandMacro(_task) => { + // TODO: Implement macro expansion for new protocol + // This would need to convert TreeWrapper to appropriate types + // and handle the expansion similar to JsonTaskExecutor + Response::ExpandMacro(Err(PanicMessage("Not Yet Implemented".into()))) + } + Request::ApiVersionCheck {} => Response::ApiVersionCheck( + proc_macro_api::legacy_protocol::msg::CURRENT_API_VERSION, + ), + Request::SetConfig(config) => Response::SetConfig(config), + }; + + // Send the response back to client + let server_msg = S2CMsg::Response(response); + server_msg.send_proto(&mut stdout)?; + } + C2SMsg::Reply => { + // Handle reply from client (in response to our query) + // For now, this is a placeholder as we don't send queries yet + // In future, this would handle bidirectional communication + } + } + } + + Ok(()) + } +}