diffler_core/

highlight.rs

//! Tree-sitter highlighting sliced into per-line styled ranges. We highlight
//! whole files so multi-line strings stay correct across hunk boundaries.

use std::ops::Range;

use tree_sitter_highlight::{HighlightEvent, Highlighter as TsHighlighter};

use crate::syntax::registry::LangEntry;
use crate::syntax::{HIGHLIGHT_NAMES, LanguageRegistry};

pub struct Highlighter {
    registry: &'static LanguageRegistry,
    theme: SyntaxTheme,
    /// The reader's `glob = language` rules, tried in order before the
    /// registry's own inference.
    rules: Vec<(String, String)>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub enum SyntaxTheme {
    #[default]
    OneHalfDark,
    OneHalfLight,
    Dracula,
    CatppuccinMocha,
    TokyoNight,
    GruvboxDark,
    Nord,
    RosePine,
    Kanagawa,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StyledRange {
    pub range: Range<usize>,
    pub fg: (u8, u8, u8),
    pub bold: bool,
    pub italic: bool,
}

impl Default for Highlighter {
    fn default() -> Self {
        Self::new(SyntaxTheme::default())
    }
}

impl Highlighter {
    pub fn new(syntax: SyntaxTheme) -> Self {
        Self {
            registry: &crate::syntax::registry::REGISTRY,
            theme: syntax,
            rules: Vec::new(),
        }
    }

    #[must_use]
    pub fn with_rules(mut self, rules: Vec<(String, String)>) -> Self {
        self.rules = rules;
        self
    }

    /// The first matching rule that names a bundled grammar, else the
    /// registry's inference.
    pub fn language(&self, path: &str, content: &str) -> Option<&'static LangEntry> {
        self.rules
            .iter()
            .filter(|(glob, _)| crate::classify::glob_match(glob, path))
            .find_map(|(_, name)| self.registry.by_name(name))
            .or_else(|| self.registry.for_file(path, content))
    }

    /// One `Vec<StyledRange>` per line, empty for an unknown language.
    pub fn highlight(&self, path: &str, content: &str) -> Vec<Vec<StyledRange>> {
        self.highlight_entry(self.language(path, content), content)
    }

    /// Highlight `content` as a markdown fence token.
    pub fn highlight_lang(&self, token: &str, content: &str) -> Vec<Vec<StyledRange>> {
        self.highlight_entry(self.registry.for_token(token), content)
    }

    fn highlight_entry(
        &self,
        entry: Option<&crate::syntax::registry::LangEntry>,
        content: &str,
    ) -> Vec<Vec<StyledRange>> {
        let bounds = crate::syntax::line_bounds(content);
        let mut out: Vec<Vec<StyledRange>> = vec![Vec::new(); bounds.len()];

        if content.len() > crate::syntax::MAX_PARSE_BYTES {
            return out;
        }
        let Some(entry) = entry else {
            return out;
        };
        let Some(config) = entry.config() else {
            return out;
        };

        let mut ts = TsHighlighter::new();
        let registry = self.registry;
        let Ok(events) = ts.highlight(config, content.as_bytes(), None, move |lang| {
            registry.config_for_injection(lang)
        }) else {
            return out;
        };

        let starts: Vec<usize> = bounds.iter().map(|&(s, _)| s).collect();
        let mut stack: Vec<usize> = Vec::new();
        for event in events {
            let Ok(event) = event else {
                return out;
            };
            match event {
                HighlightEvent::HighlightStart(h) => stack.push(h.0),
                HighlightEvent::HighlightEnd => {
                    stack.pop();
                }
                HighlightEvent::Source { start, end } => {
                    if let Some(&name_idx) = stack.last()
                        && let Some(name) = HIGHLIGHT_NAMES.get(name_idx)
                        && let Some(style) = self.theme.style(name)
                    {
                        push_styled(&mut out, &bounds, &starts, &(start..end), &style);
                    }
                }
            }
        }

        if entry.name == "markdown" {
            for (range, name) in self.registry.markdown_inline_spans(content) {
                if let Some(style) = self.theme.style(name) {
                    push_styled(&mut out, &bounds, &starts, &range, &style);
                }
            }
        }
        out
    }

    pub fn scope_index(&self, path: &str, content: &str) -> crate::syntax::ScopeIndex {
        LanguageRegistry::scope_index_as(self.language(path, content), content)
    }

    /// See [`LanguageRegistry::syntactic_emphasis`].
    pub fn syntactic_emphasis(
        &self,
        file: &mut crate::model::FileDiff,
        mark_reformat_only: bool,
    ) -> bool {
        let entry = self.language(&file.path, file.new_text.as_deref().unwrap_or_default());
        LanguageRegistry::syntactic_emphasis(entry, file, mark_reformat_only)
    }
}

struct StyleSpec {
    fg: (u8, u8, u8),
    bold: bool,
    italic: bool,
}

fn push_styled(
    out: &mut [Vec<StyledRange>],
    bounds: &[(usize, usize)],
    starts: &[usize],
    range: &Range<usize>,
    style: &StyleSpec,
) {
    crate::syntax::split_range_by_line(bounds, starts, range, |li, r| {
        if let Some(line) = out.get_mut(li) {
            line.push(StyledRange {
                range: r,
                fg: style.fg,
                bold: style.bold,
                italic: style.italic,
            });
        }
    });
}

/// Palette category and face for markdown `text.*` captures, so themes need
/// no markdown colors of their own.
fn markdown_face(name: &str) -> Option<(&'static str, bool, bool)> {
    let face = match name {
        "text.title" => ("function", true, false),
        "text.strong" => ("variable", true, false),
        "text.emphasis" => ("variable", false, true),
        "text.literal" => ("string", false, false),
        "text.uri" | "text.reference" => ("property", false, false),
        _ => return None,
    };
    Some(face)
}

/// Fold a grammar's own capture category onto the palette's. We apply it
/// before picking the face so a folded comment stays italic.
fn palette_category(category: &str) -> &str {
    match category {
        "boolean" => "constant",
        "conditional" | "storageclass" => "keyword",
        "field" => "property",
        "parameter" => "variable",
        // SQL tags comments `@comment @spell`, and the last capture wins
        "spell" => "comment",
        other => other,
    }
}

impl SyntaxTheme {
    /// Matched by the capture's leading category. `None` leaves the default fg.
    fn style(self, name: &str) -> Option<StyleSpec> {
        if let Some((category, bold, italic)) = markdown_face(name) {
            return Some(StyleSpec {
                fg: self.color(category)?,
                bold,
                italic,
            });
        }
        let category = palette_category(name.split('.').next().unwrap_or(name));
        let italic = category == "comment";
        let fg = self.color(category)?;
        Some(StyleSpec {
            fg,
            bold: false,
            italic,
        })
    }

