vtcode_commons/
ansi.rs

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

vtcode_commons/
ansi.rs
