1pub use crate::text::WidthMode;
22use crate::unicode::graphemes;
23
24#[derive(Debug, Clone, PartialEq, Eq)]
26pub enum Token<'a> {
27 Text {
29 text: &'a [u8],
31 width: u16,
33 },
34 Escape(&'a [u8]),
36 Control(u8),
39}
40
41pub fn string_width(bytes: &[u8], mode: WidthMode, eaw_wide: bool) -> usize {
45 tokenize(bytes, mode, eaw_wide)
46 .filter_map(|t| match t {
47 Token::Text { width, .. } => Some(width as usize),
48 _ => None,
49 })
50 .sum()
51}
52
53pub fn tokenize(bytes: &[u8], mode: WidthMode, eaw_wide: bool) -> Tokenizer<'_> {
55 Tokenizer {
56 bytes,
57 pos: 0,
58 mode,
59 eaw_wide,
60 }
61}
62
63pub struct Tokenizer<'a> {
65 bytes: &'a [u8],
66 pos: usize,
67 mode: WidthMode,
68 eaw_wide: bool,
69}
70
71impl<'a> Iterator for Tokenizer<'a> {
72 type Item = Token<'a>;
73
74 fn next(&mut self) -> Option<Token<'a>> {
75 if self.pos >= self.bytes.len() {
76 return None;
77 }
78
79 let b = self.bytes[self.pos];
80
81 if b == 0x1b || is_c1_introducer(b) {
83 let start = self.pos;
84 let end = scan_sequence(self.bytes, start);
85 self.pos = end;
86 return Some(Token::Escape(&self.bytes[start..end]));
87 }
88
89 if b < 0x20 || b == 0x7f || (0x80..=0x9f).contains(&b) {
92 self.pos += 1;
93 return Some(Token::Control(b));
94 }
95
96 let chunk_start = self.pos;
101 let mut chunk_end = self.pos;
102 while chunk_end < self.bytes.len() {
103 let bb = self.bytes[chunk_end];
104 if bb == 0x1b || bb < 0x20 || bb == 0x7f || (0x80..=0x9f).contains(&bb) {
105 break;
106 }
107 let n = utf8_char_len(bb);
108 if n == 0 || chunk_end + n > self.bytes.len() {
109 break;
110 }
111 chunk_end += n;
112 }
113 if chunk_end == chunk_start {
114 self.pos += 1;
117 return Some(Token::Control(b));
118 }
119 let raw = &self.bytes[chunk_start..chunk_end];
120 let valid = match std::str::from_utf8(raw) {
121 Ok(s) => s,
122 Err(e) => {
123 let up_to = e.valid_up_to();
124 if up_to == 0 {
125 self.pos += 1;
126 return Some(Token::Control(b));
127 }
128 unsafe { std::str::from_utf8_unchecked(&raw[..up_to]) }
130 }
131 };
132 let g = graphemes(valid).next()?;
133 let g_bytes = g.as_bytes();
134 self.pos = chunk_start + g_bytes.len();
135 let width = self.mode.grapheme_width(g, self.eaw_wide) as u16;
136 Some(Token::Text {
137 text: &self.bytes[chunk_start..chunk_start + g_bytes.len()],
138 width,
139 })
140 }
141}
142
143#[inline]
144fn utf8_char_len(b: u8) -> usize {
145 match b {
146 0x00..=0x7f => 1,
147 0xc2..=0xdf => 2,
148 0xe0..=0xef => 3,
149 0xf0..=0xf4 => 4,
150 _ => 0,
151 }
152}
153
154#[inline]
155fn is_c1_introducer(b: u8) -> bool {
156 matches!(b, 0x90 | 0x98 | 0x9b | 0x9d | 0x9e | 0x9f)
157}
158
159fn scan_sequence(bytes: &[u8], start: usize) -> usize {
167 let len = bytes.len();
168 let head = bytes[start];
169
170 match head {
171 0x1b => {
172 let i = start + 1;
173 if i >= len {
174 return len;
175 }
176 match bytes[i] {
177 b'[' => scan_csi(bytes, i + 1),
178 b']' | b'P' | b'X' | b'^' | b'_' => scan_string(bytes, i + 1),
179 _ => scan_esc_intermediate(bytes, i),
180 }
181 }
182 0x9b => scan_csi(bytes, start + 1),
183 0x9d | 0x90 | 0x98 | 0x9e | 0x9f => scan_string(bytes, start + 1),
184 _ => start + 1,
185 }
186}
187
188fn scan_csi(bytes: &[u8], from: usize) -> usize {
189 let len = bytes.len();
190 let mut i = from;
191 while i < len {
192 if (0x40..=0x7e).contains(&bytes[i]) {
193 return i + 1;
194 }
195 i += 1;
196 }
197 len
198}
199
200fn scan_string(bytes: &[u8], from: usize) -> usize {
201 let len = bytes.len();
202 let mut i = from;
203 while i < len {
204 let b = bytes[i];
205 if b == 0x07 || b == 0x9c {
206 return i + 1;
207 }
208 if b == 0x1b && i + 1 < len && bytes[i + 1] == b'\\' {
209 return i + 2;
210 }
211 if b == 0x1b {
214 return i;
215 }
216 i += 1;
217 }
218 len
219}
220
221fn scan_esc_intermediate(bytes: &[u8], at: usize) -> usize {
222 let len = bytes.len();
223 let mut i = at;
224 while i < len && (0x20..=0x2f).contains(&bytes[i]) {
225 i += 1;
226 }
227 if i < len { i + 1 } else { len }
228}
229
230#[cfg(test)]
231mod tests {
232 use super::*;
233
234 #[test]
235 fn tokenize_plain_text() {
236 let toks: Vec<_> = tokenize(b"abc", WidthMode::Grapheme, false).collect();
237 assert_eq!(toks.len(), 3);
238 assert!(matches!(
239 toks[0],
240 Token::Text {
241 text: b"a",
242 width: 1
243 }
244 ));
245 }
246
247 #[test]
248 fn tokenize_csi() {
249 let toks: Vec<_> = tokenize(b"\x1b[31mhi\x1b[m", WidthMode::Grapheme, false).collect();
250 assert_eq!(toks.len(), 4);
251 assert!(matches!(toks[0], Token::Escape(b"\x1b[31m")));
252 assert!(matches!(toks[3], Token::Escape(b"\x1b[m")));
253 }
254
255 #[test]
256 fn tokenize_osc_bel() {
257 let toks: Vec<_> = tokenize(b"\x1b]0;title\x07rest", WidthMode::Grapheme, false).collect();
258 assert!(matches!(toks[0], Token::Escape(b"\x1b]0;title\x07")));
259 }
260
261 #[test]
262 fn tokenize_osc_st() {
263 let toks: Vec<_> =
264 tokenize(b"\x1b]0;title\x1b\\rest", WidthMode::Grapheme, false).collect();
265 assert!(matches!(toks[0], Token::Escape(b"\x1b]0;title\x1b\\")));
266 }
267
268 #[test]
269 fn tokenize_wide_char() {
270 let toks: Vec<_> = tokenize("中".as_bytes(), WidthMode::Grapheme, false).collect();
271 assert_eq!(toks.len(), 1);
272 assert!(matches!(toks[0], Token::Text { width: 2, .. }));
273 }
274
275 #[test]
276 fn tokenize_newline_control() {
277 let toks: Vec<_> = tokenize(b"a\nb", WidthMode::Grapheme, false).collect();
278 assert_eq!(toks.len(), 3);
279 assert!(matches!(toks[1], Token::Control(b'\n')));
280 }
281
282 #[test]
283 fn string_width_ignores_escapes() {
284 assert_eq!(
285 string_width(b"\x1b[31mhello\x1b[m", WidthMode::Grapheme, false),
286 5
287 );
288 assert_eq!(
289 string_width("中文".as_bytes(), WidthMode::Grapheme, false),
290 4
291 );
292 }
293
294 #[test]
295 fn tokenize_8bit_csi() {
296 let toks: Vec<_> = tokenize(b"\x9b31mhi\x9bm", WidthMode::Grapheme, false).collect();
297 assert!(matches!(toks[0], Token::Escape(b"\x9b31m")));
298 assert!(matches!(toks[3], Token::Escape(b"\x9bm")));
299 }
300
301 #[test]
302 fn tokenize_8bit_osc_with_8bit_st() {
303 let toks: Vec<_> = tokenize(b"\x9d0;title\x9crest", WidthMode::Grapheme, false).collect();
305 assert!(matches!(toks[0], Token::Escape(b"\x9d0;title\x9c")));
306 assert_eq!(toks.len(), 5);
308 }
309
310 #[test]
311 fn tokenize_8bit_dcs() {
312 let toks: Vec<_> = tokenize(b"\x90q!data\x9cafter", WidthMode::Grapheme, false).collect();
313 assert!(matches!(toks[0], Token::Escape(b"\x90q!data\x9c")));
314 }
315
316 #[test]
317 fn tokenize_8bit_sos_pm_apc() {
318 for &intro in &[0x98u8, 0x9e, 0x9f] {
319 let mut input = vec![intro];
320 input.extend_from_slice(b"payload");
321 input.push(0x9c);
322 let toks: Vec<_> = tokenize(&input, WidthMode::Grapheme, false).collect();
323 match toks[0] {
324 Token::Escape(esc) => {
325 assert_eq!(esc[0], intro);
326 assert_eq!(*esc.last().unwrap(), 0x9c);
327 }
328 ref other => panic!("expected Escape for 0x{intro:02x}, got {other:?}"),
329 }
330 }
331 }
332
333 #[test]
334 fn tokenize_standalone_c1_is_control() {
335 let toks: Vec<_> = tokenize(b"a\x84b", WidthMode::Grapheme, false).collect();
337 assert_eq!(toks.len(), 3);
338 assert!(matches!(toks[1], Token::Control(0x84)));
339 }
340
341 #[test]
342 fn tokenize_8bit_st_outside_string_is_control() {
343 let toks: Vec<_> = tokenize(b"\x9c", WidthMode::Grapheme, false).collect();
345 assert_eq!(toks.len(), 1);
346 assert!(matches!(toks[0], Token::Control(0x9c)));
347 }
348
349 #[test]
350 fn tokenize_two_byte_esc() {
351 let toks: Vec<_> = tokenize(b"a\x1b=b", WidthMode::Grapheme, false).collect();
353 assert_eq!(toks.len(), 3);
354 assert!(matches!(toks[1], Token::Escape(b"\x1b=")));
355 }
356}