internal: Use rayon for proc-macro loading

Author: Veykril

Cargo.lock

@@ -26,10 +26,7 @@ use ide_db::{
 use itertools::Itertools;
 use proc_macro_api::{
     MacroDylib, ProcMacroClient,
-    bidirectional_protocol::{
-        msg::{SubRequest, SubResponse},
-        reject_subrequests,
-    },
+    bidirectional_protocol::msg::{SubRequest, SubResponse},
 };
 use project_model::{CargoConfig, PackageRoot, ProjectManifest, ProjectWorkspace};
 use span::{Span, SpanAnchor, SyntaxContext};
@@ -446,7 +443,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, Some(&reject_subrequests)).map_err(|e| {
+        let vec = server.load_dylib(dylib).map_err(|e| {
             ProcMacroLoadingError::ProcMacroSrvError(format!("{e}").into_boxed_str())
         })?;
         if vec.is_empty() {

crates/load-cargo/src/lib.rs

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

crates/proc-macro-api/Cargo.toml

@@ -198,12 +198,8 @@ impl ProcMacroClient {
     }
 
     /// Loads a proc-macro dylib into the server process returning a list of `ProcMacro`s loaded.
-    pub fn load_dylib(
-        &self,
-        dylib: MacroDylib,
-        callback: Option<SubCallback<'_>>,
-    ) -> Result<Vec<ProcMacro>, ServerError> {
-        self.pool.load_dylib(&dylib, callback)
+    pub fn load_dylib(&self, dylib: MacroDylib) -> Result<Vec<ProcMacro>, ServerError> {
+        self.pool.load_dylib(&dylib)
     }
 
     /// Checks if the proc-macro server has exited.

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

@@ -1,10 +1,9 @@
 //! A pool of proc-macro server processes
 use std::sync::Arc;
 
-use crate::{
-    MacroDylib, ProcMacro, ServerError, bidirectional_protocol::SubCallback,
-    process::ProcMacroServerProcess,
-};
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+
+use crate::{MacroDylib, ProcMacro, ServerError, process::ProcMacroServerProcess};
 
 #[derive(Debug, Clone)]
 pub(crate) struct ProcMacroServerPool {
@@ -50,11 +49,7 @@ impl ProcMacroServerPool {
         })
     }
 
-    pub(crate) fn load_dylib(
-        &self,
-        dylib: &MacroDylib,
-        callback: Option<SubCallback<'_>>,
-    ) -> Result<Vec<ProcMacro>, ServerError> {
+    pub(crate) fn load_dylib(&self, dylib: &MacroDylib) -> Result<Vec<ProcMacro>, ServerError> {
         let _span = tracing::info_span!("ProcMacroServer::load_dylib").entered();
 
         let dylib_path = Arc::new(dylib.path.clone());
@@ -64,14 +59,17 @@ impl ProcMacroServerPool {
         let (first, rest) = self.workers.split_first().expect("worker pool must not be empty");
 
         let macros = first
-            .find_proc_macros(&dylib.path, callback)?
+            .find_proc_macros(&dylib.path)?
             .map_err(|e| ServerError { message: e, io: None })?;
 
-        for worker in rest {
-            worker
-                .find_proc_macros(&dylib.path, callback)?
-                .map_err(|e| ServerError { message: e, io: None })?;
-        }
+        rest.into_par_iter()
+            .map(|worker| {
+                worker
+                    .find_proc_macros(&dylib.path)?
+                    .map(|_| ())
+                    .map_err(|e| ServerError { message: e, io: None })
+            })
+            .collect::<Result<(), _>>()?;
 
         Ok(macros
             .into_iter()

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

@@ -18,7 +18,11 @@ use stdx::JodChild;
 
 use crate::{
     ProcMacro, ProcMacroKind, ProtocolFormat, ServerError,
-    bidirectional_protocol::{self, SubCallback, msg::BidirectionalMessage, reject_subrequests},
+    bidirectional_protocol::{
+        self, SubCallback,
+        msg::{BidirectionalMessage, SubResponse},
+        reject_subrequests,
+    },
     legacy_protocol::{self, SpanMode},
     version,
 };
@@ -207,14 +211,18 @@ impl ProcMacroServerProcess {
     pub(crate) fn find_proc_macros(
         &self,
         dylib_path: &AbsPath,
-        callback: Option<SubCallback<'_>>,
     ) -> Result<Result<Vec<(String, ProcMacroKind)>, String>, ServerError> {
         match self.protocol {
             Protocol::LegacyJson { .. } => legacy_protocol::find_proc_macros(self, dylib_path),
 
             Protocol::BidirectionalPostcardPrototype { .. } => {
-                let cb = callback.expect("callback required for bidirectional protocol");
-                bidirectional_protocol::find_proc_macros(self, dylib_path, cb)
+                bidirectional_protocol::find_proc_macros(self, dylib_path, &|_| {
+                    Ok(SubResponse::Cancel {
+                        reason: String::from(
+                            "Server should not do a sub request when loading proc-macros",
+                        ),
+                    })
+                })
             }
         }
     }

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

@@ -88,7 +88,6 @@ impl fmt::Debug for ErasedFileAstId {
             Module,
             Static,
             Trait,
-            TraitAlias,
             Variant,
             Const,
             Fn,
@@ -129,7 +128,6 @@ enum ErasedFileAstIdKind {
     Module,
     Static,
     Trait,
-    TraitAlias,
     // Until here associated with `ErasedHasNameFileAstId`.
     // The following are associated with `ErasedAssocItemFileAstId`.
     Variant,