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, 113 insertions, 6 deletions

diff --git a/helix-view/src/editor.rs b/helix-view/src/editor.rs
index f285aa99..fbfed656 100644
--- a/helix-view/src/editor.rs
+++ b/helix-view/src/editor.rs
@@ -308,6 +308,8 @@ pub struct Config {
     pub explorer: ExplorerConfig,
     /// Whether to color modes with different colors. Defaults to `false`.
     pub color_modes: bool,
+    /// Whether to render rainbow highlights. Defaults to `false`.
+    pub rainbow_brackets: bool,
     pub soft_wrap: SoftWrap,
     /// Workspace specific lsp ceiling dirs
     pub workspace_lsp_roots: Vec<PathBuf>,
@@ -863,6 +865,7 @@ impl Default for Config {
             indent_guides: IndentGuidesConfig::default(),
             explorer: ExplorerConfig::default(),
             color_modes: false,
+            rainbow_brackets: false,
             soft_wrap: SoftWrap {
                 enable: Some(false),
                 ..SoftWrap::default()
@@ -1203,8 +1206,7 @@ impl Editor {
             return;
         }
 
-        let scopes = theme.scopes();
-        self.syn_loader.set_scopes(scopes.to_vec());
+        self.syn_loader.set_scopes(theme.scopes().to_vec());
 
         match preview {
             ThemeAction::Preview => {
diff --git a/helix-view/src/theme.rs b/helix-view/src/theme.rs
index 4acc5664..16da247f 100644
--- a/helix-view/src/theme.rs
+++ b/helix-view/src/theme.rs
@@ -216,17 +216,19 @@ pub struct Theme {
     // tree-sitter highlight styles are stored in a Vec to optimize lookups
     scopes: Vec<String>,
     highlights: Vec<Style>,
+    rainbow_length: usize,
 }
 
 impl From<Value> for Theme {
     fn from(value: Value) -> Self {
         if let Value::Table(table) = value {
-            let (styles, scopes, highlights) = build_theme_values(table);
+            let (styles, scopes, highlights, rainbow_length) = build_theme_values(table);
 
             Self {
                 styles,
                 scopes,
                 highlights,
+                rainbow_length,
                 ..Default::default()
             }
         } else {
@@ -243,12 +245,13 @@ impl<'de> Deserialize<'de> for Theme {
     {
         let values = Map::<String, Value>::deserialize(deserializer)?;
 
-        let (styles, scopes, highlights) = build_theme_values(values);
+        let (styles, scopes, highlights, rainbow_length) = build_theme_values(values);
 
         Ok(Self {
             styles,
             scopes,
             highlights,
+            rainbow_length,
             ..Default::default()
         })
     }
@@ -256,10 +259,11 @@ impl<'de> Deserialize<'de> for Theme {
 
 fn build_theme_values(
     mut values: Map<String, Value>,
-) -> (HashMap<String, Style>, Vec<String>, Vec<Style>) {
+) -> (HashMap<String, Style>, Vec<String>, Vec<Style>, usize) {
     let mut styles = HashMap::new();
     let mut scopes = Vec::new();
     let mut highlights = Vec::new();
+    let mut rainbow_length = 0;
 
     // TODO: alert user of parsing failures in editor
     let palette = values
@@ -276,6 +280,27 @@ fn build_theme_values(
     styles.reserve(values.len());
     scopes.reserve(values.len());
     highlights.reserve(values.len());
+
+    for (i, style) in values
+        .remove("rainbow")
+        .and_then(|value| match palette.parse_style_array(value) {
+            Ok(styles) => Some(styles),
+            Err(err) => {
+                warn!("{}", err);
+                None
+            }
+        })
+        .unwrap_or_else(default_rainbow)
+        .iter()
+        .enumerate()
+    {
+        let name = format!("rainbow.{}", i);
+        styles.insert(name.clone(), *style);
+        scopes.push(name);
+        highlights.push(*style);
+        rainbow_length += 1;
+    }
+
     for (name, style_value) in values {
         let mut style = Style::default();
         if let Err(err) = palette.parse_style(&mut style, style_value) {
@@ -288,7 +313,7 @@ fn build_theme_values(
         highlights.push(style);
     }
 
-    (styles, scopes, highlights)
+    (styles, scopes, highlights, rainbow_length)
 }
 
 impl Theme {
@@ -354,6 +379,21 @@ impl Theme {
                 .all(|color| !matches!(color, Some(Color::Rgb(..))))
         })
     }
+
+    pub fn rainbow_length(&self) -> usize {
+        self.rainbow_length
+    }
+}
+
+fn default_rainbow() -> Vec<Style> {
+    vec![
+        Style::default().fg(Color::Red),
+        Style::default().fg(Color::Yellow),
+        Style::default().fg(Color::Green),
+        Style::default().fg(Color::Blue),
+        Style::default().fg(Color::Cyan),
+        Style::default().fg(Color::Magenta),
+    ]
 }
 
 struct ThemePalette {
@@ -500,6 +540,24 @@ impl ThemePalette {
         }
         Ok(())
     }
+
+    /// Parses a TOML array into a [`Vec`] of [`Style`]. If the value cannot be
+    /// parsed as an array or if any style in the array cannot be parsed then an
+    /// error is returned.
+    pub fn parse_style_array(&self, value: Value) -> Result<Vec<Style>, String> {
+        let mut styles = Vec::new();
+
+        for v in value
+            .as_array()
+            .ok_or_else(|| format!("Theme: could not parse value as an array: '{}'", value))?
+        {
+            let mut style = Style::default();
+            self.parse_style(&mut style, v.clone())?;
+            styles.push(style);
+        }
+
+        Ok(styles)
+    }
 }
 
 impl TryFrom<Value> for ThemePalette {
@@ -574,4 +632,51 @@ mod tests {
                 .add_modifier(Modifier::BOLD)
         );
     }
+
+    #[test]
+    fn test_parse_valid_style_array() {
+        let theme = toml::toml! {
+            rainbow = ["#ff0000", "#ffa500", "#fff000", { fg = "#00ff00", modifiers = ["bold"] }]
+        };
+
+        let palette = ThemePalette::default();
+
+        let rainbow = theme.get("rainbow").unwrap();
+        let parse_result = palette.parse_style_array(rainbow.clone());
+
+        assert_eq!(
+            Ok(vec![
+                Style::default().fg(Color::Rgb(255, 0, 0)),
+                Style::default().fg(Color::Rgb(255, 165, 0)),
+                Style::default().fg(Color::Rgb(255, 240, 0)),
+                Style::default()
+                    .fg(Color::Rgb(0, 255, 0))
+                    .add_modifier(Modifier::BOLD),
+            ]),
+            parse_result
+        )
+    }
+
+    #[test]
+    fn test_parse_invalid_style_array() {
+        let palette = ThemePalette::default();
+
+        let theme = toml::toml! { invalid_hex_code = ["#f00"] };
+        let invalid_hex_code = theme.get("invalid_hex_code").unwrap();
+        let parse_result = palette.parse_style_array(invalid_hex_code.clone());
+
+        assert_eq!(
+            Err("Theme: malformed hexcode: #f00".to_string()),
+            parse_result
+        );
+
+        let theme = toml::toml! { not_an_array = { red = "#ff0000" } };
+        let not_an_array = theme.get("not_an_array").unwrap();
+        let parse_result = palette.parse_style_array(not_an_array.clone());
+
+        assert_eq!(
+            Err("Theme: could not parse value as an array: '{ red = \"#ff0000\" }'".to_string()),
+            parse_result
+        )
+    }
 }