From e003bc8bfd8c834fd8c744be56ab35986493170b Mon Sep 17 00:00:00 2001
From: vlad20012 <beskvlad@gmail.com>
Date: Thu, 7 Dec 2023 17:49:47 +0100
Subject: [PATCH] Add a request for setting response prefix

---
 crates/proc-macro-api/src/lib.rs      |  4 +++-
 crates/proc-macro-api/src/msg.rs      |  5 ++++-
 crates/proc-macro-api/src/process.rs  |  8 ++++++--
 crates/proc-macro-srv-cli/src/main.rs | 18 +++++++++++++++++-
 crates/proc-macro-srv/src/lib.rs      |  1 +
 5 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/crates/proc-macro-api/src/lib.rs b/crates/proc-macro-api/src/lib.rs
index f697ecd3518f..0f2bb1f25bca 100644
--- a/crates/proc-macro-api/src/lib.rs
+++ b/crates/proc-macro-api/src/lib.rs
@@ -178,7 +178,9 @@ impl ProcMacro {
             msg::Response::ExpandMacro(it) => {
                 Ok(it.map(|tree| FlatTree::to_subtree_resolved(tree, version, &span_data_table)))
             }
-            msg::Response::ListMacros(..) | msg::Response::ApiVersionCheck(..) => {
+            msg::Response::ListMacros(..)
+            | msg::Response::ApiVersionCheck(..)
+            | msg::Response::SetExpanderSettings { .. } => {
                 Err(ServerError { message: "unexpected response".to_string(), io: None })
             }
         }
diff --git a/crates/proc-macro-api/src/msg.rs b/crates/proc-macro-api/src/msg.rs
index 1d3e45aff385..e56b66b173a7 100644
--- a/crates/proc-macro-api/src/msg.rs
+++ b/crates/proc-macro-api/src/msg.rs
@@ -17,14 +17,16 @@ pub const NO_VERSION_CHECK_VERSION: u32 = 0;
 pub const VERSION_CHECK_VERSION: u32 = 1;
 pub const ENCODE_CLOSE_SPAN_VERSION: u32 = 2;
 pub const HAS_GLOBAL_SPANS: u32 = 3;
+pub const EXPAND_MACRO_RESPONSE_PREFIX: u32 = 4;
 
-pub const CURRENT_API_VERSION: u32 = HAS_GLOBAL_SPANS;
+pub const CURRENT_API_VERSION: u32 = EXPAND_MACRO_RESPONSE_PREFIX;
 
 #[derive(Debug, Serialize, Deserialize)]
 pub enum Request {
     ListMacros { dylib_path: PathBuf },
     ExpandMacro(ExpandMacro),
     ApiVersionCheck {},
+    SetExpanderSettings { response_prefix: String },
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -32,6 +34,7 @@ pub enum Response {
     ListMacros(Result<Vec<(String, ProcMacroKind)>, String>),
     ExpandMacro(Result<FlatTree, PanicMessage>),
     ApiVersionCheck(u32),
+    SetExpanderSettings {},
 }
 
 #[derive(Debug, Serialize, Deserialize)]
diff --git a/crates/proc-macro-api/src/process.rs b/crates/proc-macro-api/src/process.rs
index 9a20fa63ed70..bc7345c091f4 100644
--- a/crates/proc-macro-api/src/process.rs
+++ b/crates/proc-macro-api/src/process.rs
@@ -62,7 +62,9 @@ impl ProcMacroProcessSrv {
 
         match response {
             Response::ApiVersionCheck(version) => Ok(version),
-            Response::ExpandMacro { .. } | Response::ListMacros { .. } => {
+            Response::ExpandMacro { .. }
+            | Response::ListMacros { .. }
+            | Response::SetExpanderSettings { .. } => {
                 Err(ServerError { message: "unexpected response".to_string(), io: None })
             }
         }
@@ -78,7 +80,9 @@ impl ProcMacroProcessSrv {
 
         match response {
             Response::ListMacros(it) => Ok(it),
-            Response::ExpandMacro { .. } | Response::ApiVersionCheck { .. } => {
+            Response::ExpandMacro { .. }
+            | Response::ApiVersionCheck { .. }
+            | Response::SetExpanderSettings { .. } => {
                 Err(ServerError { message: "unexpected response".to_string(), io: None })
             }
         }
diff --git a/crates/proc-macro-srv-cli/src/main.rs b/crates/proc-macro-srv-cli/src/main.rs
index 50ce586fc429..95d4f6cf3149 100644
--- a/crates/proc-macro-srv-cli/src/main.rs
+++ b/crates/proc-macro-srv-cli/src/main.rs
@@ -26,10 +26,22 @@ fn run() -> io::Result<()> {
 #[cfg(any(feature = "sysroot-abi", rust_analyzer))]
 fn run() -> io::Result<()> {
     use proc_macro_api::msg::{self, Message};
+    use std::cell::RefCell;
+
+    let current_response_prefix = RefCell::new(String::new());
 
     let read_request = |buf: &mut String| msg::Request::read(&mut io::stdin().lock(), buf);
 
-    let write_response = |msg: msg::Response| msg.write(&mut io::stdout().lock());
+    let write_response = |msg: msg::Response| {
+        use std::io::Write;
+
+        let out = &mut io::stdout().lock();
+        let prefix = current_response_prefix.borrow();
+        if !prefix.is_empty() {
+            out.write_all(prefix.as_bytes())?
+        }
+        msg.write(out)
+    };
 
     let mut srv = proc_macro_srv::ProcMacroSrv::default();
     let mut buf = String::new();
@@ -43,6 +55,10 @@ fn run() -> io::Result<()> {
             msg::Request::ApiVersionCheck {} => {
                 msg::Response::ApiVersionCheck(proc_macro_api::msg::CURRENT_API_VERSION)
             }
+            msg::Request::SetExpanderSettings { response_prefix } => {
+                *current_response_prefix.borrow_mut() = response_prefix;
+                msg::Response::SetExpanderSettings {}
+            }
         };
         write_response(res)?
     }
diff --git a/crates/proc-macro-srv/src/lib.rs b/crates/proc-macro-srv/src/lib.rs
index 5b03813b19a8..24af7b42dd84 100644
--- a/crates/proc-macro-srv/src/lib.rs
+++ b/crates/proc-macro-srv/src/lib.rs
@@ -14,6 +14,7 @@
 #![feature(proc_macro_internals, proc_macro_diagnostic, proc_macro_span)]
 #![warn(rust_2018_idioms, unused_lifetimes)]
 #![allow(unreachable_pub)]
+#![allow(internal_features)]
 
 extern crate proc_macro;