Add rainbow tree-sitter highlights

ref: https://github.com/helix-editor/helix/issues/695
ref: https://github.com/helix-editor/helix/pull/2857

Co-authored-by: Michael Davis <mcarsondavis@gmail.com>

2 files changed, 63 insertions, 1 deletions

diff --git a/helix-term/src/health.rs b/helix-term/src/health.rs
index dff90319..978c9c49 100644
--- a/helix-term/src/health.rs
+++ b/helix-term/src/health.rs
@@ -12,11 +12,17 @@ pub enum TsFeature {
     Highlight,
     TextObject,
     AutoIndent,
+    RainbowBrackets,
 }
 
 impl TsFeature {
     pub fn all() -> &'static [Self] {
-        &[Self::Highlight, Self::TextObject, Self::AutoIndent]
+        &[
+            Self::Highlight,
+            Self::TextObject,
+            Self::AutoIndent,
+            Self::RainbowBrackets,
+        ]
     }
 
     pub fn runtime_filename(&self) -> &'static str {
@@ -24,6 +30,7 @@ impl TsFeature {
             Self::Highlight => "highlights.scm",
             Self::TextObject => "textobjects.scm",
             Self::AutoIndent => "indents.scm",
+            Self::RainbowBrackets => "rainbows.scm",
         }
     }
 
@@ -32,6 +39,7 @@ impl TsFeature {
             Self::Highlight => "Syntax Highlighting",
             Self::TextObject => "Treesitter Textobjects",
             Self::AutoIndent => "Auto Indent",
+            Self::RainbowBrackets => "Rainbow Brackets",
         }
     }
 
@@ -40,6 +48,7 @@ impl TsFeature {
             Self::Highlight => "Highlight",
             Self::TextObject => "Textobject",
             Self::AutoIndent => "Indent",
+            Self::RainbowBrackets => "Rainbow",
         }
     }
 }
diff --git a/helix-term/src/ui/editor.rs b/helix-term/src/ui/editor.rs
index a305907b..9e30cea4 100644
--- a/helix-term/src/ui/editor.rs
+++ b/helix-term/src/ui/editor.rs
@@ -97,6 +97,11 @@ impl EditorView {
         let theme = &editor.theme;
         let config = editor.config();
 
+        let should_render_rainbow_brackets = doc
+            .language_config()
+            .and_then(|lang_config| lang_config.rainbow_brackets)
+            .unwrap_or(config.rainbow_brackets);
+
         let text_annotations = view.text_annotations(doc, Some(theme));
         let mut line_decorations: Vec<Box<dyn LineDecoration>> = Vec::new();
         let mut translated_positions: Vec<TranslatedPosition> = Vec::new();
@@ -128,6 +133,12 @@ impl EditorView {
 
         let mut highlights =
             Self::doc_syntax_highlights(doc, view.offset.anchor, inner.height, theme);
+        if should_render_rainbow_brackets {
+            highlights = Box::new(syntax::merge(
+                highlights,
+                Self::doc_rainbow_highlights(doc, view.offset.anchor, inner.height, theme),
+            ));
+        }
         let overlay_highlights = Self::overlay_syntax_highlights(
             doc,
             view.offset.anchor,
@@ -336,6 +347,48 @@ impl EditorView {
         }
     }
 
+    pub fn doc_rainbow_highlights(
+        doc: &Document,
+        anchor: usize,
+        height: u16,
+        theme: &Theme,
+    ) -> Vec<(usize, std::ops::Range<usize>)> {
+        let syntax = match doc.syntax() {
+            Some(syntax) => syntax,
+            None => return Vec::new(),
+        };
+
+        let text = doc.text().slice(..);
+        let row = text.char_to_line(anchor.min(text.len_chars()));
+
+        // calculate viewport byte ranges
+        let last_line = doc.text().len_lines().saturating_sub(1);
+        let last_visible_line = (row + height as usize).saturating_sub(1).min(last_line);
+        let visible_start = text.line_to_byte(row.min(last_line));
+        let visible_end = text.line_to_byte(last_visible_line + 1);
+
+        // The calculation for the current nesting level for rainbow highlights
+        // depends on where we start the iterator from. For accuracy, we start
+        // the iterator further back than the viewport: at the start of the containing
+        // non-root syntax-tree node. Any spans that are off-screen are truncated when
+        // the spans are merged via [syntax::merge].
+        let syntax_node_start =
+            syntax::child_for_byte_range(syntax.tree().root_node(), visible_start..visible_start)
+                .map_or(visible_start, |node| node.byte_range().start);
+        let syntax_node_range = syntax_node_start..visible_end;
+
+        let mut spans = syntax.rainbow_spans(text, Some(syntax_node_range), theme.rainbow_length());
+
+        for (_highlight, range) in spans.iter_mut() {
+            let start = text.byte_to_char(ensure_grapheme_boundary_next_byte(text, range.start));
+            let end = text.byte_to_char(ensure_grapheme_boundary_next_byte(text, range.end));
+
+            *range = start..end;
+        }
+
+        spans
+    }
+
     /// Get highlight spans for document diagnostics
     pub fn doc_diagnostics_highlights(
         doc: &Document,