Skip to main content

koda_cli/
highlight.rs

1//! Syntax highlighting for code blocks using syntect.
2//!
3//! Provides terminal-colored syntax highlighting for code in
4//! fenced markdown code blocks. Uses the same engine as `bat`.
5
6use once_cell::sync::Lazy;
7use syntect::easy::HighlightLines;
8use syntect::highlighting::ThemeSet;
9use syntect::parsing::SyntaxSet;
10#[cfg(test)]
11use syntect::util::as_24_bit_terminal_escaped;
12
13/// Lazily loaded syntax definitions and theme.
14static SYNTAX_SET: Lazy<SyntaxSet> = Lazy::new(SyntaxSet::load_defaults_newlines);
15static THEME_SET: Lazy<ThemeSet> = Lazy::new(ThemeSet::load_defaults);
16
17/// A syntax highlighter for a specific language.
18///
19/// Stores a reference to the static `SyntaxReference` and creates a fresh
20/// `HighlightLines` on demand — no unsafe code needed.
21pub struct CodeHighlighter {
22    /// Persistent parse state for stateful (cross-line) highlighting.
23    state: Option<HighlightLines<'static>>,
24}
25
26impl CodeHighlighter {
27    /// Create a highlighter for the given language hint (e.g., "rust", "python").
28    ///
29    /// Maintains parse state across calls to `highlight_spans_stateful()`
30    /// so multiline strings, comments, and heredocs highlight correctly.
31    /// Use `highlight_spans()` for one-off single-line highlighting.
32    pub fn new(lang: &str) -> Self {
33        let syntax = SYNTAX_SET
34            .find_syntax_by_token(lang)
35            .or_else(|| SYNTAX_SET.find_syntax_by_extension(lang));
36
37        let state = syntax.map(|syn| {
38            let theme = &THEME_SET.themes["base16-ocean.dark"];
39            HighlightLines::new(syn, theme)
40        });
41
42        Self { state }
43    }
44
45    /// Highlight a single line of code, returning ANSI-colored output.
46    ///
47    /// Stateful — parse state carries across calls.
48    #[cfg(test)]
49    pub fn highlight_line(&mut self, line: &str) -> String {
50        match self.state.as_mut() {
51            Some(h) => {
52                let ranges = h.highlight_line(line, &SYNTAX_SET).unwrap_or_default();
53                let escaped = as_24_bit_terminal_escaped(&ranges[..], false);
54                format!("{escaped}\x1b[0m")
55            }
56            None => line.to_string(),
57        }
58    }
59
60    /// Highlight a line and return ratatui `Span`s with foreground colors.
61    ///
62    /// **Stateful** — parse state carries across calls, so multiline
63    /// strings/comments highlight correctly. Call lines in order.
64    ///
65    /// No background is set — the caller controls backgrounds for diff rendering.
66    pub fn highlight_spans(&mut self, line: &str) -> Vec<ratatui::text::Span<'static>> {
67        use ratatui::style::{Color, Style as RStyle};
68        use ratatui::text::Span;
69
70        match self.state.as_mut() {
71            Some(h) => {
72                let ranges = h.highlight_line(line, &SYNTAX_SET).unwrap_or_default();
73                ranges
74                    .into_iter()
75                    .map(|(style, text)| {
76                        let fg =
77                            Color::Rgb(style.foreground.r, style.foreground.g, style.foreground.b);
78                        Span::styled(text.to_string(), RStyle::default().fg(fg))
79                    })
80                    .collect()
81            }
82            None => vec![Span::raw(line.to_string())],
83        }
84    }
85}
86
87/// Pre-highlight an entire file, returning styled spans per line.
88///
89/// Maintains syntect parse state across lines for correct multiline
90/// string / comment / heredoc highlighting. Used by the diff renderer
91/// to look up pre-computed highlights by line number.
92pub fn pre_highlight(content: &str, ext: &str) -> Vec<Vec<ratatui::text::Span<'static>>> {
93    let mut hl = CodeHighlighter::new(ext);
94    content
95        .lines()
96        .map(|line| hl.highlight_spans(line))
97        .collect()
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn test_known_language_highlights() {
106        let mut h = CodeHighlighter::new("rust");
107        let result = h.highlight_line("fn main() {}");
108        // Should contain ANSI escape codes
109        assert!(result.contains("\x1b["));
110        assert!(result.contains("fn"));
111    }
112
113    #[test]
114    fn test_unknown_language_passthrough() {
115        let mut h = CodeHighlighter::new("nonexistent_lang_xyz");
116        let result = h.highlight_line("hello world");
117        assert_eq!(result, "hello world");
118    }
119
120    #[test]
121    fn test_python_highlights() {
122        let mut h = CodeHighlighter::new("python");
123        let result = h.highlight_line("def hello():");
124        assert!(result.contains("\x1b["));
125    }
126
127    #[test]
128    fn test_extension_lookup() {
129        // "rs" should find Rust syntax
130        let mut h = CodeHighlighter::new("rs");
131        let result = h.highlight_line("let x = 42;");
132        assert!(result.contains("\x1b["));
133    }
134}