From 6949048e3bc81aa01fcf670bcb0f34a0cf888642 Mon Sep 17 00:00:00 2001
From: Sofus Addington <sofus@addington.dk>
Date: Sat, 22 Feb 2025 17:13:01 +0100
Subject: [PATCH] Separate publish and pull diagnostic provider

---
 helix-core/src/diagnostic.rs           | 17 +++++++-
 helix-term/src/application.rs          | 17 ++++++--
 helix-term/src/commands/lsp.rs         |  9 ++--
 helix-term/src/commands/typed.rs       |  2 +-
 helix-term/src/handlers/diagnostics.rs | 19 +++++----
 helix-view/src/document.rs             | 26 ++++++++----
 helix-view/src/editor.rs               | 58 ++++++++++++++------------
 helix-view/src/gutter.rs               |  2 +-
 helix-view/src/handlers/lsp.rs         | 19 +++++----
 9 files changed, 106 insertions(+), 63 deletions(-)

diff --git a/helix-core/src/diagnostic.rs b/helix-core/src/diagnostic.rs
index 333c9409..78d48edb 100644
--- a/helix-core/src/diagnostic.rs
+++ b/helix-core/src/diagnostic.rs
@@ -50,8 +50,21 @@ pub struct Diagnostic {
     pub data: Option<serde_json::Value>,
 }
 
-// TODO turn this into an enum + feature flag when lsp becomes optional
-pub type DiagnosticProvider = LanguageServerId;
+// TODO turn this into a feature flag when lsp becomes optional
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub enum DiagnosticProvider {
+    PublishDiagnosticProvider(LanguageServerId),
+    PullDiagnosticProvider(LanguageServerId),
+}
+
+impl From<DiagnosticProvider> for LanguageServerId {
+    fn from(value: DiagnosticProvider) -> Self {
+        match value {
+            DiagnosticProvider::PublishDiagnosticProvider(id) => id,
+            DiagnosticProvider::PullDiagnosticProvider(id) => id,
+        }
+    }
+}
 
 // while I would prefer having this in helix-lsp that necessitates a bunch of
 // conversions I would rather not add. I think its fine since this just a very
diff --git a/helix-term/src/application.rs b/helix-term/src/application.rs
index f80cffd4..2d41c34e 100644
--- a/helix-term/src/application.rs
+++ b/helix-term/src/application.rs
@@ -1,6 +1,9 @@
 use arc_swap::{access::Map, ArcSwap};
 use futures_util::Stream;
