From 780f141a759f2729eeaf52ba4b1263b6d1a6cd52 Mon Sep 17 00:00:00 2001
From: Michael Davis <mcarsondavis@gmail.com>
Date: Fri, 16 Dec 2022 10:05:38 -0600
Subject: [PATCH] Overlay rainbow highlights onto syntax highlights

We call the rainbow_spans function introduced in the parent commits
over the largest node that contains the current viewport: we need to
reach far enough back in the document that we find the absolute
beginning for brackets. If we run rainbow_spans only over the current
viewport, we get a bug where the color of rainbow brackets changes as
we move the viewport.
---
 helix-core/src/syntax.rs    | 12 ++++++++++
 helix-term/src/ui/editor.rs | 48 +++++++++++++++++++++++++++++++++++--
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/helix-core/src/syntax.rs b/helix-core/src/syntax.rs
index 1a78cdef5..e80f50ff7 100644
--- a/helix-core/src/syntax.rs
+++ b/helix-core/src/syntax.rs
@@ -1672,6 +1672,18 @@ impl Syntax {
     // TODO: Folding
 }
 
+/// Finds the child of `node` which contains the given byte range `range`.
+pub fn child_for_byte_range(node: Node, range: std::ops::Range<usize>) -> Option<Node> {
+    for child in node.children(&mut node.walk()) {
+        let child_range = child.byte_range();
+        if range.start >= child_range.start && range.end <= child_range.end {
+            return Some(child);
+        }
+    }
+
+    None
+}
+
 bitflags! {
     /// Flags that track the status of a layer
     /// in the `Sytaxn::update` function
diff --git a/helix-term/src/ui/editor.rs b/helix-term/src/ui/editor.rs
index f7541fe25..56e6d40dc 100644
--- a/helix-term/src/ui/editor.rs
+++ b/helix-term/src/ui/editor.rs
@@ -92,9 +92,13 @@ impl EditorView {
         let area = view.area;
         let theme = &editor.theme;
         let config = editor.config();
-
         let view_offset = doc.view_offset(view.id);
 
+        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 decorations = DecorationManager::default();
 
@@ -120,9 +124,16 @@ impl EditorView {
             decorations.add_decoration(line_decoration);
         }
 
-        let syntax_highlights =
+        let mut syntax_highlights =
             Self::doc_syntax_highlights(doc, view_offset.anchor, inner.height, theme);
 
+        if should_render_rainbow_brackets {
+            syntax_highlights = Box::new(syntax::merge(
+                syntax_highlights,
+                Self::doc_rainbow_highlights(doc, view_offset.anchor, inner.height, theme),
+            ));
+        }
+
         let mut overlay_highlights =
             Self::empty_highlight_iter(doc, view_offset.anchor, inner.height);
         let overlay_syntax_highlights = Self::overlay_syntax_highlights(
@@ -356,6 +367,39 @@ impl EditorView {
         text_annotations.collect_overlay_highlights(range)
     }
 
+    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;
+
+        syntax.rainbow_spans(text, Some(syntax_node_range), theme.rainbow_length())
+    }
+
     /// Get highlight spans for document diagnostics
     pub fn doc_diagnostics_highlights(
         doc: &Document,