Skip to main content

vtcode_commons/
ansi.rs

1//! ANSI escape sequence parser and utilities
2
3/// Strip ANSI escape codes from text, keeping only plain text
4pub fn strip_ansi(text: &str) -> String {
5    let mut output = Vec::with_capacity(text.len());
6    let bytes = text.as_bytes();
7    let mut i = 0;
8
9    while i < bytes.len() {
10        if bytes[i] == 0x1b {
11            // Start of escape sequence
12            if i + 1 < bytes.len() {
13                match bytes[i + 1] {
14                    b'[' => {
15                        // CSI sequence - ends with a letter in range 0x40-0x7E
16                        i += 2;
17                        while i < bytes.len() && !(0x40..=0x7e).contains(&bytes[i]) {
18                            i += 1;
19                        }
20                        if i < bytes.len() {
21                            i += 1; // Include the final letter
22                        }
23                    }
24                    b']' => {
25                        // OSC sequence - ends with BEL (0x07) or ST (ESC \)
26                        i += 2;
27                        while i < bytes.len() {
28                            if bytes[i] == 0x07 {
29                                // BEL terminator
30                                i += 1;
31                                break;
32                            }
33                            if bytes[i] == 0x1b && i + 1 < bytes.len() && bytes[i + 1] == b'\\' {
34                                // ST terminator (ESC \)
35                                i += 2;
36                                break;
37                            }
38                            i += 1;
39                        }
40                    }
41                    b'P' | b'^' | b'_' | b'X' => {
42                        // DCS/PM/APC/SOS sequence - ends with ST (ESC \)
43                        i += 2;
44                        while i < bytes.len() {
45                            if bytes[i] == 0x1b && i + 1 < bytes.len() && bytes[i + 1] == b'\\' {
46                                i += 2;
47                                break;
48                            }
49                            i += 1;
50                        }
51                    }
52                    _ => {
53                        // Other 2-character escape sequences
54                        if bytes[i + 1] < 128 {
55                            i += 2;
56                        } else {
57                            i += 1;
58                        }
59                    }
60                }
61            } else {
62                i += 1;
63            }
64        } else if bytes[i] == b'\n' || bytes[i] == b'\r' || bytes[i] == b'\t' {
65            output.push(bytes[i]);
66            i += 1;
67        } else if bytes[i] < 32 {
68            i += 1;
69        } else {
70            output.push(bytes[i]);
71            i += 1;
72        }
73    }
74
75    unsafe { String::from_utf8_unchecked(output) }
76}
77
78/// Parse and determine the length of the ANSI escape sequence at the start of text
79pub fn parse_ansi_sequence(text: &str) -> Option<usize> {
80    let bytes = text.as_bytes();
81    if bytes.len() < 2 || bytes[0] != 0x1b {
82        return None;
83    }
84
85    let kind = bytes[1];
86    match kind {
87        b'[' => {
88            for (index, byte) in bytes.iter().enumerate().skip(2) {
89                if (0x40..=0x7e).contains(byte) {
90                    return Some(index + 1);
91                }
92            }
93            None
94        }
95        b']' => {
96            for index in 2..bytes.len() {
97                match bytes[index] {
98                    0x07 => return Some(index + 1),
99                    0x1b if index + 1 < bytes.len() && bytes[index + 1] == b'\\' => {
100                        return Some(index + 2);
101                    }
102                    _ => {}
103                }
104            }
105            None
106        }
107        b'P' | b'^' | b'_' | b'X' => {
108            for index in 2..bytes.len() {
109                if bytes[index] == 0x1b && index + 1 < bytes.len() && bytes[index + 1] == b'\\' {
110                    return Some(index + 2);
111                }
112            }
113            None
114        }
115        _ => {
116            if bytes[1] < 128 {
117                Some(2)
118            } else {
119                Some(1)
120            }
121        }
122    }
123}
124
125/// Fast ASCII-only ANSI stripping for performance-critical paths
126pub fn strip_ansi_ascii_only(text: &str) -> String {
127    let mut output = String::with_capacity(text.len());
128    let bytes = text.as_bytes();
129    let mut i = 0;
130    let mut last_valid = 0;
131
132    while i < bytes.len() {
133        if bytes[i] == 0x1b {
134            if last_valid < i {
135                output.push_str(&text[last_valid..i]);
136            }
137
138            if i + 1 < bytes.len() {
139                match bytes[i + 1] {
140                    b'[' => {
141                        i += 2;
142                        while i < bytes.len() && !(0x40..=0x7e).contains(&bytes[i]) {
143                            i += 1;
144                        }
145                        if i < bytes.len() {
146                            i += 1;
147                        }
148                    }
149                    b']' => {
150                        i += 2;
151                        while i < bytes.len() {
152                            if bytes[i] == 0x07 {
153                                i += 1;
154                                break;
155                            }
156                            if bytes[i] == 0x1b && i + 1 < bytes.len() && bytes[i + 1] == b'\\' {
157                                i += 2;
158                                break;
159                            }
160                            i += 1;
161                        }
162                    }
163                    b'P' | b'^' | b'_' | b'X' => {
164                        i += 2;
165                        while i < bytes.len() {
166                            if bytes[i] == 0x1b && i + 1 < bytes.len() && bytes[i + 1] == b'\\' {
167                                i += 2;
168                                break;
169                            }
170                            i += 1;
171                        }
172                    }
173                    _ => {
174                        if bytes[i + 1] < 128 {
175                            i += 2;
176                        } else {
177                            i += 1;
178                        }
179                    }
180                }
181            } else {
182                i += 1;
183            }
184            last_valid = i;
185        } else {
186            i += 1;
187        }
188    }
189
190    if last_valid < text.len() {
191        output.push_str(&text[last_valid..]);
192    }
193
194    output
195}
196
197/// Detect if text contains unicode characters that need special handling
198pub fn contains_unicode(text: &str) -> bool {
199    text.bytes().any(|b| b >= 0x80)
200}