-use helix_core::{diagnostic::Severity, pos_at_coords, syntax, Range, Selection};
+use helix_core::{
+    diagnostic::{DiagnosticProvider, Severity},
+    pos_at_coords, syntax, Range, Selection,
+};
 use helix_lsp::{
     lsp::{self, notification::Notification},
     util::lsp_range_to_range,
@@ -749,8 +752,12 @@ impl Application {
                             log::error!("Discarding publishDiagnostic notification sent by an uninitialized server: {}", language_server.name());
                             return;
                         }
+
+                        let diagnostic_provider =
+                            DiagnosticProvider::PublishDiagnosticProvider(language_server.id());
+
                         self.editor.handle_lsp_diagnostics(
-                            language_server.id(),
+                            diagnostic_provider,
                             uri,
                             params.version,
                             params.diagnostics,
@@ -864,14 +871,16 @@ impl Application {
                         // we need to clear those and remove the entries from the list if this leads to
                         // an empty diagnostic list for said files
                         for diags in self.editor.diagnostics.values_mut() {
-                            diags.retain(|(_, lsp_id)| *lsp_id != server_id);
+                            diags.retain(|(_, diagnostic_provider)| {
+                                Into::<LanguageServerId>::into(*diagnostic_provider) != server_id
+                            });
                         }
 
                         self.editor.diagnostics.retain(|_, diags| !diags.is_empty());
 
                         // Clear any diagnostics for documents with this server open.
                         for doc in self.editor.documents_mut() {
-                            doc.clear_diagnostics(Some(server_id));
+                            doc.clear_all_language_server_diagnostics(Some(server_id));
                         }
 
                         // Remove the language server from the registry.
diff --git a/helix-term/src/commands/lsp.rs b/helix-term/src/commands/lsp.rs
index c5442924..c506f7bb 100644
--- a/helix-term/src/commands/lsp.rs
+++ b/helix-term/src/commands/lsp.rs
@@ -14,7 +14,8 @@ use tui::{text::Span, widgets::Row};
 use super::{align_view, push_jump, Align, Context, Editor};
 
 use helix_core::{
-    syntax::LanguageServerFeature, text_annotations::InlineAnnotation, Selection, Uri,
+    diagnostic::DiagnosticProvider, syntax::LanguageServerFeature,
+    text_annotations::InlineAnnotation, Selection, Uri,
 };
 use helix_stdx::path;
 use helix_view::{
@@ -209,7 +210,7 @@ type DiagnosticsPicker = Picker<PickerDiagnostic, DiagnosticStyles>;
 
 fn diag_picker(
     cx: &Context,
-    diagnostics: BTreeMap<Uri, Vec<(lsp::Diagnostic, LanguageServerId)>>,
+    diagnostics: BTreeMap<Uri, Vec<(lsp::Diagnostic, DiagnosticProvider)>>,
     format: DiagnosticsFormat,
 ) -> DiagnosticsPicker {
     // TODO: drop current_path comparison and instead use workspace: bool flag?
@@ -219,8 +220,8 @@ fn diag_picker(
     for (uri, diags) in diagnostics {
         flat_diag.reserve(diags.len());
 
-        for (diag, ls) in diags {
-            if let Some(ls) = cx.editor.language_server_by_id(ls) {
+        for (diag, provider) in diags {
+            if let Some(ls) = cx.editor.language_server_by_id(provider.into()) {
                 flat_diag.push(PickerDiagnostic {
                     location: Location {
                         uri: uri.clone(),
diff --git a/helix-term/src/commands/typed.rs b/helix-term/src/commands/typed.rs
index bcb6e543..87580a7e 100644
--- a/helix-term/src/commands/typed.rs
+++ b/helix-term/src/commands/typed.rs
@@ -1623,7 +1623,7 @@ fn lsp_stop(cx: &mut compositor::Context, args: Args, event: PromptEvent) -> any
 
         for doc in cx.editor.documents_mut() {
             if let Some(client) = doc.remove_language_server_by_name(ls_name) {
-                doc.clear_diagnostics(Some(client.id()));
+                doc.clear_all_language_server_diagnostics(Some(client.id()));
                 doc.reset_all_inlay_hints();
                 doc.inlay_hints_oudated = true;
             }
diff --git a/helix-term/src/handlers/diagnostics.rs b/helix-term/src/handlers/diagnostics.rs
index c8566cfa..9e28ef52 100644
--- a/helix-term/src/handlers/diagnostics.rs
+++ b/helix-term/src/handlers/diagnostics.rs
@@ -1,10 +1,11 @@
 use std::collections::{HashMap, HashSet};
 use std::time::Duration;
 
+use helix_core::diagnostic::DiagnosticProvider;
 use helix_core::syntax::LanguageServerFeature;
 use helix_core::Uri;
 use helix_event::{register_hook, send_blocking};
-use helix_lsp::{lsp, LanguageServerId};
+use helix_lsp::lsp;
 use helix_view::document::Mode;
 use helix_view::events::{DiagnosticsDidChange, DocumentDidChange, DocumentDidOpen};
 use helix_view::handlers::diagnostics::DiagnosticEvent;
@@ -133,7 +134,7 @@ pub fn pull_diagnostics_for_document(
         return;
     };
 
-    let server_id = language_server.id();
+    let provider = DiagnosticProvider::PullDiagnosticProvider(language_server.id());
     let document_id = doc.id();
 
     tokio::spawn(async move {
@@ -145,7 +146,7 @@ pub fn pull_diagnostics_for_document(
                         Err(_) => return,
                     };
 
-                    handle_pull_diagnostics_response(editor, response, server_id, uri, document_id)
+                    handle_pull_diagnostics_response(editor, response, provider, uri, document_id)
                 })
                 .await
             }
@@ -157,7 +158,7 @@ pub fn pull_diagnostics_for_document(
 fn handle_pull_diagnostics_response(
     editor: &mut Editor,
     response: lsp::DocumentDiagnosticReport,
-    server_id: LanguageServerId,
+    provider: DiagnosticProvider,
     uri: Uri,
     document_id: DocumentId,
 ) {
@@ -169,7 +170,7 @@ fn handle_pull_diagnostics_response(
         lsp::DocumentDiagnosticReport::Full(report) => {
             // Diagnostic for requested file
             editor.handle_lsp_diagnostics(
-                server_id,
+                provider,
                 uri,
                 None,
                 report.full_document_diagnostic_report.items,
@@ -181,7 +182,7 @@ fn handle_pull_diagnostics_response(
                 editor,
                 document_id,
                 report.related_documents,
-                server_id,
+                provider,
             );
         }
         lsp::DocumentDiagnosticReport::Unchanged(report) => {
@@ -192,7 +193,7 @@ fn handle_pull_diagnostics_response(
                 editor,
                 document_id,
                 report.related_documents,
-                server_id,
+                provider,
             );
         }
     }
@@ -202,7 +203,7 @@ fn handle_document_diagnostic_report_kind(
     editor: &mut Editor,
     document_id: DocumentId,
     report: Option<HashMap<lsp::Url, lsp::DocumentDiagnosticReportKind>>,
-    server_id: LanguageServerId,
+    provider: DiagnosticProvider,
 ) {
     for (url, report) in report.into_iter().flatten() {
         match report {
@@ -211,7 +212,7 @@ fn handle_document_diagnostic_report_kind(
                     return;
                 };
 
-                editor.handle_lsp_diagnostics(server_id, uri, None, report.items, report.result_id);
+                editor.handle_lsp_diagnostics(provider, uri, None, report.items, report.result_id);
             }
             lsp::DocumentDiagnosticReportKind::Unchanged(report) => {
                 let Some(doc) = editor.document_mut(document_id) else {
diff --git a/helix-view/src/document.rs b/helix-view/src/document.rs
index 8f90aba4..5e066bb5 100644
--- a/helix-view/src/document.rs
+++ b/helix-view/src/document.rs
@@ -5,6 +5,7 @@ use futures_util::future::BoxFuture;
 use futures_util::FutureExt;
 use helix_core::auto_pairs::AutoPairs;
 use helix_core::chars::char_is_word;
+use helix_core::diagnostic::DiagnosticProvider;
 use helix_core::doc_formatter::TextFormat;
 use helix_core::encoding::Encoding;
 use helix_core::snippets::{ActiveSnippet, SnippetRenderCtx};
@@ -1934,7 +1935,7 @@ impl Document {
         text: &Rope,
         language_config: Option<&LanguageConfiguration>,
         diagnostic: &helix_lsp::lsp::Diagnostic,
-        language_server_id: LanguageServerId,
+        language_server_id: DiagnosticProvider,
         offset_encoding: helix_lsp::OffsetEncoding,
     ) -> Option<Diagnostic> {
         use helix_core::diagnostic::{Range, Severity::*};
@@ -2027,13 +2028,13 @@ impl Document {
         &mut self,
         diagnostics: impl IntoIterator<Item = Diagnostic>,
         unchanged_sources: &[String],
-        language_server_id: Option<LanguageServerId>,
+        diagnostic_provider: Option<DiagnosticProvider>,
     ) {
         if unchanged_sources.is_empty() {
-            self.clear_diagnostics(language_server_id);
+            self.clear_diagnostics(diagnostic_provider);
         } else {
             self.diagnostics.retain(|d| {
-                if language_server_id.is_some_and(|id| id != d.provider) {
+                if diagnostic_provider.is_some_and(|provider| provider != d.provider) {
                     return true;
                 }
 
@@ -2049,10 +2050,19 @@ impl Document {
             .sort_by_key(|diagnostic| (diagnostic.range, diagnostic.severity, diagnostic.provider));
     }
 
-    /// clears diagnostics for a given language server id if set, otherwise all diagnostics are cleared
-    pub fn clear_diagnostics(&mut self, language_server_id: Option<LanguageServerId>) {
-        if let Some(id) = language_server_id {
-            self.diagnostics.retain(|d| d.provider != id);
+    /// clears diagnostics for a given diagnostic provider if set, otherwise all diagnostics are cleared
+    pub fn clear_diagnostics(&mut self, provider: Option<DiagnosticProvider>) {
+        if let Some(provider) = provider {
+            self.diagnostics.retain(|d| d.provider != provider);
+        } else {
+            self.diagnostics.clear();
+        }
+    }
+
+    /// clears diagnostics for a given language_server if set, otherwise all diagnostics are cleared
+    pub fn clear_all_language_server_diagnostics(&mut self, server_id: Option<LanguageServerId>) {
+        if let Some(server_id) = server_id {
+            self.diagnostics.retain(|d| server_id != d.provider.into());
         } else {
             self.diagnostics.clear();
         }
diff --git a/helix-view/src/editor.rs b/helix-view/src/editor.rs
index 6581dd35..bcfe68d7 100644
--- a/helix-view/src/editor.rs
+++ b/helix-view/src/editor.rs
@@ -45,6 +45,7 @@ use anyhow::{anyhow, bail, Error};
 pub use helix_core::diagnostic::Severity;
 use helix_core::{
     auto_pairs::AutoPairs,
+    diagnostic::DiagnosticProvider,
     syntax::{self, AutoPairConfig, IndentationHeuristic, LanguageServerFeature, SoftWrap},
     Change, LineEnding, Position, Range, Selection, Uri, NATIVE_LINE_ENDING,
 };
@@ -1056,7 +1057,7 @@ pub struct Editor {
     pub macro_recording: Option<(char, Vec<KeyEvent>)>,
     pub macro_replaying: Vec<char>,
     pub language_servers: helix_lsp::Registry,
-    pub diagnostics: BTreeMap<Uri, Vec<(lsp::Diagnostic, LanguageServerId)>>,
+    pub diagnostics: BTreeMap<Uri, Vec<(lsp::Diagnostic, DiagnosticProvider)>>,
     pub diff_providers: DiffProviderRegistry,
 
     pub debugger: Option<dap::Client>,
@@ -2006,7 +2007,7 @@ impl Editor {
     /// Returns all supported diagnostics for the document
     pub fn doc_diagnostics<'a>(
         language_servers: &'a helix_lsp::Registry,
-        diagnostics: &'a BTreeMap<Uri, Vec<(lsp::Diagnostic, LanguageServerId)>>,
+        diagnostics: &'a BTreeMap<Uri, Vec<(lsp::Diagnostic, DiagnosticProvider)>>,
         document: &Document,
     ) -> impl Iterator<Item = helix_core::Diagnostic> + 'a {
         Editor::doc_diagnostics_with_filter(language_servers, diagnostics, document, |_, _| true)
@@ -2016,9 +2017,9 @@ impl Editor {
     /// filtered by `filter` which is invocated with the raw `lsp::Diagnostic` and the language server id it came from
     pub fn doc_diagnostics_with_filter<'a>(
         language_servers: &'a helix_lsp::Registry,
-        diagnostics: &'a BTreeMap<Uri, Vec<(lsp::Diagnostic, LanguageServerId)>>,
+        diagnostics: &'a BTreeMap<Uri, Vec<(lsp::Diagnostic, DiagnosticProvider)>>,
         document: &Document,
-        filter: impl Fn(&lsp::Diagnostic, LanguageServerId) -> bool + 'a,
+        filter: impl Fn(&lsp::Diagnostic, DiagnosticProvider) -> bool + 'a,
     ) -> impl Iterator<Item = helix_core::Diagnostic> + 'a {
         let text = document.text().clone();
         let language_config = document.language.clone();
@@ -2026,30 +2027,33 @@ impl Editor {
             .uri()
             .and_then(|uri| diagnostics.get(&uri))
             .map(|diags| {
-                diags.iter().filter_map(move |(diagnostic, lsp_id)| {
-                    let ls = language_servers.get_by_id(*lsp_id)?;
-                    language_config
-                        .as_ref()
-                        .and_then(|c| {
-                            c.language_servers.iter().find(|features| {
-                                features.name == ls.name()
-                                    && features.has_feature(LanguageServerFeature::Diagnostics)
+                diags
+                    .iter()
+                    .filter_map(move |(diagnostic, diagnostic_provider)| {
+                        let ls = language_servers
+                            .get_by_id(Into::<LanguageServerId>::into(*diagnostic_provider))?;
+                        language_config
+                            .as_ref()
+                            .and_then(|c| {
+                                c.language_servers.iter().find(|features| {
+                                    features.name == ls.name()
+                                        && features.has_feature(LanguageServerFeature::Diagnostics)
+                                })
                             })
-                        })
-                        .and_then(|_| {
-                            if filter(diagnostic, *lsp_id) {
-                                Document::lsp_diagnostic_to_diagnostic(
-                                    &text,
-                                    language_config.as_deref(),
-                                    diagnostic,
-                                    *lsp_id,
-                                    ls.offset_encoding(),
-                                )
-                            } else {
-                                None
-                            }
-                        })
-                })
+                            .and_then(|_| {
+                                if filter(diagnostic, *diagnostic_provider) {
+                                    Document::lsp_diagnostic_to_diagnostic(
+                                        &text,
+                                        language_config.as_deref(),
+                                        diagnostic,
+                                        *diagnostic_provider,
+                                        ls.offset_encoding(),
+                                    )
+                                } else {
+                                    None
+                                }
+                            })
+                    })
             })
             .into_iter()
             .flatten()
diff --git a/helix-view/src/gutter.rs b/helix-view/src/gutter.rs
index 7cd91271..b8456ffc 100644
--- a/helix-view/src/gutter.rs
+++ b/helix-view/src/gutter.rs
@@ -71,7 +71,7 @@ pub fn diagnostic<'doc>(
                     d.line == line
                         && doc
                             .language_servers_with_feature(LanguageServerFeature::Diagnostics)
-                            .any(|ls| ls.id() == d.provider)
+                            .any(|ls| ls.id() == d.provider.into())
                 });
             diagnostics_on_line.max_by_key(|d| d.severity).map(|d| {
                 write!(out, "●").ok();
diff --git a/helix-view/src/handlers/lsp.rs b/helix-view/src/handlers/lsp.rs
index 60d50f9e..989cd936 100644
--- a/helix-view/src/handlers/lsp.rs
+++ b/helix-view/src/handlers/lsp.rs
@@ -4,9 +4,10 @@ use std::fmt::Display;
 use crate::editor::Action;
 use crate::events::DiagnosticsDidChange;
 use crate::{DocumentId, Editor};
+use helix_core::diagnostic::DiagnosticProvider;
 use helix_core::Uri;
 use helix_lsp::util::generate_transaction_from_edits;
-use helix_lsp::{lsp, LanguageServerId, OffsetEncoding};
+use helix_lsp::{lsp, OffsetEncoding};
 
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub enum SignatureHelpInvoked {
@@ -280,7 +281,7 @@ impl Editor {
 
     pub fn handle_lsp_diagnostics(
         &mut self,
-        server_id: LanguageServerId,
+        diagnostic_provider: DiagnosticProvider,
         uri: Uri,
         version: Option<i32>,
         mut diagnostics: Vec<lsp::Diagnostic>,
@@ -315,7 +316,7 @@ impl Editor {
                 let old_diagnostics = old_diagnostics
                     .iter()
                     .filter(|(d, d_server)| {
-                        *d_server == server_id && d.source.as_ref() == Some(source)
+                        *d_server == diagnostic_provider && d.source.as_ref() == Some(source)
                     })
                     .map(|(d, _)| d);
                 if new_diagnostics.eq(old_diagnostics) {
@@ -324,7 +325,7 @@ impl Editor {
             }
         }
 
-        let diagnostics = diagnostics.into_iter().map(|d| (d, server_id));
+        let diagnostics = diagnostics.into_iter().map(|d| (d, diagnostic_provider));
 
         // Insert the original lsp::Diagnostics here because we may have no open document
         // for diagnostic message and so we can't calculate the exact position.
@@ -333,7 +334,7 @@ impl Editor {
             Entry::Occupied(o) => {
                 let current_diagnostics = o.into_mut();
                 // there may entries of other language servers, which is why we can't overwrite the whole entry
-                current_diagnostics.retain(|(_, lsp_id)| *lsp_id != server_id);
+                current_diagnostics.retain(|(_, lsp_id)| *lsp_id != diagnostic_provider);
                 current_diagnostics.extend(diagnostics);
                 current_diagnostics
                 // Sort diagnostics first by severity and then by line numbers.
@@ -348,7 +349,7 @@ impl Editor {
         if let Some(doc) = doc {
             let diagnostic_of_language_server_and_not_in_unchanged_sources =
                 |diagnostic: &lsp::Diagnostic, ls_id| {
-                    ls_id == server_id
+                    ls_id == diagnostic_provider
                         && diagnostic
                             .source
                             .as_ref()
@@ -360,7 +361,11 @@ impl Editor {
                 doc,
                 diagnostic_of_language_server_and_not_in_unchanged_sources,
             );
-            doc.replace_diagnostics(diagnostics, &unchanged_diag_sources, Some(server_id));
+            doc.replace_diagnostics(
+                diagnostics,
+                &unchanged_diag_sources,
+                Some(diagnostic_provider),
+            );
 
             if report_id.is_some() {
                 doc.previous_diagnostic_id = report_id;