Add rainbow highlights on top of tree-house bindings

Author: the-mikedavis

helix-core/src/syntax.rs

@@ -13,14 +13,18 @@ use std::{
 use anyhow::{Context, Result};
 use arc_swap::{ArcSwap, Guard};
 use config::{Configuration, FileType, LanguageConfiguration, LanguageServerConfiguration};
+use foldhash::HashSet;
 use helix_loader::grammar::get_language;
 use helix_stdx::rope::RopeSliceExt as _;
 use once_cell::sync::OnceCell;
 use ropey::RopeSlice;
 use tree_house::{
     highlighter,
-    query_iter::QueryIter,
-    tree_sitter::{Grammar, InactiveQueryCursor, InputEdit, Node, Query, RopeInput, Tree},
+    query_iter::{QueryIter, QueryIterEvent},
+    tree_sitter::{
+        query::{InvalidPredicateError, UserPredicate},
+        Capture, Grammar, InactiveQueryCursor, InputEdit, Node, Pattern, Query, RopeInput, Tree,
+    },
     Error, InjectionLanguageMarker, LanguageConfig as SyntaxConfig, Layer,
 };
 
@@ -37,6 +41,7 @@ pub struct LanguageData {
     syntax: OnceCell<Option<SyntaxConfig>>,
     indent_query: OnceCell<Option<IndentQuery>>,
     textobject_query: OnceCell<Option<TextObjectQuery>>,
+    rainbow_query: OnceCell<Option<RainbowQuery>>,
 }
 
 impl LanguageData {
@@ -46,6 +51,7 @@ impl LanguageData {
             syntax: OnceCell::new(),
             indent_query: OnceCell::new(),
             textobject_query: OnceCell::new(),
+            rainbow_query: OnceCell::new(),
         }
     }
 
@@ -154,6 +160,36 @@ impl LanguageData {
             .as_ref()
     }
 
+    /// Compiles the rainbows.scm query for a language.
+    /// This function should only be used by this module or the xtask crate.
+    pub fn compile_rainbow_query(
+        grammar: Grammar,
+        config: &LanguageConfiguration,
+    ) -> Result<Option<RainbowQuery>> {
+        let name = &config.language_id;
+        let text = read_query(name, "rainbows.scm");
+        if text.is_empty() {
+            return Ok(None);
+        }
+        let rainbow_query = RainbowQuery::new(grammar, &text)
+            .with_context(|| format!("Failed to compile rainbows.scm query for '{name}'"))?;
+        Ok(Some(rainbow_query))
+    }
+
+    fn rainbow_query(&self, loader: &Loader) -> Option<&RainbowQuery> {
+        self.rainbow_query
+            .get_or_init(|| {
+                let grammar = self.syntax_config(loader)?.grammar;
+                Self::compile_rainbow_query(grammar, &self.config)
+                    .map_err(|err| {
+                        log::error!("{err}");
+                    })
+                    .ok()
+                    .flatten()
+            })
+            .as_ref()
+    }
+
     fn reconfigure(&self, scopes: &[String]) {
         if let Some(Some(config)) = self.syntax.get() {
             reconfigure_highlights(config, scopes);
@@ -324,6 +360,10 @@ impl Loader {
         self.language(lang).textobject_query(self)
     }
 
+    fn rainbow_query(&self, lang: Language) -> Option<&RainbowQuery> {
+        self.language(lang).rainbow_query(self)
+    }
+
     pub fn language_server_configs(&self) -> &HashMap<String, LanguageServerConfiguration> {
         &self.language_server_configs
     }
@@ -496,6 +536,79 @@ impl Syntax {
     {
         QueryIter::new(&self.inner, source, loader, range)
     }
+
+    pub fn rainbow_highlights(
+        &self,
+        source: RopeSlice,
+        rainbow_length: usize,
+        loader: &Loader,
+        range: impl RangeBounds<u32>,
+    ) -> OverlayHighlights {
+        struct RainbowScope<'tree> {
+            end: u32,
+            node: Option<Node<'tree>>,
+            highlight: Highlight,
+        }
+
+        let mut scope_stack = Vec::<RainbowScope>::new();
+        let mut highlights = Vec::new();
+        let mut query_iter = self.query_iter::<_, (), _>(
+            source,
+            |lang| loader.rainbow_query(lang).map(|q| &q.query),
+            range,
+        );
+
+        while let Some(event) = query_iter.next() {
+            let QueryIterEvent::Match(mat) = event else {
+                continue;
+            };
+
+            let rainbow_query = loader
+                .rainbow_query(query_iter.current_language())
+                .expect("language must have a rainbow query to emit matches");
+
+            let byte_range = mat.node.byte_range();
+            // Pop any scopes that end before this capture begins.
+            while scope_stack
+                .last()
+                .is_some_and(|scope| byte_range.start >= scope.end)
+            {
+                scope_stack.pop();
+            }
+
+            let capture = Some(mat.capture);
+            if capture == rainbow_query.scope_capture {
+                scope_stack.push(RainbowScope {
+                    end: byte_range.end,
+                    node: if rainbow_query
+                        .include_children_patterns
+                        .contains(&mat.pattern)
+                    {
+                        None
+                    } else {
+                        Some(mat.node.clone())
+                    },
+                    highlight: Highlight::new((scope_stack.len() % rainbow_length) as u32),
+                });
+            } else if capture == rainbow_query.bracket_capture {
+                if let Some(scope) = scope_stack.last() {
+                    if !scope
+                        .node
+                        .as_ref()
+                        .is_some_and(|node| mat.node.parent().as_ref() != Some(node))
+                    {
+                        let start = source
+                            .byte_to_char(source.floor_char_boundary(byte_range.start as usize));
+                        let end =
+                            source.byte_to_char(source.ceil_char_boundary(byte_range.end as usize));
+                        highlights.push((scope.highlight, start..end));
+                    }
+                }
+            }
+        }
+
+        OverlayHighlights::Heterogenous { highlights }
+    }
 }
 
 pub type Highlighter<'a> = highlighter::Highlighter<'a, 'a, Loader>;
@@ -939,6 +1052,57 @@ fn pretty_print_tree_impl<W: fmt::Write>(
     Ok(())
 }
 
+/// Finds the child of `node` which contains the given byte range.
+
+pub fn child_for_byte_range<'a>(node: &Node<'a>, range: ops::Range<u32>) -> Option<Node<'a>> {
+    for child in node.children() {
+        let child_range = child.byte_range();
+
+        if range.start >= child_range.start && range.end <= child_range.end {
+            return Some(child);
+        }
+    }
+
+    None
+}
+
+#[derive(Debug)]
+pub struct RainbowQuery {
+    query: Query,
+    include_children_patterns: HashSet<Pattern>,
+    scope_capture: Option<Capture>,
+    bracket_capture: Option<Capture>,
+}
+
+impl RainbowQuery {
+    fn new(grammar: Grammar, source: &str) -> Result<Self, tree_sitter::query::ParseError> {
+        let mut include_children_patterns = HashSet::default();
+
+        let query = Query::new(grammar, source, |pattern, predicate| match predicate {
+            UserPredicate::SetProperty {
+                key: "rainbow.include-children",
+                val,
+            } => {
+                if val.is_some() {
+                    return Err(
+                        "property 'rainbow.include-children' does not take an argument".into(),
+                    );
+                }
+                include_children_patterns.insert(pattern);
+                Ok(())
+            }
+            _ => Err(InvalidPredicateError::unknown(predicate)),
+        })?;
+
+        Ok(Self {
+            include_children_patterns,
+            scope_capture: query.get_capture("rainbow.scope"),
+            bracket_capture: query.get_capture("rainbow.bracket"),
+            query,
+        })
+    }
+}
+
 #[cfg(test)]
 mod test {
     use once_cell::sync::Lazy;

helix-core/src/syntax/config.rs

@@ -98,6 +98,8 @@ pub struct LanguageConfiguration {
     pub workspace_lsp_roots: Option<Vec<PathBuf>>,
     #[serde(default)]
     pub persistent_diagnostic_sources: Vec<String>,
+    /// Overrides the `editor.rainbow-brackets` config key for the language.
+    pub rainbow_brackets: Option<bool>,
 }
 
 impl LanguageConfiguration {

helix-term/src/ui/editor.rs

@@ -127,6 +127,18 @@ impl EditorView {
             &text_annotations,
         ));
 
+        if doc
+            .language_config()
+            .and_then(|config| config.rainbow_brackets)
+            .unwrap_or(config.rainbow_brackets)
+        {
+            if let Some(overlay) =
+                Self::doc_rainbow_highlights(doc, view_offset.anchor, inner.height, theme, &loader)
+            {
+                overlays.push(overlay);
+            }
+        }
+
         Self::doc_diagnostics_highlights_into(doc, theme, &mut overlays);
 
         if is_focused {
@@ -304,6 +316,27 @@ impl EditorView {
         text_annotations.collect_overlay_highlights(range)
     }
 
+    pub fn doc_rainbow_highlights(
+        doc: &Document,
+        anchor: usize,
+        height: u16,
+        theme: &Theme,
+        loader: &syntax::Loader,
+    ) -> Option<OverlayHighlights> {
+        let syntax = doc.syntax()?;
+        let text = doc.text().slice(..);
+        let row = text.char_to_line(anchor.min(text.len_chars()));
+        let visible_range = Self::viewport_byte_range(text, row, height);
+        let start = syntax::child_for_byte_range(
+            &syntax.tree().root_node(),
+            visible_range.start as u32..visible_range.end as u32,
+        )
+        .map_or(visible_range.start as u32, |node| node.start_byte());
+        let range = start..visible_range.end as u32;
+
+        Some(syntax.rainbow_highlights(text, theme.rainbow_length(), loader, range))
+    }
+
     /// Get highlight spans for document diagnostics
     pub fn doc_diagnostics_highlights_into(
         doc: &Document,

helix-view/src/editor.rs

@@ -373,6 +373,8 @@ pub struct Config {
     /// Whether to read settings from [EditorConfig](https://editorconfig.org) files. Defaults to
     /// `true`.
     pub editor_config: bool,
+    /// Whether to render rainbow colors for matching brackets. Defaults to `false`.
+    pub rainbow_brackets: bool,
 }
 
 #[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Eq, PartialOrd, Ord)]
@@ -1020,6 +1022,7 @@ impl Default for Config {
             end_of_line_diagnostics: DiagnosticFilter::Disable,
             clipboard_provider: ClipboardProvider::default(),
             editor_config: true,
+            rainbow_brackets: false,
         }
     }
 }