Clear flycheck diagnostics more granularly

Veykril · Veykril · commit 98fcc051838e · 2024-12-20T14:09:11.000+01:00
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/diagnostics.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/diagnostics.rs
@@ -3,6 +3,7 @@ pub(crate) mod to_proto;
 
 use std::mem;
 
+use cargo_metadata::PackageId;
 use ide::FileId;
 use ide_db::FxHashMap;
 use itertools::Itertools;
@@ -13,7 +14,8 @@ use triomphe::Arc;
 
 use crate::{global_state::GlobalStateSnapshot, lsp, lsp_ext, main_loop::DiagnosticsTaskKind};
 
-pub(crate) type CheckFixes = Arc<IntMap<usize, IntMap<FileId, Vec<Fix>>>>;
+pub(crate) type CheckFixes =
+    Arc<IntMap<usize, FxHashMap<Option<PackageId>, IntMap<FileId, Vec<Fix>>>>>;
 
 #[derive(Debug, Default, Clone)]
 pub struct DiagnosticsMapConfig {
@@ -31,7 +33,8 @@ pub(crate) struct DiagnosticCollection {
     pub(crate) native_syntax: IntMap<FileId, (DiagnosticsGeneration, Vec<lsp_types::Diagnostic>)>,
     pub(crate) native_semantic: 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:
+        IntMap<usize, FxHashMap<Option<PackageId>, IntMap<FileId, Vec<lsp_types::Diagnostic>>>>,
     pub(crate) check_fixes: CheckFixes,
     changes: IntSet<FileId>,
     /// Counter for supplying a new generation number for diagnostics.
@@ -50,18 +53,19 @@ pub(crate) struct Fix {
 
 impl DiagnosticCollection {
     pub(crate) fn clear_check(&mut self, flycheck_id: usize) {
-        if let Some(it) = Arc::make_mut(&mut self.check_fixes).get_mut(&flycheck_id) {
+        if let Some(it) = self.check.get_mut(&flycheck_id) {
             it.clear();
         }
-        if let Some(it) = self.check.get_mut(&flycheck_id) {
-            self.changes.extend(it.drain().map(|(key, _value)| key));
+        if let Some(fixes) = Arc::make_mut(&mut self.check_fixes).get_mut(&flycheck_id) {
+            fixes.clear();
         }
     }
 
     pub(crate) fn clear_check_all(&mut self) {
         Arc::make_mut(&mut self.check_fixes).clear();
-        self.changes
-            .extend(self.check.values_mut().flat_map(|it| it.drain().map(|(key, _value)| key)))
+        self.changes.extend(
+            self.check.values_mut().flat_map(|it| it.drain().flat_map(|(_, v)| v.into_keys())),
+        )
     }
 
     pub(crate) fn clear_native_for(&mut self, file_id: FileId) {
@@ -70,14 +74,29 @@ impl DiagnosticCollection {
         self.changes.insert(file_id);
     }
 
+    pub(crate) fn clear_check_for_package(&mut self, flycheck_id: usize, package_id: PackageId) {
+        let Some(check) = self.check.get_mut(&flycheck_id) else {
+            return;
+        };
+        check.remove(&Some(package_id));
+    }
+
     pub(crate) fn add_check_diagnostic(
         &mut self,
         flycheck_id: usize,
+        package_id: &Option<PackageId>,
         file_id: FileId,
         diagnostic: lsp_types::Diagnostic,
         fix: Option<Box<Fix>>,
     ) {
-        let diagnostics = self.check.entry(flycheck_id).or_default().entry(file_id).or_default();
+        let diagnostics = self
+            .check
+            .entry(flycheck_id)
+            .or_default()
+            .entry(package_id.clone())
+            .or_default()
+            .entry(file_id)
+            .or_default();
         for existing_diagnostic in diagnostics.iter() {
             if are_diagnostics_equal(existing_diagnostic, &diagnostic) {
                 return;
@@ -86,7 +105,14 @@ impl DiagnosticCollection {
 
         if let Some(fix) = fix {
             let check_fixes = Arc::make_mut(&mut self.check_fixes);
-            check_fixes.entry(flycheck_id).or_default().entry(file_id).or_default().push(*fix);
+            check_fixes
+                .entry(flycheck_id)
+                .or_default()
+                .entry(package_id.clone())
+                .or_default()
+                .entry(file_id)
+                .or_default()
+                .push(*fix);
         }
         diagnostics.push(diagnostic);
         self.changes.insert(file_id);
@@ -135,7 +161,12 @@ impl DiagnosticCollection {
     ) -> impl Iterator<Item = &lsp_types::Diagnostic> {
         let native_syntax = self.native_syntax.get(&file_id).into_iter().flat_map(|(_, d)| d);
         let native_semantic = self.native_semantic.get(&file_id).into_iter().flat_map(|(_, d)| d);
-        let check = self.check.values().filter_map(move |it| it.get(&file_id)).flatten();
+        let check = self
+            .check
+            .values()
+            .flat_map(|it| it.values())
+            .filter_map(move |it| it.get(&file_id))
+            .flatten();
         native_syntax.chain(native_semantic).chain(check)
     }
 
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/flycheck.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/flycheck.rs
@@ -1,8 +1,9 @@
 //! Flycheck provides the functionality needed to run `cargo check` to provide
 //! LSP diagnostics based on the output of the command.
 
-use std::{fmt, io, process::Command, time::Duration};
+use std::{fmt, io, mem, process::Command, time::Duration};
 
+use cargo_metadata::PackageId;
 use crossbeam_channel::{select_biased, unbounded, Receiver, Sender};
 use paths::{AbsPath, AbsPathBuf, Utf8PathBuf};
 use rustc_hash::FxHashMap;
@@ -150,10 +151,19 @@ impl FlycheckHandle {
 
 pub(crate) enum FlycheckMessage {
     /// Request adding a diagnostic with fixes included to a file
-    AddDiagnostic { id: usize, workspace_root: AbsPathBuf, diagnostic: Diagnostic },
+    AddDiagnostic {
+        id: usize,
+        workspace_root: AbsPathBuf,
+        diagnostic: Diagnostic,
+        package_id: Option<PackageId>,
+    },
 
-    /// Request clearing all previous diagnostics
-    ClearDiagnostics { id: usize },
+    /// Request clearing all outdated diagnostics.
+    ClearDiagnostics {
+        id: usize,
+        /// The pacakge whose diagnostics to clear, or if unspecified, all diagnostics.
+        package_id: Option<PackageId>,
+    },
 
     /// Request check progress notification to client
     Progress {
@@ -166,15 +176,18 @@ pub(crate) enum FlycheckMessage {
 impl fmt::Debug for FlycheckMessage {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
-            FlycheckMessage::AddDiagnostic { id, workspace_root, diagnostic } => f
+            FlycheckMessage::AddDiagnostic { id, workspace_root, diagnostic, package_id } => f
                 .debug_struct("AddDiagnostic")
                 .field("id", id)
                 .field("workspace_root", workspace_root)
+                .field("package_id", package_id)
                 .field("diagnostic_code", &diagnostic.code.as_ref().map(|it| &it.code))
                 .finish(),
-            FlycheckMessage::ClearDiagnostics { id } => {
-                f.debug_struct("ClearDiagnostics").field("id", id).finish()
-            }
+            FlycheckMessage::ClearDiagnostics { id, package_id } => f
+                .debug_struct("ClearDiagnostics")
+                .field("id", id)
+                .field("package_id", package_id)
+                .finish(),
             FlycheckMessage::Progress { id, progress } => {
                 f.debug_struct("Progress").field("id", id).field("progress", progress).finish()
             }
@@ -200,6 +213,7 @@ enum StateChange {
 struct FlycheckActor {
     /// The workspace id of this flycheck instance.
     id: usize,
+
     sender: Sender<FlycheckMessage>,
     config: FlycheckConfig,
     manifest_path: Option<AbsPathBuf>,
@@ -215,8 +229,13 @@ struct FlycheckActor {
     command_handle: Option<CommandHandle<CargoCheckMessage>>,
     /// The receiver side of the channel mentioned above.
     command_receiver: Option<Receiver<CargoCheckMessage>>,
+    package_status: FxHashMap<PackageId, DiagnosticReceived>,
+}
 
-    status: FlycheckStatus,
+#[derive(PartialEq, Eq, Copy, Clone, Debug)]
+enum DiagnosticReceived {
+    Yes,
+    No,
 }
 
 #[allow(clippy::large_enum_variant)]
@@ -225,13 +244,6 @@ enum Event {
     CheckEvent(Option<CargoCheckMessage>),
 }
 
-#[derive(PartialEq)]
-enum FlycheckStatus {
-    Started,
-    DiagnosticSent,
-    Finished,
-}
-
 pub(crate) const SAVED_FILE_PLACEHOLDER: &str = "$saved_file";
 
 impl FlycheckActor {
@@ -253,7 +265,7 @@ impl FlycheckActor {
             manifest_path,
             command_handle: None,
             command_receiver: None,
-            status: FlycheckStatus::Finished,
+            package_status: FxHashMap::default(),
         }
     }
 
@@ -306,13 +318,11 @@ impl FlycheckActor {
                             self.command_handle = Some(command_handle);
                             self.command_receiver = Some(receiver);
                             self.report_progress(Progress::DidStart);
-                            self.status = FlycheckStatus::Started;
                         }
                         Err(error) => {
                             self.report_progress(Progress::DidFailToRestart(format!(
                                 "Failed to run the following command: {formatted_command} error={error}"
                             )));
-                            self.status = FlycheckStatus::Finished;
                         }
                     }
                 }
@@ -332,11 +342,25 @@ impl FlycheckActor {
                             error
                         );
                     }
-                    if self.status == FlycheckStatus::Started {
-                        self.send(FlycheckMessage::ClearDiagnostics { id: self.id });
+                    if self.package_status.is_empty() {
+                        // We finished without receiving any diagnostics.
+                        // That means all of them are stale.
+                        self.send(FlycheckMessage::ClearDiagnostics {
+                            id: self.id,
+                            package_id: None,
+                        });
+                    } else {
+                        for (package_id, status) in mem::take(&mut self.package_status) {
+                            if let DiagnosticReceived::No = status {
+                                self.send(FlycheckMessage::ClearDiagnostics {
+                                    id: self.id,
+                                    package_id: Some(package_id),
+                                });
+                            }
+                        }
                     }
+
                     self.report_progress(Progress::DidFinish(res));
-                    self.status = FlycheckStatus::Finished;
                 }
                 Event::CheckEvent(Some(message)) => match message {
                     CargoCheckMessage::CompilerArtifact(msg) => {
@@ -346,23 +370,30 @@ impl FlycheckActor {
                             "artifact received"
                         );
                         self.report_progress(Progress::DidCheckCrate(msg.target.name));
+                        self.package_status.entry(msg.package_id).or_insert(DiagnosticReceived::No);
                     }
-
-                    CargoCheckMessage::Diagnostic(msg) => {
+                    CargoCheckMessage::Diagnostic { diagnostic, package_id } => {
                         tracing::trace!(
                             flycheck_id = self.id,
-                            message = msg.message,
+                            message = diagnostic.message,
                             "diagnostic received"
                         );
-                        if self.status == FlycheckStatus::Started {
-                            self.send(FlycheckMessage::ClearDiagnostics { id: self.id });
+                        if let Some(package_id) = &package_id {
+                            if !self.package_status.contains_key(package_id) {
+                                self.package_status
+                                    .insert(package_id.clone(), DiagnosticReceived::Yes);
+                                self.send(FlycheckMessage::ClearDiagnostics {
+                                    id: self.id,
+                                    package_id: Some(package_id.clone()),
+                                });
+                            }
                         }
                         self.send(FlycheckMessage::AddDiagnostic {
                             id: self.id,
+                            package_id,
                             workspace_root: self.root.clone(),
-                            diagnostic: msg,
+                            diagnostic,
                         });
-                        self.status = FlycheckStatus::DiagnosticSent;
                     }
                 },
             }
@@ -380,7 +411,7 @@ impl FlycheckActor {
             command_handle.cancel();
             self.command_receiver.take();
             self.report_progress(Progress::DidCancel);
-            self.status = FlycheckStatus::Finished;
+            self.package_status.clear();
         }
     }
 
@@ -486,7 +517,7 @@ impl FlycheckActor {
 #[allow(clippy::large_enum_variant)]
 enum CargoCheckMessage {
     CompilerArtifact(cargo_metadata::Artifact),
-    Diagnostic(Diagnostic),
+    Diagnostic { diagnostic: Diagnostic, package_id: Option<PackageId> },
 }
 
 impl ParseFromLine for CargoCheckMessage {
@@ -501,11 +532,16 @@ impl ParseFromLine for CargoCheckMessage {
                         Some(CargoCheckMessage::CompilerArtifact(artifact))
                     }
                     cargo_metadata::Message::CompilerMessage(msg) => {
-                        Some(CargoCheckMessage::Diagnostic(msg.message))
+                        Some(CargoCheckMessage::Diagnostic {
+                            diagnostic: msg.message,
+                            package_id: Some(msg.package_id),
+                        })
                     }
                     _ => None,
                 },
-                JsonMessage::Rustc(message) => Some(CargoCheckMessage::Diagnostic(message)),
+                JsonMessage::Rustc(message) => {
+                    Some(CargoCheckMessage::Diagnostic { diagnostic: message, package_id: None })
+                }
             };
         }
 
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/handlers/request.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/handlers/request.rs
@@ -1441,7 +1441,13 @@ pub(crate) fn handle_code_action(
     }
 
     // Fixes from `cargo check`.
-    for fix in snap.check_fixes.values().filter_map(|it| it.get(&frange.file_id)).flatten() {
+    for fix in snap
+        .check_fixes
+        .values()
+        .flat_map(|it| it.values())
+        .filter_map(|it| it.get(&frange.file_id))
+        .flatten()
+    {
         // FIXME: this mapping is awkward and shouldn't exist. Refactor
         // `snap.check_fixes` to not convert to LSP prematurely.
         let intersect_fix_range = fix
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
@@ -956,7 +956,7 @@ impl GlobalState {
 
     fn handle_flycheck_msg(&mut self, message: FlycheckMessage) {
         match message {
-            FlycheckMessage::AddDiagnostic { id, workspace_root, diagnostic } => {
+            FlycheckMessage::AddDiagnostic { id, workspace_root, diagnostic, package_id } => {
                 let snap = self.snapshot();
                 let diagnostics = crate::diagnostics::to_proto::map_rust_diagnostic_to_lsp(
                     &self.config.diagnostics_map(None),
@@ -968,6 +968,7 @@ impl GlobalState {
                     match url_to_file_id(&self.vfs.read().0, &diag.url) {
                         Ok(file_id) => self.diagnostics.add_check_diagnostic(
                             id,
+                            &package_id,
                             file_id,
                             diag.diagnostic,
                             diag.fix,
@@ -981,9 +982,12 @@ impl GlobalState {
                     };
                 }
             }
-
-            FlycheckMessage::ClearDiagnostics { id } => self.diagnostics.clear_check(id),
-
+            FlycheckMessage::ClearDiagnostics { id, package_id: None } => {
+                self.diagnostics.clear_check(id)
+            }
+            FlycheckMessage::ClearDiagnostics { id, package_id: Some(package_id) } => {
+                self.diagnostics.clear_check_for_package(id, package_id)
+            }
             FlycheckMessage::Progress { id, progress } => {
                 let (state, message) = match progress {
                     flycheck::Progress::DidStart => (Progress::Begin, None),
