Abtract away json protocol for proc-macro-srv

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

@@ -9,6 +9,7 @@ extern crate rustc_driver as _;
 
 #[cfg(any(feature = "sysroot-abi", rust_analyzer))]
 mod main_loop;
+use clap::{Command, ValueEnum};
 #[cfg(any(feature = "sysroot-abi", rust_analyzer))]
 use main_loop::run;
 
@@ -23,12 +24,46 @@ fn main() -> std::io::Result<()> {
         );
         std::process::exit(122);
     }
+    let matches = Command::new("proc-macro-srv")
+        .args(&[clap::Arg::new("format")
+            .long("format")
+            .action(clap::ArgAction::Set)
+            .default_value("json")
+            .value_parser(clap::builder::EnumValueParser::<ProtocolFormat>::new())])
+        .get_matches();
+    let &format =
+        matches.get_one::<ProtocolFormat>("format").expect("format value should always be present");
+    run(format)
+}
 
-    run()
+#[derive(Copy, Clone)]
+enum ProtocolFormat {
+    Json,
+    Postcard,
+}
+
+impl ValueEnum for ProtocolFormat {
+    fn value_variants<'a>() -> &'a [Self] {
+        &[ProtocolFormat::Json]
+    }
+
+    fn to_possible_value(&self) -> Option<clap::builder::PossibleValue> {
+        match self {
+            ProtocolFormat::Json => Some(clap::builder::PossibleValue::new("json")),
+            ProtocolFormat::Postcard => Some(clap::builder::PossibleValue::new("postcard")),
+        }
+    }
+    fn from_str(input: &str, _ignore_case: bool) -> Result<Self, String> {
+        match input {
+            "json" => Ok(ProtocolFormat::Json),
+            "postcard" => Ok(ProtocolFormat::Postcard),
+            _ => Err(format!("unknown protocol format: {input}")),
+        }
+    }
 }
 
 #[cfg(not(any(feature = "sysroot-abi", rust_analyzer)))]
-fn run() -> std::io::Result<()> {
+fn run(_: ProtocolFormat) -> std::io::Result<()> {
     Err(std::io::Error::new(
         std::io::ErrorKind::Unsupported,
         "proc-macro-srv-cli needs to be compiled with the `sysroot-abi` feature to function"

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

@@ -1,16 +1,47 @@
 //! 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 proc_macro_api::{
+    legacy_protocol::{
+        json::{read_json, write_json},
+        msg::{
+            self, ExpandMacroData, ExpnGlobals, Message, SpanMode, SpanTransformer,
+            deserialize_span_data_index_map, serialize_span_data_index_map,
+        },
     },
+    version::CURRENT_API_VERSION,
 };
-use proc_macro_srv::EnvSnapshot;
+use proc_macro_srv::{EnvSnapshot, SpanId};
 
-pub(crate) fn run() -> io::Result<()> {
+use crate::ProtocolFormat;
+
+struct SpanTrans;
+
+impl SpanTransformer for SpanTrans {
+    type Table = ();
+    type Span = SpanId;
+    fn token_id_of(
+        _: &mut Self::Table,
+        span: Self::Span,
+    ) -> proc_macro_api::legacy_protocol::SpanId {
+        proc_macro_api::legacy_protocol::SpanId(span.0 as u32)
+    }
+    fn span_for_token_id(
+        _: &Self::Table,
+        id: proc_macro_api::legacy_protocol::SpanId,
+    ) -> Self::Span {
+        SpanId(id.0 as u32)
+    }
+}
+
+pub(crate) fn run(format: ProtocolFormat) -> io::Result<()> {
+    match format {
+        ProtocolFormat::Json => run_json(),
+        ProtocolFormat::Postcard => unimplemented!(),
+    }
+}
+
+fn run_json() -> io::Result<()> {
     fn macro_kind_to_api(kind: proc_macro_srv::ProcMacroKind) -> proc_macro_api::ProcMacroKind {
         match kind {
             proc_macro_srv::ProcMacroKind::CustomDerive => {
@@ -54,13 +85,14 @@ pub(crate) fn run() -> io::Result<()> {
                 } = *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 def_site = SpanId(def_site as u32);
+                        let call_site = SpanId(call_site as u32);
+                        let mixed_site = SpanId(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));
+                        let macro_body =
+                            macro_body.to_subtree_unresolved::<SpanTrans>(CURRENT_API_VERSION);
+                        let attributes = attributes
+                            .map(|it| it.to_subtree_unresolved::<SpanTrans>(CURRENT_API_VERSION));
 
                         srv.expand(
                             lib,
@@ -74,7 +106,10 @@ pub(crate) fn run() -> io::Result<()> {
                             mixed_site,
                         )
                         .map(|it| {
-                            msg::FlatTree::new_raw(tt::SubtreeView::new(&it), CURRENT_API_VERSION)
+                            msg::FlatTree::new_raw::<SpanTrans>(
+                                tt::SubtreeView::new(&it),
+                                CURRENT_API_VERSION,
+                            )
                         })
                         .map_err(msg::PanicMessage)
                     }),