Skip to main content

render_ansi/
lib.rs

1use highlight_spans::{Grammar, HighlightError, HighlightResult, SpanHighlighter};
2use theme_engine::{Style, Theme};
3use thiserror::Error;
4
5const SGR_RESET: &str = "\x1b[0m";
6
7#[derive(Debug, Clone, Copy, Eq, PartialEq)]
8pub struct StyledSpan {
9    pub start_byte: usize,
10    pub end_byte: usize,
11    pub style: Option<Style>,
12}
13
14#[derive(Debug, Error)]
15pub enum RenderError {
16    #[error("highlighting failed: {0}")]
17    Highlight(#[from] HighlightError),
18    #[error("invalid span range {start_byte}..{end_byte} for source length {source_len}")]
19    SpanOutOfBounds {
20        start_byte: usize,
21        end_byte: usize,
22        source_len: usize,
23    },
24    #[error(
25        "spans must be sorted and non-overlapping: prev_end={prev_end}, next_start={next_start}"
26    )]
27    OverlappingSpans { prev_end: usize, next_start: usize },
28    #[error("invalid attr_id {attr_id}; attrs length is {attrs_len}")]
29    InvalidAttrId { attr_id: usize, attrs_len: usize },
30}
31
32pub fn resolve_styled_spans(
33    highlight: &HighlightResult,
34    theme: &Theme,
35) -> Result<Vec<StyledSpan>, RenderError> {
36    let mut out = Vec::with_capacity(highlight.spans.len());
37    for span in &highlight.spans {
38        let Some(attr) = highlight.attrs.get(span.attr_id) else {
39            return Err(RenderError::InvalidAttrId {
40                attr_id: span.attr_id,
41                attrs_len: highlight.attrs.len(),
42            });
43        };
44        out.push(StyledSpan {
45            start_byte: span.start_byte,
46            end_byte: span.end_byte,
47            style: theme.resolve(&attr.capture_name).copied(),
48        });
49    }
50    Ok(out)
51}
52
53pub fn render_ansi(source: &[u8], spans: &[StyledSpan]) -> Result<String, RenderError> {
54    validate_spans(source.len(), spans)?;
55
56    let mut out = String::new();
57    let mut cursor = 0usize;
58    for span in spans {
59        if cursor < span.start_byte {
60            out.push_str(&String::from_utf8_lossy(&source[cursor..span.start_byte]));
61        }
62        append_styled_segment(
63            &mut out,
64            &source[span.start_byte..span.end_byte],
65            span.style,
66        );
67        cursor = span.end_byte;
68    }
69
70    if cursor < source.len() {
71        out.push_str(&String::from_utf8_lossy(&source[cursor..]));
72    }
73
74    Ok(out)
75}
76
77pub fn render_ansi_lines(source: &[u8], spans: &[StyledSpan]) -> Result<Vec<String>, RenderError> {
78    validate_spans(source.len(), spans)?;
79
80    let line_ranges = compute_line_ranges(source);
81    let mut lines = Vec::with_capacity(line_ranges.len());
82    let mut span_cursor = 0usize;
83
84    for (line_start, line_end) in line_ranges {
85        while span_cursor < spans.len() && spans[span_cursor].end_byte <= line_start {
86            span_cursor += 1;
87        }
88
89        let mut line = String::new();
90        let mut cursor = line_start;
91        let mut i = span_cursor;
92        while i < spans.len() {
93            let span = spans[i];
94            if span.start_byte >= line_end {
95                break;
96            }
97
98            let seg_start = span.start_byte.max(line_start);
99            let seg_end = span.end_byte.min(line_end);
100            if cursor < seg_start {
101                line.push_str(&String::from_utf8_lossy(&source[cursor..seg_start]));
102            }
103            append_styled_segment(&mut line, &source[seg_start..seg_end], span.style);
104            cursor = seg_end;
105            i += 1;
106        }
107
108        if cursor < line_end {
109            line.push_str(&String::from_utf8_lossy(&source[cursor..line_end]));
110        }
111
112        lines.push(line);
113    }
114
115    Ok(lines)
116}
117
118pub fn highlight_to_ansi(
119    source: &[u8],
120    flavor: Grammar,
121    theme: &Theme,
122) -> Result<String, RenderError> {
123    let mut highlighter = SpanHighlighter::new()?;
124    highlight_to_ansi_with_highlighter(&mut highlighter, source, flavor, theme)
125}
126
127pub fn highlight_to_ansi_with_highlighter(
128    highlighter: &mut SpanHighlighter,
129    source: &[u8],
130    flavor: Grammar,
131    theme: &Theme,
132) -> Result<String, RenderError> {
133    let highlight = highlighter.highlight(source, flavor)?;
134    let styled = resolve_styled_spans(&highlight, theme)?;
135    render_ansi(source, &styled)
136}
137
138pub fn highlight_lines_to_ansi_lines<S: AsRef<str>>(
139    lines: &[S],
140    flavor: Grammar,
141    theme: &Theme,
142) -> Result<Vec<String>, RenderError> {
143    let mut highlighter = SpanHighlighter::new()?;
144    highlight_lines_to_ansi_lines_with_highlighter(&mut highlighter, lines, flavor, theme)
145}
146
147pub fn highlight_lines_to_ansi_lines_with_highlighter<S: AsRef<str>>(
148    highlighter: &mut SpanHighlighter,
149    lines: &[S],
150    flavor: Grammar,
151    theme: &Theme,
152) -> Result<Vec<String>, RenderError> {
153    let highlight = highlighter.highlight_lines(lines, flavor)?;
154    let source = lines
155        .iter()
156        .map(AsRef::as_ref)
157        .collect::<Vec<_>>()
158        .join("\n");
159    let styled = resolve_styled_spans(&highlight, theme)?;
160    render_ansi_lines(source.as_bytes(), &styled)
161}
162
163fn append_styled_segment(out: &mut String, text: &[u8], style: Option<Style>) {
164    if text.is_empty() {
165        return;
166    }
167
168    if let Some(open) = style_open_sgr(style) {
169        out.push_str(&open);
170        out.push_str(&String::from_utf8_lossy(text));
171        out.push_str(SGR_RESET);
172        return;
173    }
174
175    out.push_str(&String::from_utf8_lossy(text));
176}
177
178fn style_open_sgr(style: Option<Style>) -> Option<String> {
179    let style = style?;
180    let mut parts = Vec::new();
181    if let Some(fg) = style.fg {
182        parts.push(format!("38;2;{};{};{}", fg.r, fg.g, fg.b));
183    }
184    if let Some(bg) = style.bg {
185        parts.push(format!("48;2;{};{};{}", bg.r, bg.g, bg.b));
186    }
187    if style.bold {
188        parts.push("1".to_string());
189    }
190    if style.italic {
191        parts.push("3".to_string());
192    }
193    if style.underline {
194        parts.push("4".to_string());
195    }
196
197    if parts.is_empty() {
198        return None;
199    }
200
201    Some(format!("\x1b[{}m", parts.join(";")))
202}
203
204fn compute_line_ranges(source: &[u8]) -> Vec<(usize, usize)> {
205    let mut ranges = Vec::new();
206    let mut line_start = 0usize;
207    for (i, byte) in source.iter().enumerate() {
208        if *byte == b'\n' {
209            ranges.push((line_start, i));
210            line_start = i + 1;
211        }
212    }
213    ranges.push((line_start, source.len()));
214    ranges
215}
216
217fn validate_spans(source_len: usize, spans: &[StyledSpan]) -> Result<(), RenderError> {
218    let mut prev_end = 0usize;
219    for (i, span) in spans.iter().enumerate() {
220        if span.start_byte > span.end_byte || span.end_byte > source_len {
221            return Err(RenderError::SpanOutOfBounds {
222                start_byte: span.start_byte,
223                end_byte: span.end_byte,
224                source_len,
225            });
226        }
227        if i > 0 && span.start_byte < prev_end {
228            return Err(RenderError::OverlappingSpans {
229                prev_end,
230                next_start: span.start_byte,
231            });
232        }
233        prev_end = span.end_byte;
234    }
235    Ok(())
236}
237
238#[cfg(test)]
239mod tests {
240    use super::{
241        highlight_lines_to_ansi_lines, highlight_to_ansi, render_ansi, render_ansi_lines,
242        RenderError, StyledSpan,
243    };
244    use highlight_spans::Grammar;
245    use theme_engine::{load_theme, Rgb, Style, Theme};
246
247    #[test]
248    fn renders_basic_styled_segment() {
249        let source = b"abc";
250        let spans = [StyledSpan {
251            start_byte: 1,
252            end_byte: 2,
253            style: Some(Style {
254                fg: Some(Rgb::new(255, 0, 0)),
255                bold: true,
256                ..Style::default()
257            }),
258        }];
259        let out = render_ansi(source, &spans).expect("failed to render");
260        assert_eq!(out, "a\x1b[38;2;255;0;0;1mb\x1b[0mc");
261    }
262
263    #[test]
264    fn renders_per_line_output_for_multiline_span() {
265        let source = b"ab\ncd";
266        let spans = [StyledSpan {
267            start_byte: 1,
268            end_byte: 5,
269            style: Some(Style {
270                fg: Some(Rgb::new(1, 2, 3)),
271                ..Style::default()
272            }),
273        }];
274
275        let lines = render_ansi_lines(source, &spans).expect("failed to render lines");
276        assert_eq!(lines.len(), 2);
277        assert_eq!(lines[0], "a\x1b[38;2;1;2;3mb\x1b[0m");
278        assert_eq!(lines[1], "\x1b[38;2;1;2;3mcd\x1b[0m");
279    }
280
281    #[test]
282    fn rejects_overlapping_spans() {
283        let spans = [
284            StyledSpan {
285                start_byte: 0,
286                end_byte: 2,
287                style: None,
288            },
289            StyledSpan {
290                start_byte: 1,
291                end_byte: 3,
292                style: None,
293            },
294        ];
295        let err = render_ansi(b"abcd", &spans).expect_err("expected overlap error");
296        assert!(matches!(err, RenderError::OverlappingSpans { .. }));
297    }
298
299    #[test]
300    fn highlights_source_to_ansi() {
301        let theme = Theme::from_json_str(
302            r#"{
303  "styles": {
304    "normal": { "fg": { "r": 220, "g": 220, "b": 220 } },
305    "number": { "fg": { "r": 255, "g": 180, "b": 120 } }
306  }
307}"#,
308        )
309        .expect("theme parse failed");
310
311        let out = highlight_to_ansi(b"set x = 42", Grammar::ObjectScript, &theme)
312            .expect("highlight+render failed");
313        assert!(out.contains("42"));
314        assert!(out.contains("\x1b["));
315    }
316
317    #[test]
318    fn highlights_lines_to_ansi_lines() {
319        let theme = load_theme("tokyo-night").expect("failed to load built-in theme");
320        let lines = vec!["set x = 1", "set y = 2"];
321        let rendered = highlight_lines_to_ansi_lines(&lines, Grammar::ObjectScript, &theme)
322            .expect("failed to highlight lines");
323        assert_eq!(rendered.len(), 2);
324    }
325}