vtcode_commons/
ansi.rs

1#![expect(
2    clippy::indexing_slicing,
3    clippy::string_slice,
4    reason = "The ANSI parser validates byte positions while walking escape-sequence boundaries."
5)]
6
7//! Shared ANSI escape parser and stripping utilities for VT Code.
8//!
9//! See `docs/reference/ansi-in-vtcode.md` for the workspace usage map.
10
11use crate::ansi_codes::{BEL_BYTE, ESC_BYTE};
12use memchr::memchr;
13
14const ESC: u8 = ESC_BYTE;
15const BEL: u8 = BEL_BYTE;
16const DEL: u8 = 0x7f;
17const C1_ST: u8 = 0x9c;
18const C1_DCS: u8 = 0x90;
19const C1_SOS: u8 = 0x98;
20const C1_CSI: u8 = 0x9b;
21const C1_OSC: u8 = 0x9d;
22const C1_PM: u8 = 0x9e;
23const C1_APC: u8 = 0x9f;
24const CAN: u8 = 0x18;
25const SUB: u8 = 0x1a;
26const MAX_STRING_SEQUENCE_BYTES: usize = 4096;
27const MAX_CSI_SEQUENCE_BYTES: usize = 64;
28
29#[derive(Clone, Copy)]
30enum StringSequenceTerminator {
31    StOnly,
32    BelOrSt,
33}
34
35impl StringSequenceTerminator {
36    #[inline]
37    const fn allows_bel(self) -> bool {
38        matches!(self, Self::BelOrSt)
39    }
40}
41
42#[inline]
43fn parse_c1_at(bytes: &[u8], start: usize) -> Option<(u8, usize)> {
44    let first = *bytes.get(start)?;
45    if (0x80..=0x9f).contains(&first) {
46        return Some((first, 1));
47    }
48    None
49}
50
51#[inline]
52fn parse_csi(bytes: &[u8], start: usize) -> Option<usize> {
53    // ECMA-48 / ISO 6429 CSI grammar:
54    // - parameter bytes: 0x30..0x3F
55    // - intermediate bytes: 0x20..0x2F
56    // - final byte: 0x40..0x7E
57    // (See ANSI escape code article on Wikipedia, CSI section.)
58    let mut index = start;
59    let mut phase = 0u8; // 0=parameter, 1=intermediate
60    let mut consumed = 0usize;
61
62    while index < bytes.len() {
63        let byte = bytes[index];
64        if byte == ESC {
65            // VT100: ESC aborts current control sequence and starts a new one.
66            return Some(index);
67        }
68        if byte == CAN || byte == SUB {
69            // VT100: CAN/SUB abort current control sequence.
70            return Some(index + 1);
71        }
72
73        consumed += 1;
74        if consumed > MAX_CSI_SEQUENCE_BYTES {
75            // Bound malformed or hostile input.
76            return Some(index + 1);
77        }
78
79        if phase == 0 && (0x30..=0x3f).contains(&byte) {
80            index += 1;
81            continue;
82        }
83        if (0x20..=0x2f).contains(&byte) {
84            phase = 1;
85            index += 1;
86            continue;
87        }
88        if (0x40..=0x7e).contains(&byte) {
89            return Some(index + 1);
90        }
91
92        // Invalid CSI byte: abort sequence without consuming this byte.
93        return Some(index);
94    }
95
96    None
97}
98
99#[inline]
100fn parse_string_sequence(bytes: &[u8], start: usize, terminator: StringSequenceTerminator) -> Option<usize> {
101    let mut consumed = 0usize;
102    for index in start..bytes.len() {
103        if bytes[index] == ESC && !(index + 1 < bytes.len() && bytes[index + 1] == b'\\') {
104            // VT100: ESC aborts current sequence and begins a new one.
105            return Some(index);
106        }
107        if bytes[index] == CAN || bytes[index] == SUB {
108            return Some(index + 1);
109        }
110
111        if let Some((c1, len)) = parse_c1_at(bytes, index)
112            && c1 == C1_ST
113        {
114            return Some(index + len);
115        }
116
117        match bytes[index] {
118            BEL if terminator.allows_bel() => return Some(index + 1),
119            ESC if index + 1 < bytes.len() && bytes[index + 1] == b'\\' => return Some(index + 2),
120            _ => {}
121        }
122
123        consumed += 1;
124        if consumed > MAX_STRING_SEQUENCE_BYTES {
125            // Cap unbounded strings when terminator is missing.
126            return Some(index + 1);
127        }
128    }
129    None
130}
131
132#[inline]
133fn push_visible_byte(output: &mut Vec<u8>, byte: u8) {
134    if matches!(byte, b'\n' | b'\r' | b'\t') || !(byte < 32 || byte == DEL) {
135        output.push(byte);
136    }
137}
138
139#[inline]
140fn parse_ansi_sequence_bytes(bytes: &[u8]) -> Option<usize> {
141    if bytes.is_empty() {
142        return None;
143    }
144
145    if let Some((c1, c1_len)) = parse_c1_at(bytes, 0) {
146        return match c1 {
147            C1_CSI => parse_csi(bytes, c1_len),
148            C1_OSC => parse_string_sequence(bytes, c1_len, StringSequenceTerminator::BelOrSt),
149            C1_DCS | C1_SOS | C1_PM | C1_APC => parse_string_sequence(bytes, c1_len, StringSequenceTerminator::StOnly),
150            _ => Some(c1_len),
151        };
152    }
153
154    match bytes[0] {
155        ESC => {
156            if bytes.len() < 2 {
157                return None;
158            }
159
160            match bytes[1] {
161                b'[' => parse_csi(bytes, 2),
162                b']' => parse_string_sequence(bytes, 2, StringSequenceTerminator::BelOrSt),
163                b'P' | b'^' | b'_' | b'X' => parse_string_sequence(bytes, 2, StringSequenceTerminator::StOnly),
164                // Three-byte sequences: ESC + intermediate + final
165                // ESC SP {F,G,L,M,N} — 7/8-bit controls, ANSI conformance
166                // ESC # {3,4,5,6,8} — DEC line attributes / screen alignment
167                // ESC % {@ ,G} — character set selection (ISO 2022)
168                // ESC ( C / ESC ) C / ESC * C / ESC + C — G0-G3 designation
169                b' ' | b'#' | b'%' | b'(' | b')' | b'*' | b'+' => {
170                    if bytes.len() > 2 {
171                        Some(3)
172                    } else {
173                        None
174                    }
175                }
176                next if next < 128 => Some(2),
177                _ => Some(1),
178            }
179        }
180        _ => None,
181    }
182}
183
184/// Strip ANSI escape codes from text, returning a borrowed `Cow` when the
185/// input contains no ESC byte.  This is the preferred API for call-sites that
186/// want to avoid allocation on the common "no ANSI codes" path.
187pub fn strip_ansi_codes(text: &str) -> std::borrow::Cow<'_, str> {
188    if !text.contains('\x1b') {
189        return std::borrow::Cow::Borrowed(text);
190    }
191    std::borrow::Cow::Owned(strip_ansi(text))
192}
193
194/// Strip ANSI escape codes from text, keeping only plain text
195pub fn strip_ansi(text: &str) -> String {
196    let mut output = Vec::with_capacity(text.len());
197    let bytes = text.as_bytes();
198    let mut i = 0;
199
200    while i < bytes.len() {
201        let next_esc = memchr(ESC, &bytes[i..]).map_or(bytes.len(), |offset| i + offset);
202        // Pre-slice to avoid bounds checks in the inner loop — the range
203        // i..next_esc is provably within bytes[..].
204        for &b in &bytes[i..next_esc] {
205            push_visible_byte(&mut output, b);
206        }
207        i = next_esc;
208
209        if i >= bytes.len() {
210            break;
211        }
212
213        if let Some(len) = parse_ansi_sequence_bytes(&bytes[i..]) {
214            i += len;
215            continue;
216        } else {
217            // Incomplete/unterminated control sequence at end of available text.
218            break;
219        }
220    }
221
222    // The output is always valid UTF-8: `push_visible_byte` only filters ASCII
223    // control bytes (< 32 except \n/\r/\t, and DEL=127), which cannot be part of
224    // a multi-byte UTF-8 sequence (those bytes are all >= 0x80). `from_utf8`
225    // reuses the Vec allocation directly; `from_utf8_lossy().into_owned()` would
226    // copy even on the valid-UTF-8 fast path.
227    String::from_utf8(output).unwrap_or_else(|e| {
228        // Defensive fallback — should not happen given the invariant above.
229        String::from_utf8_lossy(&e.into_bytes()).into_owned()
230    })
231}
232
233/// Strip ANSI escape codes from arbitrary bytes, preserving non-control bytes.
234///
235/// This is the preferred API when input may contain raw C1 (8-bit) controls.
236fn strip_ansi_bytes(input: &[u8]) -> Vec<u8> {
237    let mut output = Vec::with_capacity(input.len());
238    let bytes = input;
239    let mut i = 0;
240
241    while i < bytes.len() {
242        // Pre-slice to the remaining portion so all indexing below shares one bounds edge.
243        let rest = &bytes[i..];
244
245        if (rest[0] == ESC || parse_c1_at(bytes, i).is_some())
246            && let Some(len) = parse_ansi_sequence_bytes(rest)
247        {
248            i += len;
249            continue;
250        }
251        if rest[0] == ESC || parse_c1_at(bytes, i).is_some() {
252            // Incomplete/unterminated control sequence at end of available text.
253            break;
254        }
255
256        push_visible_byte(&mut output, rest[0]);
257        i += 1;
258    }
259    output
260}
261
262/// Parse and determine the length of the ANSI escape sequence at the start of text
263pub fn parse_ansi_sequence(text: &str) -> Option<usize> {
264    let bytes = text.as_bytes();
265    parse_ansi_sequence_bytes(bytes)
266}
267
268/// Fast ASCII-only ANSI stripping for performance-critical paths
269pub fn strip_ansi_ascii_only(text: &str) -> String {
270    let mut output = String::with_capacity(text.len());
271    let bytes = text.as_bytes();
272    let mut search_start = 0;
273    let mut copy_start = 0;
274
275    while let Some(offset) = memchr(ESC, &bytes[search_start..]) {
276        let esc_index = search_start + offset;
277        if let Some(len) = parse_ansi_sequence_bytes(&bytes[esc_index..]) {
278            if copy_start < esc_index {
279                output.push_str(&text[copy_start..esc_index]);
280            }
281            copy_start = esc_index + len;
282            search_start = copy_start;
283        } else {
284            search_start = esc_index + 1;
285        }
286    }
287
288    if copy_start < text.len() {
289        output.push_str(&text[copy_start..]);
290    }
291
292    output
293}
294
295/// Detect if text contains unicode characters that need special handling
296#[must_use]
297pub fn contains_unicode(text: &str) -> bool {
298    text.bytes().any(|b| b >= 0x80)
299}
300
301#[cfg(test)]
302mod tests {
303    use super::{CAN, SUB, strip_ansi, strip_ansi_ascii_only};
304
305    #[test]
306    fn strips_esc_csi_sequences() {
307        let input = "a\x1b[31mred\x1b[0mz";
308        assert_eq!(strip_ansi(input), "aredz");
309        assert_eq!(strip_ansi_ascii_only(input), "aredz");
310    }
311
312    #[test]
313    fn utf8_encoded_c1_is_not_reprocessed_as_control() {
314        // XTerm/ECMA-48: controls are processed once; decoded UTF-8 text is not reprocessed as C1.
315        let input = "a\u{009b}31mred";
316        assert_eq!(strip_ansi(input), input);
317    }
318
319    #[test]
320    fn strip_removes_ascii_del_control() {
321        let input = format!("a{}b", char::from(0x7f));
322        assert_eq!(strip_ansi(&input), "ab");
323    }
324
325    #[test]
326    fn csi_aborts_on_esc_then_new_sequence_parses() {
327        let input = "a\x1b[31\x1b[32mgreen\x1b[0mz";
328        assert_eq!(strip_ansi(input), "agreenz");
329    }
330
331    #[test]
332    fn csi_aborts_on_can_and_sub() {
333        let can = format!("a\x1b[31{}b", char::from(CAN));
334        let sub = format!("a\x1b[31{}b", char::from(SUB));
335        assert_eq!(strip_ansi(&can), "ab");
336        assert_eq!(strip_ansi(&sub), "ab");
337    }
338
339    #[test]
340    fn osc_aborts_on_esc_non_st() {
341        let input = "a\x1b]title\x1b[31mred\x1b[0mz";
342        assert_eq!(strip_ansi(input), "aredz");
343    }
344
345    #[test]
346    fn incomplete_sequence_drops_tail() {
347        let input = "text\x1b[31";
348        assert_eq!(strip_ansi(input), "text");
349    }
350
351    #[test]
352    fn ascii_only_incomplete_sequence_keeps_tail() {
353        let input = "text\x1b[31";
354        assert_eq!(strip_ansi_ascii_only(input), input);
355    }
356
357    #[test]
358    fn strips_common_progress_redraw_sequences() {
359        // Common pattern for dynamic CLI updates:
360        // carriage return + erase line + redraw text.
361        let input = "\r\x1b[2KProgress 10%\r\x1b[2KDone\n";
362        assert_eq!(strip_ansi(input), "\rProgress 10%\rDone\n");
363    }
364
365    #[test]
366    fn strips_cursor_navigation_sequences() {
367        let input = "left\x1b[1D!\nup\x1b[1Arow";
368        assert_eq!(strip_ansi(input), "left!\nuprow");
369    }
370
371    #[test]
372    fn strip_ansi_bytes_supports_raw_c1_csi() {
373        let input = [b'a', 0x9b, b'3', b'1', b'm', b'r', b'e', b'd', 0x9b, b'0', b'm', b'z'];
374        let out = super::strip_ansi_bytes(&input);
375        assert_eq!(out, b"aredz");
376    }
377
378    #[test]
379    fn strip_ansi_bytes_supports_raw_c1_osc_and_st() {
380        let mut input = b"pre".to_vec();
381        input.extend_from_slice(&[0x9d]);
382        input.extend_from_slice(b"8;;https://example.com");
383        input.extend_from_slice(&[0x9c]);
384        input.extend_from_slice(b"link");
385        input.extend_from_slice(&[0x9d]);
386        input.extend_from_slice(b"8;;");
387        input.extend_from_slice(&[0x9c]);
388        input.extend_from_slice(b"post");
389        let out = super::strip_ansi_bytes(&input);
390        assert_eq!(out, b"prelinkpost");
391    }
392
393    #[test]
394    fn csi_respects_parameter_intermediate_final_grammar() {
395        // Parameter bytes ("1;2"), intermediate bytes (" "), then final ("m")
396        let input = "a\x1b[1;2 mred\x1b[0mz";
397        assert_eq!(strip_ansi(input), "aredz");
398    }
399
400    #[test]
401    fn malformed_csi_does_not_consume_following_text() {
402        // 0x10 is not valid CSI parameter/intermediate/final.
403        let malformed = format!("a\x1b[12{}visible", char::from(0x10));
404        assert_eq!(strip_ansi(&malformed), "avisible");
405    }
406
407    #[test]
408    fn strips_wikipedia_sgr_8bit_color_pattern() {
409        let input = "x\x1b[38;5;196mred\x1b[0my";
410        assert_eq!(strip_ansi(input), "xredy");
411    }
412
413    #[test]
414    fn strips_wikipedia_sgr_truecolor_pattern() {
415        let input = "x\x1b[48;2;12;34;56mblock\x1b[0my";
416        assert_eq!(strip_ansi(input), "xblocky");
417    }
418
419    #[test]
420    fn strips_wikipedia_osc8_hyperlink_pattern() {
421        let input = "go \x1b]8;;https://example.com\x1b\\here\x1b]8;;\x1b\\ now";
422        assert_eq!(strip_ansi(input), "go here now");
423    }
424
425    #[test]
426    fn strips_dec_private_mode_csi() {
427        let input = "a\x1b[?25lb\x1b[?25hc";
428        assert_eq!(strip_ansi(input), "abc");
429    }
430
431    #[test]
432    fn strips_three_byte_esc_sequences() {
433        // ESC # 8 = DEC screen alignment test
434        let input = "a\x1b#8b";
435        assert_eq!(strip_ansi(input), "ab");
436
437        // ESC ( B = designate US ASCII as G0
438        let input2 = "a\x1b(Bb";
439        assert_eq!(strip_ansi(input2), "ab");
440
441        // ESC SP F = 7-bit controls
442        let input3 = "a\x1b Fb";
443        assert_eq!(strip_ansi(input3), "ab");
444
445        // ESC % G = select UTF-8
446        let input4 = "a\x1b%Gb";
447        assert_eq!(strip_ansi(input4), "ab");
448    }
449
450    #[test]
451    fn incomplete_three_byte_esc_sequence_drops_tail() {
452        // ESC # at end — incomplete, should not consume past end
453        let input = "text\x1b#";
454        assert_eq!(strip_ansi(input), "text");
455    }
456}
vtcode_commons/ansi.rs

vtcode_commons/
ansi.rs
