Add bidirectional messaging proc-macro-srv

.github/workflows/ci.yaml

@@ -70,7 +70,7 @@ jobs:
 
       - name: Test
         run: cargo test --features sysroot-abi -p proc-macro-srv -p proc-macro-srv-cli -p proc-macro-api -- --quiet
-
+      
       - name: Check salsa dependency
         run: "! (cargo tree -p proc-macro-srv-cli | grep -q salsa)"
 

crates/hir-def/src/macro_expansion_tests/mod.rs

@@ -16,7 +16,7 @@ mod proc_macros;
 
 use std::{any::TypeId, iter, ops::Range, sync};
 
-use base_db::RootQueryDb;
+use base_db::{RootQueryDb, SourceDatabase};
 use expect_test::Expect;
 use hir_expand::{
     AstId, InFile, MacroCallId, MacroCallKind, MacroKind,
@@ -374,6 +374,7 @@ struct IdentityWhenValidProcMacroExpander;
 impl ProcMacroExpander for IdentityWhenValidProcMacroExpander {
     fn expand(
         &self,
+        _: &dyn SourceDatabase,
         subtree: &TopSubtree,
         _: Option<&TopSubtree>,
         _: &base_db::Env,

crates/hir-expand/src/proc_macro.rs

@@ -4,7 +4,7 @@ use core::fmt;
 use std::any::Any;
 use std::{panic::RefUnwindSafe, sync};
 
-use base_db::{Crate, CrateBuilderId, CratesIdMap, Env, ProcMacroLoadingError};
+use base_db::{Crate, CrateBuilderId, CratesIdMap, Env, ProcMacroLoadingError, SourceDatabase};
 use intern::Symbol;
 use rustc_hash::FxHashMap;
 use span::Span;
@@ -25,6 +25,7 @@ pub trait ProcMacroExpander: fmt::Debug + Send + Sync + RefUnwindSafe + Any {
     /// [`ProcMacroKind::Attr`]), environment variables, and span information.
     fn expand(
         &self,
+        db: &dyn SourceDatabase,
         subtree: &tt::TopSubtree,
         attrs: Option<&tt::TopSubtree>,
         env: &Env,
@@ -309,6 +310,7 @@ impl CustomProcMacroExpander {
                 let current_dir = calling_crate.data(db).proc_macro_cwd.to_string();
 
                 match proc_macro.expander.expand(
+                    db,
                     tt,
                     attr_arg,
                     env,

crates/load-cargo/src/lib.rs

@@ -17,15 +17,23 @@ use hir_expand::proc_macro::{
 };
 use ide_db::{
     ChangeWithProcMacros, FxHashMap, RootDatabase,
-    base_db::{CrateGraphBuilder, Env, ProcMacroLoadingError, SourceRoot, SourceRootId},
+    base_db::{
+        CrateGraphBuilder, Env, ProcMacroLoadingError, SourceDatabase, SourceRoot, SourceRootId,
+    },
     prime_caches,
 };
 use itertools::Itertools;
-use proc_macro_api::{MacroDylib, ProcMacroClient};
+use proc_macro_api::{
+    MacroDylib, ProcMacroClient,
+    bidirectional_protocol::{
+        msg::{SubRequest, SubResponse},
+        reject_subrequests,
+    },
+};
 use project_model::{CargoConfig, PackageRoot, ProjectManifest, ProjectWorkspace};
 use span::Span;
 use vfs::{
-    AbsPath, AbsPathBuf, VfsPath,
+    AbsPath, AbsPathBuf, FileId, VfsPath,
     file_set::FileSetConfig,
     loader::{Handle, LoadingProgress},
 };
@@ -425,7 +433,7 @@ pub fn load_proc_macro(
 ) -> ProcMacroLoadResult {
     let res: Result<Vec<_>, _> = (|| {
         let dylib = MacroDylib::new(path.to_path_buf());
-        let vec = server.load_dylib(dylib).map_err(|e| {
+        let vec = server.load_dylib(dylib, Some(&mut reject_subrequests)).map_err(|e| {
             ProcMacroLoadingError::ProcMacroSrvError(format!("{e}").into_boxed_str())
         })?;
         if vec.is_empty() {
@@ -522,6 +530,7 @@ struct Expander(proc_macro_api::ProcMacro);
 impl ProcMacroExpander for Expander {
     fn expand(
         &self,
+        db: &dyn SourceDatabase,
         subtree: &tt::TopSubtree<Span>,
         attrs: Option<&tt::TopSubtree<Span>>,
         env: &Env,
@@ -530,6 +539,14 @@ impl ProcMacroExpander for Expander {
         mixed_site: Span,
         current_dir: String,
     ) -> Result<tt::TopSubtree<Span>, ProcMacroExpansionError> {
+        let mut cb = |req| match req {
+            SubRequest::SourceText { file_id, start, end } => {
+                let file = FileId::from_raw(file_id);
+                let text = db.file_text(file).text(db);
+                let slice = text.get(start as usize..end as usize).map(ToOwned::to_owned);
+                Ok(SubResponse::SourceTextResult { text: slice })
+            }
+        };
         match self.0.expand(
             subtree.view(),
             attrs.map(|attrs| attrs.view()),
@@ -538,6 +555,7 @@ impl ProcMacroExpander for Expander {
             call_site,
             mixed_site,
             current_dir,
+            Some(&mut cb),
         ) {
             Ok(Ok(subtree)) => Ok(subtree),
             Ok(Err(err)) => Err(ProcMacroExpansionError::Panic(err)),

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

@@ -0,0 +1,226 @@
+//! Bidirectional protocol methods
+
+use std::{
+    io::{self, BufRead, Write},
+    sync::Arc,
+};
+
+use paths::AbsPath;
+use span::Span;
+
+use crate::{
+    Codec, ProcMacro, ProcMacroKind, ServerError,
+    bidirectional_protocol::msg::{
+        BidirectionalMessage, ExpandMacro, ExpandMacroData, ExpnGlobals, Request, Response,
+        SubRequest, SubResponse,
+    },
+    legacy_protocol::{
+        SpanMode,
+        msg::{
+            FlatTree, ServerConfig, SpanDataIndexMap, deserialize_span_data_index_map,
+            serialize_span_data_index_map,
+        },
+    },
+    process::ProcMacroServerProcess,
+    transport::codec::{json::JsonProtocol, postcard::PostcardProtocol},
+    version,
+};
+
+pub mod msg;
+
+pub type SubCallback<'a> = &'a mut dyn FnMut(SubRequest) -> Result<SubResponse, ServerError>;
+
+pub fn run_conversation<C: Codec>(
+    writer: &mut dyn Write,
+    reader: &mut dyn BufRead,
+    buf: &mut C::Buf,
+    msg: BidirectionalMessage,
+    callback: SubCallback<'_>,
+) -> Result<BidirectionalMessage, ServerError> {
+    let encoded = C::encode(&msg).map_err(wrap_encode)?;
+    C::write(writer, &encoded).map_err(wrap_io("failed to write initial request"))?;
+
+    loop {
+        let maybe_buf = C::read(reader, buf).map_err(wrap_io("failed to read message"))?;
+        let Some(b) = maybe_buf else {
+            return Err(ServerError {
+                message: "proc-macro server closed the stream".into(),
+                io: Some(Arc::new(io::Error::new(io::ErrorKind::UnexpectedEof, "closed"))),
+            });
+        };
+
+        let msg: BidirectionalMessage = C::decode(b).map_err(wrap_decode)?;
+
+        match msg {
+            BidirectionalMessage::Response(response) => {
+                return Ok(BidirectionalMessage::Response(response));
+            }
+            BidirectionalMessage::SubRequest(sr) => {
+                let resp = callback(sr)?;
+                let reply = BidirectionalMessage::SubResponse(resp);
+                let encoded = C::encode(&reply).map_err(wrap_encode)?;
+                C::write(writer, &encoded).map_err(wrap_io("failed to write sub-response"))?;
+            }
+            _ => {
+                return Err(ServerError {
+                    message: format!("unexpected message {:?}", msg),
+                    io: None,
+                });
+            }
+        }
+    }
+}
+
+fn wrap_io(msg: &'static str) -> impl Fn(io::Error) -> ServerError {
+    move |err| ServerError { message: msg.into(), io: Some(Arc::new(err)) }
+}
+
+fn wrap_encode(err: io::Error) -> ServerError {
+    ServerError { message: "failed to encode message".into(), io: Some(Arc::new(err)) }
+}
+
+fn wrap_decode(err: io::Error) -> ServerError {
+    ServerError { message: "failed to decode message".into(), io: Some(Arc::new(err)) }
+}
+
+pub(crate) fn version_check(
+    srv: &ProcMacroServerProcess,
+    callback: SubCallback<'_>,
+) -> Result<u32, ServerError> {
+    let request = BidirectionalMessage::Request(Request::ApiVersionCheck {});
+
+    let response_payload = run_request(srv, request, callback)?;
+
+    match response_payload {
+        BidirectionalMessage::Response(Response::ApiVersionCheck(version)) => Ok(version),
+        other => {
+            Err(ServerError { message: format!("unexpected response: {:?}", other), io: None })
+        }
+    }
+}
+
+/// Enable support for rust-analyzer span mode if the server supports it.
+pub(crate) fn enable_rust_analyzer_spans(
+    srv: &ProcMacroServerProcess,
+    callback: SubCallback<'_>,
+) -> Result<SpanMode, ServerError> {
+    let request = BidirectionalMessage::Request(Request::SetConfig(ServerConfig {
+        span_mode: SpanMode::RustAnalyzer,
+    }));
+
+    let response_payload = run_request(srv, request, callback)?;
+
+    match response_payload {
+        BidirectionalMessage::Response(Response::SetConfig(ServerConfig { span_mode })) => {
+            Ok(span_mode)
+        }
+        _ => Err(ServerError { message: "unexpected response".to_owned(), io: None }),
+    }
+}
+
+/// Finds proc-macros in a given dynamic library.
+pub(crate) fn find_proc_macros(
+    srv: &ProcMacroServerProcess,
+    dylib_path: &AbsPath,
+    callback: SubCallback<'_>,
+) -> Result<Result<Vec<(String, ProcMacroKind)>, String>, ServerError> {
+    let request = BidirectionalMessage::Request(Request::ListMacros {
+        dylib_path: dylib_path.to_path_buf().into(),
+    });
+
+    let response_payload = run_request(srv, request, callback)?;
+
+    match response_payload {
+        BidirectionalMessage::Response(Response::ListMacros(it)) => Ok(it),
+        _ => Err(ServerError { message: "unexpected response".to_owned(), io: None }),
+    }
+}
+
+pub(crate) fn expand(
+    proc_macro: &ProcMacro,
+    subtree: tt::SubtreeView<'_, Span>,
+    attr: Option<tt::SubtreeView<'_, Span>>,
+    env: Vec<(String, String)>,
+    def_site: Span,
+    call_site: Span,
+    mixed_site: Span,
+    current_dir: String,
+    callback: SubCallback<'_>,
+) -> Result<Result<tt::TopSubtree<span::SpanData<span::SyntaxContext>>, String>, crate::ServerError>
+{
+    let version = proc_macro.process.version();
+    let mut span_data_table = SpanDataIndexMap::default();
+    let def_site = span_data_table.insert_full(def_site).0;
+    let call_site = span_data_table.insert_full(call_site).0;
+    let mixed_site = span_data_table.insert_full(mixed_site).0;
+    let task = BidirectionalMessage::Request(Request::ExpandMacro(Box::new(ExpandMacro {
+        data: ExpandMacroData {
+            macro_body: FlatTree::from_subtree(subtree, version, &mut span_data_table),
+            macro_name: proc_macro.name.to_string(),
+            attributes: attr
+                .map(|subtree| FlatTree::from_subtree(subtree, version, &mut span_data_table)),
+            has_global_spans: ExpnGlobals {
+                serialize: version >= version::HAS_GLOBAL_SPANS,
+                def_site,
+                call_site,
+                mixed_site,
+            },
+            span_data_table: if proc_macro.process.rust_analyzer_spans() {
+                serialize_span_data_index_map(&span_data_table)
+            } else {
+                Vec::new()
+            },
+        },
+        lib: proc_macro.dylib_path.to_path_buf().into(),
+        env,
+        current_dir: Some(current_dir),
+    })));
+
+    let response_payload = run_request(&proc_macro.process, task, callback)?;
+
+    match response_payload {
+        BidirectionalMessage::Response(Response::ExpandMacro(it)) => Ok(it
+            .map(|tree| {
+                let mut expanded = FlatTree::to_subtree_resolved(tree, version, &span_data_table);
+                if proc_macro.needs_fixup_change() {
+                    proc_macro.change_fixup_to_match_old_server(&mut expanded);
+                }
+                expanded
+            })
+            .map_err(|msg| msg.0)),
+        BidirectionalMessage::Response(Response::ExpandMacroExtended(it)) => Ok(it
+            .map(|resp| {
+                let mut expanded = FlatTree::to_subtree_resolved(
+                    resp.tree,
+                    version,
+                    &deserialize_span_data_index_map(&resp.span_data_table),
+                );
+                if proc_macro.needs_fixup_change() {
+                    proc_macro.change_fixup_to_match_old_server(&mut expanded);
+                }
+                expanded
+            })
+            .map_err(|msg| msg.0)),
+        _ => Err(ServerError { message: "unexpected response".to_owned(), io: None }),
+    }
+}
+
+fn run_request(
+    srv: &ProcMacroServerProcess,
+    msg: BidirectionalMessage,
+    callback: SubCallback<'_>,
+) -> Result<BidirectionalMessage, ServerError> {
+    if let Some(server_error) = srv.exited() {
+        return Err(server_error.clone());
+    }
+
+    if srv.use_postcard() {
+        srv.run_bidirectional::<PostcardProtocol>(msg, callback)
+    } else {
+        srv.run_bidirectional::<JsonProtocol>(msg, callback)
+    }
+}
+
+pub fn reject_subrequests(req: SubRequest) -> Result<SubResponse, ServerError> {
+    Err(ServerError { message: format!("{req:?} sub-request not supported here"), io: None })
+}