feat: Compute native diagnostics in parallel

3 files changed, 73 insertions, 26 deletions

diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/diagnostics.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/diagnostics.rs
index 6798e058dbf..434c7620d3d 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/diagnostics.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/diagnostics.rs
@@ -8,6 +8,7 @@ use ide_db::FxHashMap;
 use itertools::Itertools;
 use nohash_hasher::{IntMap, IntSet};
 use rustc_hash::FxHashSet;
+use stdx::iter_eq_by;
 use triomphe::Arc;
 
 use crate::{global_state::GlobalStateSnapshot, lsp, lsp_ext};
@@ -22,14 +23,21 @@ pub struct DiagnosticsMapConfig {
     pub check_ignore: FxHashSet<String>,
 }
 
+pub(crate) type DiagnosticsGeneration = usize;
+
 #[derive(Debug, Default, Clone)]
 pub(crate) struct DiagnosticCollection {
     // FIXME: should be IntMap<FileId, Vec<ra_id::Diagnostic>>
-    pub(crate) native: IntMap<FileId, Vec<lsp_types::Diagnostic>>,
+    pub(crate) native: IntMap<FileId, (DiagnosticsGeneration, Vec<lsp_types::Diagnostic>)>,
     // FIXME: should be Vec<flycheck::Diagnostic>
     pub(crate) check: IntMap<usize, IntMap<FileId, Vec<lsp_types::Diagnostic>>>,
     pub(crate) check_fixes: CheckFixes,
     changes: IntSet<FileId>,
+    /// Counter for supplying a new generation number for diagnostics.
+    /// This is used to keep track of when to clear the diagnostics for a given file as we compute
+    /// diagnostics on multiple worker threads simultaneously which may result in multiple diagnostics
+    /// updates for the same file in a single generation update (due to macros affecting multiple files).
+    generation: DiagnosticsGeneration,
 }
 
 #[derive(Debug, Clone)]
