refactor: make diagnostic sources pluggable

src/global_state.rs

@@ -16,7 +16,7 @@ pub(crate) mod task;
 mod trace;
 mod workspace_state;
 
-use std::time::Instant;
+use std::{sync::Arc as StdArc, time::Instant};
 
 use crossbeam_channel::{Receiver, Sender, unbounded};
 use hir::base_db::{
@@ -44,7 +44,7 @@ pub(crate) use self::workspace_state::{
 use self::{
     diagnostics::{
         DiagnosticCommitFreshness, DiagnosticFileRevision, DiagnosticPublishFreshness,
-        publisher::DiagnosticPublishKey,
+        DiagnosticSource, publisher::DiagnosticPublishKey,
     },
     mem_docs::MemDocs,
     snapshot::GlobalStateSnapshot,
@@ -72,6 +72,7 @@ pub(crate) struct GlobalState {
     pub(crate) diagnostics: DiagnosticsState,
     pub(crate) workspace: WorkspaceState,
     pub(crate) qihe: qihe::Qihe,
+    pub(crate) external_sources: Vec<StdArc<dyn DiagnosticSource>>,
     pub(crate) tasks: TaskState,
 }
 
@@ -147,6 +148,11 @@ impl GlobalState {
             .raw_db_mut()
             .set_diagnostics_config_with_durability(diagnostics_config, Durability::HIGH);
 
+        let qihe_diagnostics = qihe::QiheDiagnostics::new();
+        let qihe = qihe::Qihe::new(qihe_diagnostics);
+        let qihe_source: StdArc<dyn DiagnosticSource> = StdArc::new(qihe.diagnostics_snapshot());
+        let external_sources = vec![qihe_source];
+
         GlobalState {
             client: ClientState {
                 sender,
@@ -181,7 +187,8 @@ impl GlobalState {
                 fetch_workspaces_task: ExclTask::default(),
                 registered_client_file_watcher_globs: None,
             },
-            qihe: qihe::Qihe::new(),
+            qihe,
+            external_sources,
             tasks: TaskState { task_pool },
         }
     }
@@ -201,7 +208,7 @@ impl GlobalState {
             vfs: Arc::clone(&self.workspace.vfs),
             mem_docs: self.analysis.mem_docs.clone(),
             sema_tokens_cache: Arc::clone(&self.analysis.semantic_tokens_cache),
-            qihe_diagnostics: self.qihe.diagnostics_snapshot(),
+            external_sources: self.external_sources.clone(),
             diagnostic_publish_freshness: self.diagnostic_publish_freshness(),
             diagnostic_file_revisions: self.diagnostics.diagnostic_file_revisions.clone(),
             cancellation,

src/global_state/diagnostics.rs

@@ -1,5 +1,6 @@
 use hir::base_db::{project::CompilationProfileId, source_root::SourceRootId};
 use lsp_types::Url;
+use rustc_hash::FxHashSet;
 use vfs::FileId;
 
 pub(crate) mod publisher;
@@ -20,6 +21,22 @@ impl DiagnosticCommitFreshness {
     }
 }
 
+pub(crate) trait DiagnosticSource: Send + Sync {
+    fn diagnostics(
+        &self,
+        file_id: FileId,
+        freshness: &DiagnosticCommitFreshness,
+    ) -> Vec<lsp_types::Diagnostic>;
+
+    fn external_revision(
+        &self,
+        file_id: FileId,
+        freshness: &DiagnosticCommitFreshness,
+    ) -> Option<DiagnosticExternalRevision>;
+
+    fn remove_deleted(&self, files: &FxHashSet<FileId>);
+}
+
 /// Freshness token for a diagnostics publish batch.
 ///
 /// Diagnostic contents and diagnostic publish targets can change
@@ -56,7 +73,7 @@ pub(crate) enum DiagnosticOwner {
     File(FileId),
     SourceRoot(SourceRootId),
     CompilationProfile(CompilationProfileId),
-    ExternalQihe { file: FileId },
+    External { source: &'static str, file: FileId },
 }
 
 impl DiagnosticOwner {
@@ -69,7 +86,7 @@ impl DiagnosticOwner {
             DiagnosticOwner::CompilationProfile(profile_id) => {
                 format!("compilation-profile:{}", profile_id.0)
             }
-            DiagnosticOwner::ExternalQihe { file } => format!("external-qihe:{}", file.0),
+            DiagnosticOwner::External { source, file } => format!("external-{source}:{}", file.0),
         }
     }
 }

src/global_state/handlers/request/diagnostics.rs

@@ -67,7 +67,10 @@ pub(crate) fn handle_workspace_diagnostic(
             .into_iter()
             .map(|diag| to_proto::diagnostic(snap.config.i18n, &line_info, diag))
             .collect::<Vec<_>>();
-        diag_items.extend(snap.qihe_diagnostics(file_id));
+        let freshness = snap.diagnostic_commit_freshness();
+        diag_items.extend(
+            snap.external_sources.iter().flat_map(|source| source.diagnostics(file_id, &freshness)),
+        );
 
         for target in targets {
             let uri = target.uri().clone();

src/global_state/process_changes.rs

@@ -117,7 +117,9 @@ impl GlobalState {
         if diagnostic_targets_changed {
             self.diagnostics.diagnostic_target_revision += 1;
         }
-        self.qihe.remove_deleted(&deleted_file_ids);
+        for source in &self.external_sources {
+            source.remove_deleted(&deleted_file_ids);
+        }
         self.clear_deleted_push_diagnostics(&deleted_push_diagnostics);
         if has_structure_changes {
             self.invalidate_diagnostics(DiagnosticInvalidation::WorkspaceChanged);

src/global_state/qihe.rs

@@ -5,7 +5,7 @@ use std::{
     panic::{self, AssertUnwindSafe},
     path::{Path, PathBuf},
     process::{Command, Stdio},
-    sync::LazyLock,
+    sync::{Arc as StdArc, LazyLock},
     thread::{self, JoinHandle},
     time::{SystemTime, UNIX_EPOCH},
 };
@@ -17,7 +17,7 @@ use lsp_types::{
     Diagnostic, DiagnosticRelatedInformation, DiagnosticSeverity, NumberOrString, notification,
     request,
 };
-use parking_lot::Mutex;
+use parking_lot::{Mutex, MutexGuard};
 use project_model::project_manifest::{ProjectManifest, ProjectManifestFileName};
 use regex::Regex;
 use rustc_hash::{FxHashMap, FxHashSet};
@@ -37,7 +37,8 @@ use super::{
     AnalysisState, ClientState, ConfigState, DEFAULT_REQ_HANDLER, DiagnosticsState, GlobalState,
     QiheDiagnosticState, TaskState, WorkspaceState,
     diagnostics::{
-        DiagnosticCommitFreshness, DiagnosticPublishFreshness,
+        DiagnosticCommitFreshness, DiagnosticExternalRevision, DiagnosticOwner,
+        DiagnosticPublishFreshness, DiagnosticSource,
         publisher::{DiagnosticsPublisher, PublishDiagnosticsBatch, PublishDiagnosticsTask},
     },
     respond::Progress,
@@ -71,6 +72,59 @@ const QIHE_LOG_BATCH_BYTES: usize = 8 * 1024;
 static ANSI_ESCAPE_RE: LazyLock<Regex> =
     LazyLock::new(|| Regex::new(r"\x1B\[[0-?]*[ -/]*[@-~]").unwrap());
 
+#[derive(Clone)]
+pub(crate) struct QiheDiagnostics {
+    states: Arc<Mutex<FxHashMap<FileId, QiheDiagnosticState>>>,
+}
+
+impl QiheDiagnostics {
+    pub(crate) fn new() -> Self {
+        Self { states: Arc::new(Mutex::new(FxHashMap::default())) }
+    }
+
+    fn lock(&self) -> MutexGuard<'_, FxHashMap<FileId, QiheDiagnosticState>> {
+        self.states.lock()
+    }
+}
+
+impl DiagnosticSource for QiheDiagnostics {
+    fn diagnostics(
+        &self,
+        file_id: FileId,
+        freshness: &DiagnosticCommitFreshness,
+    ) -> Vec<Diagnostic> {
+        self.lock()
+            .get(&file_id)
+            .filter(|state| state.freshness == *freshness)
+            .map(|state| state.diagnostics.clone())
+            .unwrap_or_default()
+    }
+
+    fn external_revision(
+        &self,
+        file_id: FileId,
+        freshness: &DiagnosticCommitFreshness,
+    ) -> Option<DiagnosticExternalRevision> {
+        self.lock().get(&file_id).filter(|state| state.freshness == *freshness).map(|state| {
+            DiagnosticExternalRevision::new(
+                DiagnosticOwner::External { source: QIHE, file: file_id },
+                state.generation,
+            )
+        })
+    }
+
+    fn remove_deleted(&self, files: &FxHashSet<FileId>) {
+        if files.is_empty() {
+            return;
+        }
+
+        let mut diagnostics = self.lock();
+        for file_id in files {
+            diagnostics.remove(file_id);
+        }
+    }
+}
+
 /// Monotonic identity for a Qihe analysis run.
 ///
 /// The server keeps latest-run semantics for Qihe: logs, diagnostics, and
@@ -151,23 +205,21 @@ pub(crate) struct Qihe {
     run_generation: QiheRunId,
     active_progress_token: Option<String>,
     active_cancel_token: Option<CancellationToken>,
-    diagnostics: Arc<Mutex<FxHashMap<FileId, QiheDiagnosticState>>>,
+    diagnostics: QiheDiagnostics,
 }
 
 impl Qihe {
-    pub(crate) fn new() -> Self {
+    pub(crate) fn new(diagnostics: QiheDiagnostics) -> Self {
         Self {
             run_generation: QiheRunId::default(),
             active_progress_token: None,
             active_cancel_token: None,
-            diagnostics: Arc::new(Mutex::new(FxHashMap::default())),
+            diagnostics,
         }
     }
 
-    pub(crate) fn diagnostics_snapshot(
-        &self,
-    ) -> Arc<Mutex<FxHashMap<FileId, QiheDiagnosticState>>> {
-        Arc::clone(&self.diagnostics)
+    pub(crate) fn diagnostics_snapshot(&self) -> QiheDiagnostics {
+        self.diagnostics.clone()
     }
 
     pub(crate) fn start<C: QiheCtx>(&mut self, params: RunQiheAnalysisParams, ctx: &mut C) {
@@ -177,7 +229,7 @@ impl Qihe {
         let progress_token = qihe_progress_token(run_id, &params.uri);
         let progress_label = params.uri.path().to_string();
         let cancellation = ctx.task_cancel_token();
-        let snapshot = ctx.make_snapshot(cancellation.clone(), self.diagnostics_snapshot());
+        let snapshot = ctx.make_snapshot(cancellation.clone());
 
         self.active_progress_token = Some(progress_token.clone());
         self.active_cancel_token = Some(cancellation.clone());
@@ -292,23 +344,12 @@ impl Qihe {
         }
     }
 
-    pub(crate) fn remove_deleted(&mut self, deleted_file_ids: &FxHashSet<FileId>) {
-        if deleted_file_ids.is_empty() {
-            return;
-        }
-
-        let mut diagnostics = self.diagnostics.lock();
-        for file_id in deleted_file_ids {
-            diagnostics.remove(file_id);
-        }
-    }
-
     pub(crate) fn publish_diagnostics<C: QiheCtx>(
         &mut self,
         changed_files: FxHashSet<FileId>,
         ctx: &mut C,
     ) {
-        ctx.publish_qihe_diagnostics(changed_files, self.diagnostics_snapshot());
+        ctx.publish_qihe_diagnostics(changed_files);
     }
 
     fn end_superseded<C: QiheCtx>(&mut self, ctx: &mut C) {
@@ -363,11 +404,7 @@ impl Qihe {
 pub(crate) trait QiheCtx {
     fn i18n_text(&self, key: QiheI18nKey) -> &str;
     fn diagnostic_commit_freshness(&self) -> DiagnosticCommitFreshness;
-    fn make_snapshot(
-        &self,
-        cancellation: CancellationToken,
-        diagnostics: Arc<Mutex<FxHashMap<FileId, QiheDiagnosticState>>>,
-    ) -> GlobalStateSnapshot;
+    fn make_snapshot(&self, cancellation: CancellationToken) -> GlobalStateSnapshot;
     fn spawn_qihe_task<F>(&mut self, task: F)
     where
         F: FnOnce(crossbeam_channel::Sender<Task>) + Send + 'static;
@@ -381,11 +418,7 @@ pub(crate) trait QiheCtx {
         fraction: Option<f64>,
         token: String,
     );
-    fn publish_qihe_diagnostics(
-        &mut self,
-        changed_files: FxHashSet<FileId>,
-        diagnostics: Arc<Mutex<FxHashMap<FileId, QiheDiagnosticState>>>,
-    );
+    fn publish_qihe_diagnostics(&mut self, changed_files: FxHashSet<FileId>);
 }
 
 #[derive(Clone, Copy)]
@@ -402,6 +435,7 @@ pub(super) struct QiheGlobalCtx<'a> {
     analysis: &'a mut AnalysisState,
     diagnostics: &'a mut DiagnosticsState,
     workspace: &'a mut WorkspaceState,
+    external_sources: &'a [StdArc<dyn DiagnosticSource>],
     tasks: &'a mut TaskState,
 }
 
@@ -469,19 +503,15 @@ impl QiheCtx for QiheGlobalCtx<'_> {
         self.diagnostic_publish_freshness().commit()
     }
 
-    fn make_snapshot(
-        &self,
-        cancellation: CancellationToken,
-        diagnostics: Arc<Mutex<FxHashMap<FileId, QiheDiagnosticState>>>,
-    ) -> GlobalStateSnapshot {
+    fn make_snapshot(&self, cancellation: CancellationToken) -> GlobalStateSnapshot {
         GlobalStateSnapshot {
             config: Arc::clone(&self.config_state.config),
             workspaces: Arc::clone(&self.workspace.workspaces),
             analysis: self.analysis.analysis_host.make_analysis(),
             vfs: Arc::clone(&self.workspace.vfs),
             mem_docs: self.analysis.mem_docs.clone(),
             sema_tokens_cache: Arc::clone(&self.analysis.semantic_tokens_cache),
-            qihe_diagnostics: diagnostics,
+            external_sources: self.external_sources.to_vec(),
             diagnostic_publish_freshness: self.diagnostic_publish_freshness(),
             diagnostic_file_revisions: self.diagnostics.diagnostic_file_revisions.clone(),
             cancellation,
@@ -562,11 +592,7 @@ impl QiheCtx for QiheGlobalCtx<'_> {
         });
     }
 
-    fn publish_qihe_diagnostics(
-        &mut self,
-        changed_files: FxHashSet<FileId>,
-        diagnostics: Arc<Mutex<FxHashMap<FileId, QiheDiagnosticState>>>,
-    ) {
+    fn publish_qihe_diagnostics(&mut self, changed_files: FxHashSet<FileId>) {
         if changed_files.is_empty() {
             return;
         }
@@ -576,7 +602,7 @@ impl QiheCtx for QiheGlobalCtx<'_> {
             return;
         }
 
-        let snapshot = self.make_snapshot(self.task_cancel_token(), diagnostics);
+        let snapshot = self.make_snapshot(self.task_cancel_token());
         let mut publish_tasks = Vec::with_capacity(changed_files.len());
         let mut touched_file_ids = FxHashSet::default();
         for file_id in changed_files.iter().copied() {
@@ -646,8 +672,25 @@ pub(super) fn with_global_ctx<T>(
     state: &mut GlobalState,
     f: impl FnOnce(&mut Qihe, &mut QiheGlobalCtx<'_>) -> T,
 ) -> T {
-    let GlobalState { client, config_state, analysis, diagnostics, workspace, qihe, tasks } = state;
-    let mut ctx = QiheGlobalCtx { client, config_state, analysis, diagnostics, workspace, tasks };
+    let GlobalState {
+        client,
+        config_state,
+        analysis,
+        diagnostics,
+        workspace,
+        qihe,
+        external_sources,
+        tasks,
+    } = state;
+    let mut ctx = QiheGlobalCtx {
+        client,
+        config_state,
+        analysis,
+        diagnostics,
+        workspace,
+        external_sources,
+        tasks,
+    };
     f(qihe, &mut ctx)
 }