    #[allow(clippy::too_many_lines)]
    fn color(self, category: &str) -> Option<(u8, u8, u8)> {
        let c = match self {
            SyntaxTheme::OneHalfDark => match category {
                "keyword" | "label" => (198, 120, 221),
                "function" => (97, 175, 239),
                "type" | "constructor" => (229, 192, 123),
                "string" => (152, 195, 121),
                // brighter than One Dark's default so comments stay legible on
                // the diff backgrounds
                "comment" => (126, 134, 145),
                "constant" | "number" | "attribute" => (209, 154, 102),
                "operator" | "escape" => (86, 182, 194),
                "property" | "tag" => (224, 108, 117),
                "variable" | "punctuation" => (171, 178, 191),
                _ => return None,
            },
            SyntaxTheme::OneHalfLight => match category {
                "keyword" | "label" => (166, 38, 164),
                "function" => (64, 120, 242),
                "type" | "constructor" => (193, 132, 1),
                "string" => (80, 161, 79),
                "comment" => (160, 161, 167),
                "constant" | "number" | "attribute" => (152, 104, 1),
                "operator" | "escape" => (1, 132, 188),
                "property" | "tag" => (228, 86, 73),
                "variable" | "punctuation" => (56, 58, 66),
                _ => return None,
            },
            SyntaxTheme::Dracula => match category {
                "keyword" | "label" | "tag" | "operator" => (255, 121, 198),
                "function" | "property" => (80, 250, 123),
                "type" | "constructor" => (139, 233, 253),
                "string" => (241, 250, 140),
                "comment" => (98, 114, 164),
                "constant" | "number" => (189, 147, 249),
                "escape" | "attribute" => (255, 184, 108),
                "variable" | "punctuation" => (248, 248, 242),
                _ => return None,
            },
            SyntaxTheme::CatppuccinMocha => match category {
                "keyword" | "label" => (203, 166, 247),
                "function" => (137, 180, 250),
                "type" | "constructor" => (249, 226, 175),
                "string" => (166, 227, 161),
                "comment" => (127, 132, 156),
                "constant" | "number" | "attribute" => (250, 179, 135),
                "operator" | "escape" => (137, 220, 235),
                "property" | "tag" => (243, 139, 168),
                "variable" | "punctuation" => (205, 214, 244),
                _ => return None,
            },
            SyntaxTheme::TokyoNight => match category {
                "keyword" | "label" => (187, 154, 247),
                "function" => (122, 162, 247),
                "type" | "constructor" => (42, 195, 222),
                "string" => (158, 206, 106),
                "comment" => (99, 109, 150),
                "constant" | "number" | "attribute" => (255, 158, 100),
                "operator" | "escape" => (137, 221, 255),
                "property" | "tag" => (247, 118, 142),
                "variable" | "punctuation" => (192, 202, 245),
                _ => return None,
            },
            SyntaxTheme::GruvboxDark => match category {
                "keyword" | "label" => (251, 73, 52),
                "function" => (184, 187, 38),
                "type" | "constructor" => (250, 189, 47),
                "string" => (142, 192, 124),
                "comment" => (146, 131, 116),
                "constant" | "number" => (211, 134, 155),
                "operator" | "escape" | "attribute" => (254, 128, 25),
                "property" | "tag" => (131, 165, 152),
                "variable" | "punctuation" => (235, 219, 178),
                _ => return None,
            },
            SyntaxTheme::Nord => match category {
                "keyword" | "label" | "operator" | "escape" => (129, 161, 193),
                "function" => (136, 192, 208),
                "type" | "constructor" | "property" | "tag" => (143, 188, 187),
                "string" => (163, 190, 140),
                "comment" => (123, 136, 161),
                "constant" | "number" | "attribute" => (180, 142, 173),
                "variable" | "punctuation" => (216, 222, 233),
                _ => return None,
            },
            SyntaxTheme::RosePine => match category {
                "keyword" | "label" | "operator" | "escape" => (49, 116, 143),
                "function" => (235, 188, 186),
                "type" | "constructor" | "property" | "tag" => (156, 207, 216),
                "string" => (246, 193, 119),
                "comment" => (129, 124, 153),
                "constant" | "number" | "attribute" => (196, 167, 231),
                "variable" | "punctuation" => (224, 222, 244),
                _ => return None,
            },
            SyntaxTheme::Kanagawa => match category {
                "keyword" | "label" => (149, 127, 184),
                "function" => (126, 156, 216),
                "type" | "constructor" => (122, 168, 159),
                "string" => (152, 187, 108),
                "comment" => (144, 140, 128),
                "constant" | "number" | "attribute" => (210, 126, 153),
                "operator" | "escape" => (127, 180, 202),
                "property" | "tag" => (106, 149, 137),
                "variable" | "punctuation" => (220, 215, 186),
                _ => return None,
            },
        };
        Some(c)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn the_readers_rules_outrank_what_the_registry_infers() {
        let rules = |rules: &[(&str, &str)]| {
            Highlighter::default().with_rules(
                rules
                    .iter()
                    .map(|&(glob, lang)| (glob.to_owned(), lang.to_owned()))
                    .collect(),
            )
        };
        let name = |hl: &Highlighter, path| hl.language(path, "").map(|entry| entry.name);
        let hl = rules(&[("/main.rs", "toml"), ("*.tpl", "nope"), ("*.tpl", "yaml")]);
        assert_eq!(name(&hl, "main.rs"), Some("toml"));
        assert_eq!(
            name(&hl, "src/main.rs"),
            Some("rust"),
            "a leading / anchors"
        );
        assert_eq!(
            name(&hl, "chart/values.tpl"),
            Some("yaml"),
            "an unknown name is skipped"
        );
    }

    /// A grammar crate with a broken highlight query links fine and renders
    /// plain, so we check each one colours a sample.
    #[test]
    fn every_language_colours_a_sample() {
        let samples: &[(&str, &str)] = &[
            ("a.rs", "fn main() { let x = 1; }\n"),
            ("a.py", "def f():\n    return 1\n"),
            ("a.js", "const x = 1;\n"),
            ("a.ts", "const x: number = 1;\n"),
            ("a.tsx", "const A = () => <div />;\n"),
            ("a.go", "package main\nfunc main() {}\n"),
            ("a.c", "int main(void) { return 0; }\n"),
            ("a.cpp", "int main() { return 0; }\n"),
            ("A.java", "class A { void f() {} }\n"),
            ("a.cs", "class A { void F() {} }\n"),
            ("a.rb", "def f\n  1\nend\n"),
            ("a.php", "<?php function f() { return 1; }\n"),
            ("a.sh", "set -e\necho hi\n"),
            ("a.json", "{\"a\": 1}\n"),
            ("a.html", "<p>hi</p>\n"),
            ("a.css", "a { color: red; }\n"),
            ("a.yml", "name: CI\non: push\n"),
            ("a.sql", "SELECT id FROM users;\n"),
            ("a.md", "# Title\n\ntext\n"),
            ("a.toml", "[package]\nname = \"x\"\n"),
            (
                "main.tf",
                "resource \"aws_s3_bucket\" \"b\" {\n  bucket = var.name\n}\n",
            ),
            ("Dockerfile", "FROM alpine:3\nRUN echo hi\n"),
            ("a.mk", "all:\n\techo hi\n"),
            ("a.lua", "local function f() return 1 end\n"),
            ("a.nix", "{ pkgs }: pkgs.hello\n"),
            ("a.xml", "<root><a b=\"c\"/></root>\n"),
            ("a.swift", "func f() -> Int { return 1 }\n"),
            ("a.scala", "object A { def f = 1 }\n"),
            ("a.ex", "defmodule A do\n  def f, do: 1\nend\n"),
            ("a.zig", "pub fn main() void {}\n"),
            ("a.hs", "main :: IO ()\nmain = return ()\n"),
            ("a.dart", "void main() { var x = 1; }\n"),
            ("a.ps1", "function Get-Thing { param($x) $x }\n"),
            ("a.svelte", "<script>let x = 1;</script>\n<p>{x}</p>\n"),
        ];
        let hl = Highlighter::default();
        let plain: Vec<&str> = samples
            .iter()
            .filter(|(path, content)| hl.highlight(path, content).iter().all(Vec::is_empty))
            .map(|(path, _)| *path)
            .collect();
        assert!(plain.is_empty(), "rendered plain: {plain:?}");
    }

    #[test]
    fn a_build_tool_file_resolves_by_its_name() {
        let hl = Highlighter::default();
        for path in ["Makefile", "GNUmakefile", "Dockerfile", "Containerfile"] {
            let lines = hl.highlight(path, "FROM alpine\nall:\n\techo hi\n");
            assert!(
                lines.iter().any(|line| !line.is_empty()),
                "{path} rendered plain"
            );
        }
    }

    #[test]
    fn terraform_fences_and_extensions_reach_the_hcl_grammar() {
        let hl = Highlighter::default();
        let source = "variable \"name\" {\n  type = string\n}\n";
        for path in ["main.tf", "vars.tfvars", "config.hcl"] {
            assert!(
                hl.highlight(path, source)
                    .iter()
                    .any(|line| !line.is_empty()),
                "{path} rendered plain"
            );
        }
        for token in ["terraform", "hcl", "tf"] {
            assert!(
                hl.highlight_lang(token, source)
                    .iter()
                    .any(|line| !line.is_empty()),
                "fence {token} rendered plain"
            );
        }
    }

    #[test]
    fn python_keywords_get_distinct_color() {
        let hl = Highlighter::default();
        let lines = hl.highlight("a.py", "def f():\n    return 1\n");
        assert_eq!(lines.len(), 2);
        let colors: std::collections::HashSet<(u8, u8, u8)> =
            lines[0].iter().map(|r| r.fg).collect();
        assert!(colors.len() > 1, "expected multiple colors, got {colors:?}");
    }

    #[test]
    fn yaml_is_highlighted() {
        let hl = Highlighter::default();
        let lines = hl.highlight("ci.yml", "name: CI\non: push\njobs:\n  lint: {}\n");
        assert!(
            lines.iter().any(|line| !line.is_empty()),
            "expected styled ranges for a .yml file"
        );
    }

    #[test]
    fn sql_is_highlighted() {
        let hl = Highlighter::default();
        let lines = hl.highlight("q.sql", "SELECT id FROM users WHERE active = true;\n");
        assert!(
            lines.iter().any(|line| !line.is_empty()),
            "expected styled ranges for a .sql file"
        );
    }

    #[test]
    fn sql_line_and_block_comments_style_alike() {
        let hl = Highlighter::default();
        let lines = hl.highlight("q.sql", "-- one\n/* two */\n");
        let style = |line: &[StyledRange]| line.first().map(|r| (r.fg, r.italic));
        assert_eq!(
            style(&lines[0]),
            style(&lines[1]),
            "a -- comment styles like a block comment"
        );
        assert!(style(&lines[0]).is_some(), "comments are styled at all");
    }

    #[test]
    fn sql_numbers_are_not_styled_as_strings() {
        let hl = Highlighter::default();
        let lines = hl.highlight("q.sql", "SELECT 42, 1.5, 'txt';\n");
        let colors: Vec<(u8, u8, u8)> = lines[0].iter().map(|r| r.fg).collect();
        let string_fg = (152, 195, 121);
        let number_fg = (209, 154, 102);
        assert!(colors.contains(&number_fg), "numbers get the number color");
        assert_eq!(
            colors.iter().filter(|c| **c == string_fg).count(),
            1,
            "only the quoted literal is a string: {colors:?}"
        );
    }

    #[test]
    fn ranges_cover_within_line_bounds() {
        let hl = Highlighter::default();
        let src = "fn main() { let x = \"hi\"; }\n";
        let lines = hl.highlight("a.rs", src);
        let visible = src.trim_end();
        for r in &lines[0] {
            assert!(r.range.end <= visible.len());
            assert!(r.range.start < r.range.end);
        }
    }

    #[test]
    fn multiline_string_state_carries_across_lines() {
        let hl = Highlighter::default();
        let src = "s = \"\"\"first\nsecond\nthird\"\"\"\nx = 1\n";
        let lines = hl.highlight("a.py", src);
        let string_color = lines[0].iter().last().map(|r| r.fg).expect("line 0 styled");
        assert!(
            lines[1].iter().all(|r| r.fg == string_color),
            "inside-string line must keep string color"
        );
    }

    #[test]
    fn markdown_highlights_headings_and_inline_code() {
        let hl = Highlighter::default();
        let src = "# Title\n\nSome `code` and **bold** text.\n";
        let lines = hl.highlight("readme.md", src);
        assert!(!lines[0].is_empty(), "heading line should be styled");
        assert!(
            lines[2].iter().any(|r| r.fg == (152, 195, 121)),
            "inline `code` should get the string color"
        );
    }

    #[test]
    fn markdown_inline_code_offset_is_absolute_not_range_relative() {
        let hl = Highlighter::default();
        let src = "# A longer heading here\n\nintro line\n\nthen `code` appears.\n";
        let lines = hl.highlight("readme.md", src);
        let code_line = "then `code` appears.";
        let styled: Vec<_> = lines[4]
            .iter()
            .filter(|r| r.fg == (152, 195, 121))
            .collect();
        assert!(!styled.is_empty(), "code span should be styled on line 4");
        for r in styled {
            assert!(
                r.range.end <= code_line.len(),
                "range {:?} escapes the line (offsets not absolute)",
                r.range
            );
            assert_eq!(&code_line[r.range.clone()], "`code`");
        }
    }

    #[test]
    fn markdown_fenced_code_block_gets_language_highlight() {
        let hl = Highlighter::default();
        let src = "text\n\n```rust\nfn f() {}\n```\n";
        let lines = hl.highlight("readme.md", src);
        assert!(
            lines[3].iter().any(|r| r.fg == (198, 120, 221)),
            "fenced rust `fn` should get the keyword color"
        );
    }

    #[test]
    fn markdown_fence_tag_resolves_by_extension() {
        let hl = Highlighter::default();
        let src = "text\n\n```rs\nfn f() {}\n```\n";
        let lines = hl.highlight("readme.md", src);
        assert!(
            lines[3].iter().any(|r| r.fg == (198, 120, 221)),
            "an `rs` fence should resolve to rust via by_ext"
        );
    }

    #[test]
    fn unknown_extension_yields_plain_lines() {
        let hl = Highlighter::default();
        let lines = hl.highlight("file.zzz-unknown", "a\nb\n");
        assert_eq!(lines, vec![Vec::new(), Vec::new()]);
    }

    #[test]
    fn syntax_theme_changes_the_foreground_palette() {
        let src = "fn main() { let x = 1; }\n";
        let dark = Highlighter::new(SyntaxTheme::OneHalfDark).highlight("a.rs", src);
        let light = Highlighter::new(SyntaxTheme::OneHalfLight).highlight("a.rs", src);
        assert_ne!(dark, light, "a different syntax theme recolors the line");
    }
}