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}