@@ -82,21 +90,31 @@ impl DiagnosticCollection {
 
     pub(crate) fn set_native_diagnostics(
         &mut self,
+        generation: DiagnosticsGeneration,
         file_id: FileId,
-        diagnostics: Vec<lsp_types::Diagnostic>,
+        mut diagnostics: Vec<lsp_types::Diagnostic>,
     ) {
-        if let Some(existing_diagnostics) = self.native.get(&file_id) {
+        diagnostics.sort_by_key(|it| (it.range.start, it.range.end));
+        if let Some((old_gen, existing_diagnostics)) = self.native.get_mut(&file_id) {
             if existing_diagnostics.len() == diagnostics.len()
-                && diagnostics
-                    .iter()
-                    .zip(existing_diagnostics)
-                    .all(|(new, existing)| are_diagnostics_equal(new, existing))
+                && iter_eq_by(&diagnostics, &*existing_diagnostics, |new, existing| {
+                    are_diagnostics_equal(new, existing)
+                })
             {
+                // don't signal an update if the diagnostics are the same
                 return;
             }
+            if *old_gen < generation || generation == 0 {
+                self.native.insert(file_id, (generation, diagnostics));
+            } else {
+                existing_diagnostics.extend(diagnostics);
+                // FIXME: Doing the merge step of a merge sort here would be a bit more performant
+                // but eh
+                existing_diagnostics.sort_by_key(|it| (it.range.start, it.range.end))
+            }
+        } else {
+            self.native.insert(file_id, (generation, diagnostics));
         }
-
-        self.native.insert(file_id, diagnostics);
         self.changes.insert(file_id);
     }
 
@@ -104,7 +122,7 @@ impl DiagnosticCollection {
         &self,
         file_id: FileId,
     ) -> impl Iterator<Item = &lsp_types::Diagnostic> {
-        let native = self.native.get(&file_id).into_iter().flatten();
+        let native = self.native.get(&file_id).into_iter().map(|(_, d)| d).flatten();
         let check = self.check.values().filter_map(move |it| it.get(&file_id)).flatten();
         native.chain(check)
     }
@@ -115,6 +133,11 @@ impl DiagnosticCollection {
         }
         Some(mem::take(&mut self.changes))
     }
+
+    pub(crate) fn next_generation(&mut self) -> usize {
+        self.generation += 1;
+        self.generation
+    }
 }
 
 fn are_diagnostics_equal(left: &lsp_types::Diagnostic, right: &lsp_types::Diagnostic) -> bool {
@@ -126,7 +149,8 @@ fn are_diagnostics_equal(left: &lsp_types::Diagnostic, right: &lsp_types::Diagno
 
 pub(crate) fn fetch_native_diagnostics(
     snapshot: GlobalStateSnapshot,
-    subscriptions: Vec<FileId>,
+    subscriptions: std::sync::Arc<[FileId]>,
+    slice: std::ops::Range<usize>,
 ) -> Vec<(FileId, Vec<lsp_types::Diagnostic>)> {
     let _p = tracing::info_span!("fetch_native_diagnostics").entered();
     let _ctx = stdx::panic_context::enter("fetch_native_diagnostics".to_owned());
@@ -149,7 +173,7 @@ pub(crate) fn fetch_native_diagnostics(
     // the diagnostics produced may point to different files not requested by the concrete request,
     // put those into here and filter later
     let mut odd_ones = Vec::new();
-    let mut diagnostics = subscriptions
+    let mut diagnostics = subscriptions[slice]
         .iter()
         .copied()
         .filter_map(|file_id| {
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/global_state.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/global_state.rs
index 59431d7d420..3d5f525aaf9 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/global_state.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/global_state.rs
@@ -163,7 +163,9 @@ pub(crate) struct GlobalStateSnapshot {
     pub(crate) semantic_tokens_cache: Arc<Mutex<FxHashMap<Url, SemanticTokens>>>,
     vfs: Arc<RwLock<(vfs::Vfs, IntMap<FileId, LineEndings>)>>,
     pub(crate) workspaces: Arc<Vec<ProjectWorkspace>>,
-    // used to signal semantic highlighting to fall back to syntax based highlighting until proc-macros have been loaded
+    // used to signal semantic highlighting to fall back to syntax based highlighting until
+    // proc-macros have been loaded
+    // FIXME: Can we derive this from somewhere else?
     pub(crate) proc_macros_loaded: bool,
     pub(crate) flycheck: Arc<[FlycheckHandle]>,
 }
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/main_loop.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/main_loop.rs
index 9b19e58eaa6..e2909be41a2 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/main_loop.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/main_loop.rs
@@ -17,7 +17,7 @@ use vfs::FileId;
 
 use crate::{
     config::Config,
-    diagnostics::fetch_native_diagnostics,
+    diagnostics::{fetch_native_diagnostics, DiagnosticsGeneration},
     dispatch::{NotificationDispatcher, RequestDispatcher},
     global_state::{file_id_to_url, url_to_file_id, GlobalState},
     hack_recover_crate_name,
@@ -87,7 +87,7 @@ pub(crate) enum Task {
     Response(lsp_server::Response),
     ClientNotification(lsp_ext::UnindexedProjectParams),
     Retry(lsp_server::Request),
-    Diagnostics(Vec<(FileId, Vec<lsp_types::Diagnostic>)>),
+    Diagnostics(DiagnosticsGeneration, Vec<(FileId, Vec<lsp_types::Diagnostic>)>),
     DiscoverTest(lsp_ext::DiscoverTestResults),
     PrimeCaches(PrimeCachesProgress),
     FetchWorkspace(ProjectWorkspaceProgress),
@@ -479,7 +479,7 @@ impl GlobalState {
 
     fn update_diagnostics(&mut self) {
         let db = self.analysis_host.raw_database();
-        // spawn a task per subscription?
+        let generation = self.diagnostics.next_generation();
         let subscriptions = {
             let vfs = &self.vfs.read().0;
             self.mem_docs
@@ -494,16 +494,37 @@ impl GlobalState {
                     // forever if we emitted them here.
                     !db.source_root(source_root).is_library
                 })
-                .collect::<Vec<_>>()
+                .collect::<std::sync::Arc<_>>()
         };
         tracing::trace!("updating notifications for {:?}", subscriptions);
-
-        // Diagnostics are triggered by the user typing
-        // so we run them on a latency sensitive thread.
-        self.task_pool.handle.spawn(ThreadIntent::LatencySensitive, {
-            let snapshot = self.snapshot();
-            move || Task::Diagnostics(fetch_native_diagnostics(snapshot, subscriptions))
-        });
+        // Split up the work on multiple threads, but we don't wanna fill the entire task pool with
+        // diagnostic tasks, so we limit the number of tasks to a quarter of the total thread pool.
+        let max_tasks = self.config.main_loop_num_threads() / 4;
+        let chunk_length = subscriptions.len() / max_tasks;
+        let remainder = subscriptions.len() % max_tasks;
+
+        let mut start = 0;
+        for task_idx in 0..max_tasks {
+            let extra = if task_idx < remainder { 1 } else { 0 };
+            let end = start + chunk_length + extra;
+            let slice = start..end;
+            if slice.is_empty() {
+                break;
+            }
+            // Diagnostics are triggered by the user typing
+            // so we run them on a latency sensitive thread.
+            self.task_pool.handle.spawn(ThreadIntent::LatencySensitive, {
+                let snapshot = self.snapshot();
+                let subscriptions = subscriptions.clone();
+                move || {
+                    Task::Diagnostics(
+                        generation,
+                        fetch_native_diagnostics(snapshot, subscriptions, slice),
+                    )
+                }
+            });
+            start = end;
+        }
     }
 
     fn update_tests(&mut self) {
@@ -590,9 +611,9 @@ impl GlobalState {
             // Only retry requests that haven't been cancelled. Otherwise we do unnecessary work.
             Task::Retry(req) if !self.is_completed(&req) => self.on_request(req),
             Task::Retry(_) => (),
-            Task::Diagnostics(diagnostics_per_file) => {
+            Task::Diagnostics(generation, diagnostics_per_file) => {
                 for (file_id, diagnostics) in diagnostics_per_file {
-                    self.diagnostics.set_native_diagnostics(file_id, diagnostics)
+                    self.diagnostics.set_native_diagnostics(generation, file_id, diagnostics)
                 }
             }
             Task::PrimeCaches(progress) => match progress {