Skip to main content

uncurses/ansi/
text.rs

1//! ANSI-aware byte-stream tokenizer for text utilities.
2//!
3//! ## Category
4//!
5//! The tokenizer classifies an input byte slice as visible grapheme clusters,
6//! complete ANSI escape/string sequences, or standalone control bytes. Width,
7//! stripping, truncation, and wrapping utilities all build on this stream.
8//!
9//! ## 7-bit and 8-bit controls
10//!
11//! Both 7-bit forms (`ESC [`, `ESC ]`, `ESC P`, `ESC _`) and 8-bit C1 forms
12//! (`0x9b`, `0x9d`, `0x90`, `0x9f`) are recognized. String controls terminate on
13//! BEL, `ST` (`ESC \\`), or 8-bit ST where applicable.
14//!
15//! ## Mode interaction
16//!
17//! This module does not interpret terminal modes or sequence semantics. Escape
18//! bytes are passed through as zero-width tokens so callers can preserve or drop
19//! them according to their own policy.
20
21pub use crate::text::WidthMode;
22use crate::unicode::graphemes;
23
24/// A single token produced by [`tokenize`].
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub enum Token<'a> {
27    /// A grapheme of visible text with its display width.
28    Text {
29        /// Grapheme bytes.
30        text: &'a [u8],
31        /// Display width in terminal cells.
32        width: u16,
33    },
34    /// An ANSI escape sequence (passed through verbatim, no width).
35    Escape(&'a [u8]),
36    /// A single control byte (C0, DEL, or a non-introducer C1) that isn't
37    /// part of an escape sequence.
38    Control(u8),
39}
40
41/// Return the display width of `bytes` ignoring ANSI escapes.
42///
43/// Non-UTF-8 bytes contribute no width.
44pub 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
53/// Tokenize an input byte slice into ANSI-aware tokens.
54pub 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
63/// Iterator returned by [`tokenize`].
64pub 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        // 7-bit escape (ESC) and 8-bit C1 introducers open a sequence.
82        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        // C0 controls (incl. DEL) and non-introducer C1 bytes are emitted as
90        // single control bytes.
91        if b < 0x20 || b == 0x7f || (0x80..=0x9f).contains(&b) {
92            self.pos += 1;
93            return Some(Token::Control(b));
94        }
95
96        // Plain text — walk forward one UTF-8 codepoint at a time, stopping at
97        // any byte that would start an escape or control token. Stepping per
98        // codepoint guarantees we never confuse a UTF-8 continuation byte
99        // (which can be in 0x80..=0xBF) with a C1 control.
100        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            // Invalid UTF-8 leading byte that wasn't a control — emit it
115            // verbatim as a Control byte to keep forward progress.
116            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                // SAFETY: validated up to `up_to` by from_utf8.
129                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
159/// Return the byte index past the end of an escape sequence starting at
160/// `start`.
161///
162/// `bytes[start]` is either `0x1B` (ESC) or an 8-bit C1 sequence introducer
163/// (`0x9B`, `0x9D`, `0x90`, `0x98`, `0x9E`, `0x9F`). The returned index points
164/// to the first byte after the sequence. If the sequence is incomplete the
165/// index is the end of `bytes`.
166fn 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        // Lone ESC (without trailing `\`) terminates the string but is left
212        // to be re-parsed as the next sequence.
213        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        // 0x9d "0;title" 0x9c "rest"
304        let toks: Vec<_> = tokenize(b"\x9d0;title\x9crest", WidthMode::Grapheme, false).collect();
305        assert!(matches!(toks[0], Token::Escape(b"\x9d0;title\x9c")));
306        // "rest" tokenizes to 4 single-grapheme Text tokens.
307        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        // IND (0x84) standalone — not an introducer, emitted as Control.
336        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        // Bare 0x9c with no preceding string is just a control byte.
344        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        // ESC = (DECKPAM)
352        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}