Color swatches ( 🟩 green 🟥 #ffaaaa ) (#12308)

helix-term/src/handlers.rs

@@ -10,9 +10,12 @@ use crate::handlers::signature_help::SignatureHelpHandler;
 
 pub use helix_view::handlers::Handlers;
 
+use self::document_colors::DocumentColorsHandler;
+
 mod auto_save;
 pub mod completion;
 mod diagnostics;
+mod document_colors;
 mod signature_help;
 mod snippet;
 
@@ -22,11 +25,13 @@ pub fn setup(config: Arc<ArcSwap<Config>>) -> Handlers {
     let event_tx = completion::CompletionHandler::new(config).spawn();
     let signature_hints = SignatureHelpHandler::new().spawn();
     let auto_save = AutoSaveHandler::new().spawn();
+    let document_colors = DocumentColorsHandler::default().spawn();
 
     let handlers = Handlers {
         completions: helix_view::handlers::completion::CompletionHandler::new(event_tx),
         signature_hints,
         auto_save,
+        document_colors,
     };
 
     helix_view::handlers::register_hooks(&handlers);
@@ -35,5 +40,6 @@ pub fn setup(config: Arc<ArcSwap<Config>>) -> Handlers {
     auto_save::register_hooks(&handlers);
     diagnostics::register_hooks(&handlers);
     snippet::register_hooks(&handlers);
+    document_colors::register_hooks(&handlers);
     handlers
 }

helix-term/src/handlers/document_colors.rs

@@ -0,0 +1,204 @@
+use std::{collections::HashSet, time::Duration};
+
+use futures_util::{stream::FuturesOrdered, StreamExt};
+use helix_core::{syntax::LanguageServerFeature, text_annotations::InlineAnnotation};
+use helix_event::{cancelable_future, register_hook};
+use helix_lsp::lsp;
+use helix_view::{
+    document::DocumentColorSwatches,
+    events::{DocumentDidChange, DocumentDidOpen, LanguageServerExited, LanguageServerInitialized},
+    handlers::{lsp::DocumentColorsEvent, Handlers},
+    DocumentId, Editor, Theme,
+};
+use tokio::time::Instant;
+
+use crate::job;
+
+#[derive(Default)]
+pub(super) struct DocumentColorsHandler {
+    docs: HashSet<DocumentId>,
+}
+
+const DOCUMENT_CHANGE_DEBOUNCE: Duration = Duration::from_millis(250);
+
+impl helix_event::AsyncHook for DocumentColorsHandler {
+    type Event = DocumentColorsEvent;
+
+    fn handle_event(&mut self, event: Self::Event, _timeout: Option<Instant>) -> Option<Instant> {
+        let DocumentColorsEvent(doc_id) = event;
+        self.docs.insert(doc_id);
+        Some(Instant::now() + DOCUMENT_CHANGE_DEBOUNCE)
+    }
+
+    fn finish_debounce(&mut self) {
+        let docs = std::mem::take(&mut self.docs);
+
+        job::dispatch_blocking(move |editor, _compositor| {
+            for doc in docs {
+                request_document_colors(editor, doc);
+            }
+        });
+    }
+}
+
+fn request_document_colors(editor: &mut Editor, doc_id: DocumentId) {
+    if !editor.config().lsp.display_color_swatches {
+        return;
+    }
+
+    let Some(doc) = editor.document_mut(doc_id) else {
+        return;
+    };
+
+    let cancel = doc.color_swatch_controller.restart();
+
+    let mut seen_language_servers = HashSet::new();
+    let mut futures: FuturesOrdered<_> = doc
+        .language_servers_with_feature(LanguageServerFeature::DocumentColors)
+        .filter(|ls| seen_language_servers.insert(ls.id()))
+        .map(|language_server| {
+            let text = doc.text().clone();
+            let offset_encoding = language_server.offset_encoding();
+            let future = language_server
+                .text_document_document_color(doc.identifier(), None)
+                .unwrap();
+
+            async move {
+                let colors: Vec<_> = future
+                    .await?
+                    .into_iter()
+                    .filter_map(|color_info| {
+                        let pos = helix_lsp::util::lsp_pos_to_pos(
+                            &text,
+                            color_info.range.start,
+                            offset_encoding,
+                        )?;
+                        Some((pos, color_info.color))
+                    })
+                    .collect();
+                anyhow::Ok(colors)
+            }
+        })
+        .collect();
+
+    tokio::spawn(async move {
+        let mut all_colors = Vec::new();
+        loop {
+            match cancelable_future(futures.next(), &cancel).await {
+                Some(Some(Ok(items))) => all_colors.extend(items),
+                Some(Some(Err(err))) => log::error!("document color request failed: {err}"),
+                Some(None) => break,
+                // The request was cancelled.
+                None => return,
+            }
+        }
+        job::dispatch(move |editor, _| attach_document_colors(editor, doc_id, all_colors)).await;
+    });
+}
+
+fn attach_document_colors(
+    editor: &mut Editor,
+    doc_id: DocumentId,
+    mut doc_colors: Vec<(usize, lsp::Color)>,
+) {
+    if !editor.config().lsp.display_color_swatches {
+        return;
+    }
+
+    let Some(doc) = editor.documents.get_mut(&doc_id) else {
+        return;
+    };
+
+    if doc_colors.is_empty() {
+        doc.color_swatches.take();
+        return;
+    }
+
+    doc_colors.sort_by_key(|(pos, _)| *pos);
+
+    let mut color_swatches = Vec::with_capacity(doc_colors.len());
+    let mut color_swatches_padding = Vec::with_capacity(doc_colors.len());
+    let mut colors = Vec::with_capacity(doc_colors.len());
+
+    for (pos, color) in doc_colors {
+        color_swatches_padding.push(InlineAnnotation::new(pos, " "));
+        color_swatches.push(InlineAnnotation::new(pos, "■"));
+        colors.push(Theme::rgb_highlight(
+            (color.red * 255.) as u8,
+            (color.green * 255.) as u8,
+            (color.blue * 255.) as u8,
+        ));
+    }
+
+    doc.color_swatches = Some(DocumentColorSwatches {
+        color_swatches,
+        colors,
+        color_swatches_padding,
+    });
+}
+
+pub(super) fn register_hooks(handlers: &Handlers) {
+    register_hook!(move |event: &mut DocumentDidOpen<'_>| {
+        // when a document is initially opened, request colors for it
+        request_document_colors(event.editor, event.doc);
+
+        Ok(())
+    });
+
+    let tx = handlers.document_colors.clone();
+    register_hook!(move |event: &mut DocumentDidChange<'_>| {
+        // Update the color swatch' positions, helping ensure they are displayed in the
+        // proper place.
+        let apply_color_swatch_changes = |annotations: &mut Vec<InlineAnnotation>| {
+            event.changes.update_positions(
+                annotations
+                    .iter_mut()
+                    .map(|annotation| (&mut annotation.char_idx, helix_core::Assoc::After)),
+            );
+        };
+
+        if let Some(DocumentColorSwatches {
+            color_swatches,
+            colors: _colors,
+            color_swatches_padding,
+        }) = &mut event.doc.color_swatches
+        {
+            apply_color_swatch_changes(color_swatches);
+            apply_color_swatch_changes(color_swatches_padding);
+        }
+
+        // Cancel the ongoing request, if present.
+        event.doc.color_swatch_controller.cancel();
+
+        helix_event::send_blocking(&tx, DocumentColorsEvent(event.doc.id()));
+
+        Ok(())
+    });
+
+    register_hook!(move |event: &mut LanguageServerInitialized<'_>| {
+        let doc_ids: Vec<_> = event.editor.documents().map(|doc| doc.id()).collect();
+
+        for doc_id in doc_ids {
+            request_document_colors(event.editor, doc_id);
+        }
+
+        Ok(())
+    });
+
+    register_hook!(move |event: &mut LanguageServerExited<'_>| {
+        // Clear and re-request all color swatches when a server exits.
+        for doc in event.editor.documents_mut() {
+            if doc.supports_language_server(event.server_id) {
+                doc.color_swatches.take();
+            }
+        }
+
+        let doc_ids: Vec<_> = event.editor.documents().map(|doc| doc.id()).collect();
+
+        for doc_id in doc_ids {
+            request_document_colors(event.editor, doc_id);
+        }
+
+        Ok(())
+    });
+}