Simplify diagnostics handling

3 files changed, 50 insertions, 70 deletions

diff --git a/crates/rust-analyzer/src/diagnostics.rs b/crates/rust-analyzer/src/diagnostics.rs
index 290609e7f..f3cdb842b 100644
--- a/crates/rust-analyzer/src/diagnostics.rs
+++ b/crates/rust-analyzer/src/diagnostics.rs
@@ -1,14 +1,15 @@
 //! Book keeping for keeping diagnostics easily in sync with the client.
 pub(crate) mod to_proto;
 
-use std::{collections::HashMap, sync::Arc};
+use std::{collections::HashMap, mem, sync::Arc};
 
 use lsp_types::{Diagnostic, Range};
 use ra_ide::FileId;
+use rustc_hash::FxHashSet;
 
 use crate::lsp_ext;
 
-pub type CheckFixes = Arc<HashMap<FileId, Vec<Fix>>>;
+pub(crate) type CheckFixes = Arc<HashMap<FileId, Vec<Fix>>>;
 
 #[derive(Debug, Default, Clone)]
 pub struct DiagnosticsConfig {
@@ -17,32 +18,26 @@ pub struct DiagnosticsConfig {
 }
 
 #[derive(Debug, Default, Clone)]
-pub struct DiagnosticCollection {
-    pub native: HashMap<FileId, Vec<Diagnostic>>,
-    pub check: HashMap<FileId, Vec<Diagnostic>>,
-    pub check_fixes: CheckFixes,
+pub(crate) struct DiagnosticCollection {
+    pub(crate) native: HashMap<FileId, Vec<Diagnostic>>,
+    pub(crate) check: HashMap<FileId, Vec<Diagnostic>>,
+    pub(crate) check_fixes: CheckFixes,
+    changes: FxHashSet<FileId>,
 }
 
 #[derive(Debug, Clone)]
-pub struct Fix {
-    pub range: Range,
-    pub action: lsp_ext::CodeAction,
-}
-
-#[derive(Debug)]
-pub enum DiagnosticTask {
-    ClearCheck,
-    AddCheck(FileId, Diagnostic, Vec<lsp_ext::CodeAction>),
-    SetNative(FileId, Vec<Diagnostic>),
+pub(crate) struct Fix {
+    pub(crate) range: Range,
+    pub(crate) action: lsp_ext::CodeAction,
 }
 
 impl DiagnosticCollection {
-    pub fn clear_check(&mut self) -> Vec<FileId> {
+    pub(crate) fn clear_check(&mut self) {
         Arc::make_mut(&mut self.check_fixes).clear();
-        self.check.drain().map(|(key, _value)| key).collect()
+        self.changes.extend(self.check.drain().map(|(key, _value)| key))
     }
 
-    pub fn add_check_diagnostic(
+    pub(crate) fn add_check_diagnostic(
         &mut self,
         file_id: FileId,
         diagnostic: Diagnostic,
@@ -61,30 +56,25 @@ impl DiagnosticCollection {
             .or_default()
             .extend(fixes.into_iter().map(|action| Fix { range: diagnostic.range, action }));
         diagnostics.push(diagnostic);
+        self.changes.insert(file_id);
     }
 
-    pub fn set_native_diagnostics(&mut self, file_id: FileId, diagnostics: Vec<Diagnostic>) {
+    pub(crate) fn set_native_diagnostics(&mut self, file_id: FileId, diagnostics: Vec<Diagnostic>) {
         self.native.insert(file_id, diagnostics);
+        self.changes.insert(file_id);
     }
 
-    pub fn diagnostics_for(&self, file_id: FileId) -> impl Iterator<Item = &Diagnostic> {
+    pub(crate) fn diagnostics_for(&self, file_id: FileId) -> impl Iterator<Item = &Diagnostic> {
         let native = self.native.get(&file_id).into_iter().flatten();
         let check = self.check.get(&file_id).into_iter().flatten();
         native.chain(check)
     }
 
-    pub fn handle_task(&mut self, task: DiagnosticTask) -> Vec<FileId> {
-        match task {
-            DiagnosticTask::ClearCheck => self.clear_check(),
-            DiagnosticTask::AddCheck(file_id, diagnostic, fixes) => {
-                self.add_check_diagnostic(file_id, diagnostic, fixes);
-                vec![file_id]
-            }
-            DiagnosticTask::SetNative(file_id, diagnostics) => {
-                self.set_native_diagnostics(file_id, diagnostics);
-                vec![file_id]
-            }
+    pub(crate) fn take_changes(&mut self) -> Option<FxHashSet<FileId>> {
+        if self.changes.is_empty() {
+            return None;
         }
+        Some(mem::take(&mut self.changes))
     }
 }
 
diff --git a/crates/rust-analyzer/src/handlers.rs b/crates/rust-analyzer/src/handlers.rs
index b2ff9a157..12b494496 100644
--- a/crates/rust-analyzer/src/handlers.rs
+++ b/crates/rust-analyzer/src/handlers.rs
@@ -31,7 +31,6 @@ use stdx::{format_to, split_delim};
 use crate::{
     cargo_target_spec::CargoTargetSpec,
     config::RustfmtConfig,
-    diagnostics::DiagnosticTask,
     from_json, from_proto,
     global_state::GlobalStateSnapshot,
     lsp_ext::{self, InlayHint, InlayHintsParams},
@@ -950,7 +949,7 @@ pub(crate) fn handle_ssr(
 pub(crate) fn publish_diagnostics(
     snap: &GlobalStateSnapshot,
     file_id: FileId,
-) -> Result<DiagnosticTask> {
+) -> Result<Vec<Diagnostic>> {
     let _p = profile("publish_diagnostics");
     let line_index = snap.analysis.file_line_index(file_id)?;
     let diagnostics: Vec<Diagnostic> = snap
@@ -967,7 +966,7 @@ pub(crate) fn publish_diagnostics(
             tags: None,
         })
         .collect();
-    Ok(DiagnosticTask::SetNative(file_id, diagnostics))
+    Ok(diagnostics)
 }
 
 pub(crate) fn handle_inlay_hints(
diff --git a/crates/rust-analyzer/src/main_loop.rs b/crates/rust-analyzer/src/main_loop.rs
index 76dea3e22..1bd9d6389 100644
--- a/crates/rust-analyzer/src/main_loop.rs
+++ b/crates/rust-analyzer/src/main_loop.rs
@@ -15,7 +15,6 @@ use ra_project_model::{PackageRoot, ProjectWorkspace};
 
 use crate::{
     config::{Config, FilesWatcher, LinkedProject},
-    diagnostics::DiagnosticTask,
     dispatch::{NotificationDispatcher, RequestDispatcher},
     from_proto,
     global_state::{file_id_to_url, GlobalState, Status},
@@ -132,6 +131,13 @@ enum Event {
     Flycheck(flycheck::Message),
 }
 
+#[derive(Debug)]
+pub(crate) enum Task {
+    Response(Response),
+    Diagnostics(Vec<(FileId, Vec<lsp_types::Diagnostic>)>),
+    Unit,
+}
+
 impl fmt::Debug for Event {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let debug_verbose_not = |not: &Notification, f: &mut fmt::Formatter| {
@@ -219,8 +225,10 @@ impl GlobalState {
             Event::Task(task) => {
                 match task {
                     Task::Response(response) => self.respond(response),
-                    Task::Diagnostics(tasks) => {
-                        tasks.into_iter().for_each(|task| on_diagnostic_task(task, self))
+                    Task::Diagnostics(diagnostics_per_file) => {
+                        for (file_id, diagnostics) in diagnostics_per_file {
+                            self.diagnostics.set_native_diagnostics(file_id, diagnostics)
+                        }
                     }
                     Task::Unit => (),
                 }
@@ -257,9 +265,7 @@ impl GlobalState {
                 }
             },
             Event::Flycheck(task) => match task {
-                flycheck::Message::ClearDiagnostics => {
-                    on_diagnostic_task(DiagnosticTask::ClearCheck, self)
-                }
+                flycheck::Message::ClearDiagnostics => self.diagnostics.clear_check(),
 
                 flycheck::Message::AddDiagnostic { workspace_root, diagnostic } => {
                     let diagnostics = crate::diagnostics::to_proto::map_rust_diagnostic_to_lsp(
@@ -279,15 +285,7 @@ impl GlobalState {
                                 return Ok(());
                             }
                         };
-
-                        on_diagnostic_task(
-                            DiagnosticTask::AddCheck(
-                                file_id,
-                                diag.diagnostic,
-                                diag.fixes.into_iter().map(|it| it.into()).collect(),
-                            ),
-                            self,
-                        )
+                        self.diagnostics.add_check_diagnostic(file_id, diag.diagnostic, diag.fixes)
                     }
                 }
 
@@ -322,9 +320,20 @@ impl GlobalState {
             self.update_file_notifications_on_threadpool(subscriptions);
         }
 
+        if let Some(diagnostic_changes) = self.diagnostics.take_changes() {
+            for file_id in diagnostic_changes {
+                let url = file_id_to_url(&self.vfs.read().0, file_id);
+                let diagnostics = self.diagnostics.diagnostics_for(file_id).cloned().collect();
+                let params =
+                    lsp_types::PublishDiagnosticsParams { uri: url, diagnostics, version: None };
+                let not = notification_new::<lsp_types::notification::PublishDiagnostics>(params);
+                self.send(not.into());
+            }
+        }
+
         let loop_duration = loop_start.elapsed();
         if loop_duration > Duration::from_millis(100) {
-            log::error!("overly long loop turn: {:?}", loop_duration);
+            log::warn!("overly long loop turn: {:?}", loop_duration);
             if env::var("RA_PROFILE").is_ok() {
                 self.show_message(
                     lsp_types::MessageType::Error,
@@ -516,6 +525,7 @@ impl GlobalState {
                                 ()
                             })
                             .ok()
+                            .map(|diags| (file_id, diags))
                     })
                     .collect::<Vec<_>>();
                 Task::Diagnostics(diagnostics)
@@ -532,29 +542,10 @@ impl GlobalState {
     }
 }
 
-#[derive(Debug)]
-pub(crate) enum Task {
-    Response(Response),
-    Diagnostics(()),
-    Unit,
-}
-
 pub(crate) type ReqHandler = fn(&mut GlobalState, Response);
 pub(crate) type ReqQueue = lsp_server::ReqQueue<(String, Instant), ReqHandler>;
 const DO_NOTHING: ReqHandler = |_, _| ();
 
-fn on_diagnostic_task(task: DiagnosticTask, global_state: &mut GlobalState) {
-    let subscriptions = global_state.diagnostics.handle_task(task);
-
-    for file_id in subscriptions {
-        let url = file_id_to_url(&global_state.vfs.read().0, file_id);
-        let diagnostics = global_state.diagnostics.diagnostics_for(file_id).cloned().collect();
-        let params = lsp_types::PublishDiagnosticsParams { uri: url, diagnostics, version: None };
-        let not = notification_new::<lsp_types::notification::PublishDiagnostics>(params);
-        global_state.send(not.into());
-    }
-}
-
 #[derive(Eq, PartialEq)]
 enum Progress {
     Begin,