rust-analyzer

diff --git a/crates/proc-macro-api/src/bidirectional_protocol.rs b/crates/proc-macro-api/src/bidirectional_protocol.rs
index 137f2dafc0..cd1f6f6f1f 100644
--- a/crates/proc-macro-api/src/bidirectional_protocol.rs
+++ b/crates/proc-macro-api/src/bidirectional_protocol.rs

@@ -138,6 +138,7 @@ pub(crate) fn find_proc_macros(

pub(crate) fn expand(

proc_macro: &ProcMacro,

+ process: &ProcMacroServerProcess,

subtree: tt::SubtreeView<'_>,

attr: Option<tt::SubtreeView<'_>>,

env: Vec<(String, String)>,

@@ -147,7 +148,7 @@ pub(crate) fn expand(

current_dir: String,

callback: SubCallback<'_>,

) -> Result<Result<tt::TopSubtree, String>, crate::ServerError> {

- let version = proc_macro.process.version();

+ let version = 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;

@@ -164,7 +165,7 @@ pub(crate) fn expand(

call_site,

mixed_site,

- span_data_table: if proc_macro.process.rust_analyzer_spans() {

+ span_data_table: if process.rust_analyzer_spans() {

serialize_span_data_index_map(&span_data_table)

} else {

Vec::new()

@@ -175,13 +176,13 @@ pub(crate) fn expand(

current_dir: Some(current_dir),

})));

- let response_payload = run_request(proc_macro.process.as_ref(), task, callback)?;

+ let response_payload = run_request(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() {

+ if proc_macro.needs_fixup_change(process) {

proc_macro.change_fixup_to_match_old_server(&mut expanded);

}

expanded

@@ -194,7 +195,7 @@ pub(crate) fn expand(

version,

&deserialize_span_data_index_map(&resp.span_data_table),

);

- if proc_macro.needs_fixup_change() {

+ if proc_macro.needs_fixup_change(process) {

proc_macro.change_fixup_to_match_old_server(&mut expanded);

}

expanded

diff --git a/crates/proc-macro-api/src/legacy_protocol.rs b/crates/proc-macro-api/src/legacy_protocol.rs
index 7b546cf7ae..412d207303 100644
--- a/crates/proc-macro-api/src/legacy_protocol.rs
+++ b/crates/proc-macro-api/src/legacy_protocol.rs

@@ -77,6 +77,7 @@ pub(crate) fn find_proc_macros(

pub(crate) fn expand(

proc_macro: &ProcMacro,

+ process: &ProcMacroServerProcess,

subtree: tt::SubtreeView<'_>,

attr: Option<tt::SubtreeView<'_>>,

env: Vec<(String, String)>,

@@ -85,7 +86,7 @@ pub(crate) fn expand(

mixed_site: Span,

current_dir: String,

) -> Result<Result<tt::TopSubtree, String>, crate::ServerError> {

- let version = proc_macro.process.version();

+ let version = 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;

@@ -102,7 +103,7 @@ pub(crate) fn expand(

call_site,

mixed_site,

- span_data_table: if proc_macro.process.rust_analyzer_spans() {

+ span_data_table: if process.rust_analyzer_spans() {

serialize_span_data_index_map(&span_data_table)

} else {

Vec::new()

@@ -113,13 +114,13 @@ pub(crate) fn expand(

current_dir: Some(current_dir),

};

- let response = send_task(proc_macro.process.as_ref(), Request::ExpandMacro(Box::new(task)))?;

+ let response = send_task(process, Request::ExpandMacro(Box::new(task)))?;

match 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() {

+ if proc_macro.needs_fixup_change(process) {

proc_macro.change_fixup_to_match_old_server(&mut expanded);

}

expanded

@@ -132,7 +133,7 @@ pub(crate) fn expand(

version,

&deserialize_span_data_index_map(&resp.span_data_table),

);

- if proc_macro.needs_fixup_change() {

+ if proc_macro.needs_fixup_change(process) {

proc_macro.change_fixup_to_match_old_server(&mut expanded);

}

expanded

diff --git a/crates/proc-macro-api/src/lib.rs b/crates/proc-macro-api/src/lib.rs
index ffae28f92c..fe17e14024 100644
--- a/crates/proc-macro-api/src/lib.rs
+++ b/crates/proc-macro-api/src/lib.rs

@@ -112,7 +112,7 @@ impl MacroDylib {

/// we share a single expander process for all macros within a workspace.

#[derive(Debug, Clone)]

pub struct ProcMacro {

- process: Arc<ProcMacroServerProcess>,

+ process: ProcMacroServerPool,

dylib_path: Arc<AbsPathBuf>,

name: Box<str>,

kind: ProcMacroKind,

@@ -126,7 +126,6 @@ impl PartialEq for ProcMacro {

&& self.kind == other.kind

&& self.dylib_path == other.dylib_path

&& self.dylib_last_modified == other.dylib_last_modified

- && Arc::ptr_eq(&self.process, &other.process)

}

@@ -210,8 +209,8 @@ impl ProcMacro {

self.kind

}

- fn needs_fixup_change(&self) -> bool {

- let version = self.process.version();

+ fn needs_fixup_change(&self, process: &ProcMacroServerProcess) -> bool {

+ let version = process.version();

(version::RUST_ANALYZER_SPAN_SUPPORT..version::HASHED_AST_ID).contains(&version)

}

@@ -241,20 +240,6 @@ impl ProcMacro {

current_dir: String,

callback: Option<SubCallback<'_>>,

) -> Result<Result<tt::TopSubtree, String>, ServerError> {

- let (mut subtree, mut attr) = (subtree, attr);

- let (mut subtree_changed, mut attr_changed);

- if self.needs_fixup_change() {

- subtree_changed = tt::TopSubtree::from_subtree(subtree);

- self.change_fixup_to_match_old_server(&mut subtree_changed);

- subtree = subtree_changed.view();

- if let Some(attr) = &mut attr {

- attr_changed = tt::TopSubtree::from_subtree(*attr);

- self.change_fixup_to_match_old_server(&mut attr_changed);

- *attr = attr_changed.view();

- }

self.process.expand(

self,

subtree,

diff --git a/crates/proc-macro-api/src/pool.rs b/crates/proc-macro-api/src/pool.rs
index 685bc05be6..4639374f3e 100644
--- a/crates/proc-macro-api/src/pool.rs
+++ b/crates/proc-macro-api/src/pool.rs

@@ -1,29 +1,38 @@

use std::sync::Arc;

+use tt::Span;

use crate::{

MacroDylib, ProcMacro, ServerError, bidirectional_protocol::SubCallback,

process::ProcMacroServerProcess,

};

-#[derive(Debug)]

+#[derive(Debug, Clone)]

pub(crate) struct ProcMacroServerPool {

- workers: Vec<Arc<ProcMacroServerProcess>>,

+ workers: Arc<[ProcMacroServerProcess]>,

}

impl ProcMacroServerPool {

- pub(crate) fn new(workers: Vec<Arc<ProcMacroServerProcess>>) -> Self {

- Self { workers }

+ pub(crate) fn new(workers: Vec<ProcMacroServerProcess>) -> Self {

+ Self { workers: workers.into() }

}

impl ProcMacroServerPool {

pub(crate) fn exited(&self) -> Option<&ServerError> {

- for worker in &self.workers {

- if let Some(e) = worker.exited() {

- return Some(e);

+ for worker in &*self.workers {

+ worker.exited()?;

+ }

+ self.workers[0].exited()

+ }

+ fn pick_process(&self) -> &ProcMacroServerProcess {

+ for workers in &*self.workers {

+ if workers.can_use() {

+ return workers;

}

- None

+ &self.workers[0]

}

pub(crate) fn load_dylib(

@@ -34,16 +43,16 @@ impl ProcMacroServerPool {

let _p = tracing::info_span!("ProcMacroServer::load_dylib").entered();

let mut all_macros = Vec::new();

- for worker in &self.workers {

+ for worker in &*self.workers {

let dylib_path = Arc::new(dylib.path.clone());

let dylib_last_modified = std::fs::metadata(dylib_path.as_path())

.ok()

.and_then(|metadata| metadata.modified().ok());

- let macros = worker.load_dylib(&dylib.path, None)?;

+ let macros = worker.find_proc_macros(&dylib.path, None)?.unwrap();

for (name, kind) in macros {

all_macros.push(ProcMacro {

- process: worker.clone(),

+ process: self.clone(),

name: name.into(),

kind,

dylib_path: Arc::new(dylib.path.clone()),

@@ -54,6 +63,47 @@ impl ProcMacroServerPool {

Ok(all_macros)

}

+ pub(crate) fn expand(

+ &self,

+ proc_macro: &ProcMacro,

+ subtree: tt::SubtreeView<'_>,

+ attr: Option<tt::SubtreeView<'_>>,

+ env: Vec<(String, String)>,

+ def_site: Span,

+ call_site: Span,

+ mixed_site: Span,

+ current_dir: String,

+ callback: Option<SubCallback<'_>>,

+ ) -> Result<Result<tt::TopSubtree, String>, ServerError> {

+ let process = self.pick_process();

+ let (mut subtree, mut attr) = (subtree, attr);

+ let (mut subtree_changed, mut attr_changed);

+ if proc_macro.needs_fixup_change(process) {

+ subtree_changed = tt::TopSubtree::from_subtree(subtree);

+ proc_macro.change_fixup_to_match_old_server(&mut subtree_changed);

+ subtree = subtree_changed.view();

+ if let Some(attr) = &mut attr {

+ attr_changed = tt::TopSubtree::from_subtree(*attr);

+ proc_macro.change_fixup_to_match_old_server(&mut attr_changed);

+ *attr = attr_changed.view();

+ }

+ process.expand(

+ proc_macro,

+ subtree,

+ attr,

+ env,

+ def_site,

+ call_site,

+ mixed_site,

+ current_dir,

+ callback,

+ )

+ }

}

pub(crate) fn default_pool_size() -> usize {

diff --git a/crates/proc-macro-api/src/process.rs b/crates/proc-macro-api/src/process.rs
index 30877c5cf4..a41bb58e74 100644
--- a/crates/proc-macro-api/src/process.rs
+++ b/crates/proc-macro-api/src/process.rs

@@ -5,7 +5,10 @@ use std::{

io::{self, BufRead, BufReader, Read, Write},

panic::AssertUnwindSafe,

process::{Child, ChildStdin, ChildStdout, Command, Stdio},

- sync::{Arc, Mutex, OnceLock},

+ sync::{

+ Arc, Mutex, OnceLock,

+ atomic::{AtomicBool, Ordering},

+ },

};

use paths::AbsPath;

@@ -29,7 +32,7 @@ pub(crate) struct ProcMacroServerProcess {

protocol: Protocol,

/// Populated when the server exits.

exited: OnceLock<AssertUnwindSafe<ServerError>>,

- single_use: bool,

+ can_use: AtomicBool,

}

impl std::fmt::Debug for ProcMacroServerProcess {

@@ -149,11 +152,7 @@ impl ProcMacroServerProcess {

let (process, stdin, stdout) = spawn(format)?;

io::Result::Ok(ProcMacroServerProcess {

- state: Mutex::new(ProcessSrvState {

- process,

- stdin,

- stdout,

- }),

+ state: Mutex::new(ProcessSrvState { process, stdin, stdout }),

version: 0,

protocol: match format {

Some(ProtocolFormat::BidirectionalPostcardPrototype) => {

@@ -164,7 +163,7 @@ impl ProcMacroServerProcess {

}

exited: OnceLock::new(),

- single_use,

+ can_use: AtomicBool::new(true),

})

};

let mut srv = create_srv()?;

@@ -204,34 +203,20 @@ impl ProcMacroServerProcess {

Err(err.unwrap())

}

- pub(crate) fn load_dylib(

+ /// Finds proc-macros in a given dynamic library.

+ pub(crate) fn find_proc_macros(

&self,

dylib_path: &AbsPath,

callback: Option<SubCallback<'_>>,

- ) -> Result<Vec<(String, ProcMacroKind)>, ServerError> {

- let _state = self.state.lock().unwrap();

- // if state.loaded_dylibs.contains(dylib_path) {

- // // Already loaded in this worker

- // return Ok(Vec::new());

- // }

- let result = match self.protocol {

+ ) -> Result<Result<Vec<(String, ProcMacroKind)>, String>, ServerError> {

+ match self.protocol {

Protocol::LegacyJson { .. } | Protocol::LegacyPostcard { .. } => {

- legacy_protocol::find_proc_macros(self, dylib_path)?

+ legacy_protocol::find_proc_macros(self, dylib_path)

}

Protocol::BidirectionalPostcardPrototype { .. } => {

- let cb = callback.expect("callback required");

- bidirectional_protocol::find_proc_macros(self, dylib_path, cb)?

- }

- };

- match result {

- Ok(macros) => {

- // state.loaded_dylibs.insert(dylib_path.to_owned());

- Ok(macros)

+ let cb = callback.expect("callback required for bidirectional protocol");

+ bidirectional_protocol::find_proc_macros(self, dylib_path, cb)

}

- Err(message) => Err(ServerError { message, io: None }),

}

@@ -290,9 +275,11 @@ impl ProcMacroServerProcess {

current_dir: String,

callback: Option<SubCallback<'_>>,

) -> Result<Result<tt::TopSubtree, String>, ServerError> {

+ self.can_use.store(false, Ordering::Release);

let result = match self.protocol {

Protocol::LegacyJson { .. } => legacy_protocol::expand(

proc_macro,

+ self,

subtree,

attr,

env,

@@ -303,6 +290,7 @@ impl ProcMacroServerProcess {

Protocol::BidirectionalPostcardPrototype { .. } => bidirectional_protocol::expand(

proc_macro,

+ self,

subtree,

attr,

env,

@@ -314,23 +302,10 @@ impl ProcMacroServerProcess {

};

- if !self.is_reusable() {

- self.terminate();

- }

+ self.can_use.store(true, Ordering::Release);

result

}

- fn is_reusable(&self) -> bool {

- self.single_use

- }

- fn terminate(&self) {

- if let Ok(mut state) = self.state.lock() {

- let _ = state.process.child.kill();

- }

pub(crate) fn send_task<Request, Response, C: Codec>(

&self,

send: impl FnOnce(

@@ -391,6 +366,10 @@ impl ProcMacroServerProcess {

bidirectional_protocol::run_conversation::<C>(writer, reader, buf, initial, callback)

})

}

+ pub(crate) fn can_use(&self) -> bool {

+ self.can_use.load(Ordering::Acquire)

+ }

}

/// Manages the execution of the proc-macro server process.