console/
ansi.rs

1#[cfg(feature = "alloc")]
2use alloc::{borrow::Cow, string::String};
3use core::{
4    fmt::Display,
5    iter::{FusedIterator, Peekable},
6    str::CharIndices,
7};
8
9#[derive(Clone, Copy, Debug)]
10struct OscSequence;
11
12impl EscSequence for OscSequence {
13    const START: char = ']';
14
15    fn on_escape(next: Option<char>) -> EscAction {
16        let end = matches!(next, Some('\\'));
17        EscAction {
18            consume_next: end,
19            end,
20        }
21    }
22
23    fn on_char(c: char) -> EscAction {
24        EscAction {
25            consume_next: false,
26            end: c == '\u{07}',
27        }
28    }
29}
30
31#[derive(Clone, Copy, Debug)]
32struct DcsSequence;
33
34impl EscSequence for DcsSequence {
35    const START: char = 'P';
36
37    fn on_escape(next: Option<char>) -> EscAction {
38        match next {
39            Some('\\') => EscAction {
40                consume_next: true,
41                end: true,
42            },
43            None => EscAction {
44                consume_next: false,
45                end: true,
46            },
47            Some('\u{1b}') => EscAction {
48                consume_next: true,
49                end: false,
50            },
51            _ => EscAction {
52                consume_next: false,
53                end: false,
54            },
55        }
56    }
57
58    fn on_char(_c: char) -> EscAction {
59        EscAction {
60            consume_next: false,
61            end: false,
62        }
63    }
64}
65
66trait EscSequence {
67    fn on_escape(next: Option<char>) -> EscAction;
68    fn on_char(c: char) -> EscAction;
69    const START: char;
70}
71
72struct EscAction {
73    consume_next: bool,
74    end: bool,
75}
76
77fn consume_end_exclusive<S: EscSequence>(
78    it: &mut Peekable<CharIndices<'_>>,
79    start: usize,
80) -> usize {
81    let mut end = start + 1;
82    let Some((idx, start_char)) = it.next() else {
83        return end;
84    };
85    if start_char != S::START {
86        return end;
87    }
88    end = idx + start_char.len_utf8();
89
90    while let Some((idx, c)) = it.next() {
91        end = idx + c.len_utf8();
92
93        if c == '\u{1b}' {
94            let action = S::on_escape(it.peek().map(|(_, next)| *next));
95            if action.consume_next {
96                if let Some((next_idx, next)) = it.next() {
97                    end = next_idx + next.len_utf8();
98                }
99            }
100            if action.end {
101                return end;
102            }
103            continue;
104        }
105
106        match c {
107            '\u{9c}' => return end,
108            _ if S::on_char(c).end => return end,
109            _ => {}
110        }
111    }
112
113    end
114}
115
116fn find_dfa_end_exclusive_after_entry(it: &mut Peekable<CharIndices<'_>>) -> Option<usize> {
117    let mut state = State::S1;
118    let mut maybe_end = None;
119
120    loop {
121        let item = it.peek();
122
123        if let Some((idx, c)) = item {
124            state.transition(*c);
125
126            if state.is_final() {
127                maybe_end = Some(*idx);
128            }
129        }
130
131        if state.is_trapped() || item.is_none() {
132            return maybe_end.map(|end| end + 1);
133        }
134
135        it.next();
136    }
137}
138
139#[derive(Debug, Clone, Copy, Default)]
140enum State {
141    #[default]
142    Start,
143    S1,
144    S2,
145    S3,
146    S4,
147    S5,
148    S6,
149    S7,
150    S8,
151    S9,
152    S10,
153    S11,
154    Trap,
155}
156
157impl State {
158    fn is_final(&self) -> bool {
159        #[allow(clippy::match_like_matches_macro)]
160        match self {
161            Self::S3 | Self::S5 | Self::S6 | Self::S7 | Self::S8 | Self::S9 | Self::S11 => true,
162            _ => false,
163        }
164    }
165
166    fn is_trapped(&self) -> bool {
167        #[allow(clippy::match_like_matches_macro)]
168        match self {
169            Self::Trap => true,
170            _ => false,
171        }
172    }
173
174    fn transition(&mut self, c: char) {
175        *self = match c {
176            '\u{1b}' | '\u{9b}' => match self {
177                Self::Start => Self::S1,
178                _ => Self::Trap,
179            },
180            '(' | ')' => match self {
181                Self::S1 => Self::S2,
182                Self::S2 | Self::S4 => Self::S4,
183                _ => Self::Trap,
184            },
185            ';' => match self {
186                Self::S1 | Self::S2 | Self::S4 => Self::S4,
187                Self::S5 | Self::S6 | Self::S7 | Self::S8 | Self::S10 => Self::S10,
188                _ => Self::Trap,
189            },
190
191            '[' | '#' | '?' => match self {
192                Self::S1 | Self::S2 | Self::S4 => Self::S4,
193                _ => Self::Trap,
194            },
195            '0'..='2' => match self {
196                Self::S1 | Self::S4 => Self::S5,
197                Self::S2 => Self::S3,
198                Self::S5 => Self::S6,
199                Self::S6 => Self::S7,
200                Self::S7 => Self::S8,
201                Self::S8 => Self::S9,
202                Self::S10 => Self::S5,
203                _ => Self::Trap,
204            },
205            '3'..='9' => match self {
206                Self::S1 | Self::S4 => Self::S5,
207                Self::S2 => Self::S5,
208                Self::S5 => Self::S6,
209                Self::S6 => Self::S7,
210                Self::S7 => Self::S8,
211                Self::S8 => Self::S9,
212                Self::S10 => Self::S5,
213                _ => Self::Trap,
214            },
215            'A'..='P' | 'R' | 'Z' | 'c' | 'f'..='n' | 'q' | 'r' | 'y' | '=' | '>' | '<' => {
216                match self {
217                    Self::S1
218                    | Self::S2
219                    | Self::S4
220                    | Self::S5
221                    | Self::S6
222                    | Self::S7
223                    | Self::S8
224                    | Self::S10 => Self::S11,
225                    _ => Self::Trap,
226                }
227            }
228            _ => Self::Trap,
229        };
230    }
231}
232
233#[derive(Debug)]
234struct Matches<'a> {
235    s: &'a str,
236    it: Peekable<CharIndices<'a>>,
237}
238
239impl<'a> Matches<'a> {
240    fn new(s: &'a str) -> Self {
241        let it = s.char_indices().peekable();
242        Self { s, it }
243    }
244}
245
246#[derive(Debug)]
247struct Match<'a> {
248    text: &'a str,
249    start: usize,
250    end: usize,
251}
252
253impl<'a> Match<'a> {
254    #[inline]
255    pub(crate) fn as_str(&self) -> &'a str {
256        &self.text[self.start..self.end]
257    }
258}
259
260impl<'a> Iterator for Matches<'a> {
261    type Item = Match<'a>;
262
263    fn next(&mut self) -> Option<Self::Item> {
264        find_ansi_code_exclusive(&mut self.it).map(|(start, end)| Match {
265            text: self.s,
266            start,
267            end,
268        })
269    }
270}
271
272impl FusedIterator for Matches<'_> {}
273
274fn find_ansi_code_exclusive(it: &mut Peekable<CharIndices>) -> Option<(usize, usize)> {
275    'outer: loop {
276        let (start, entry) = *it.peek()?;
277        match entry {
278            '\u{1b}' => {
279                it.next();
280                match it.peek() {
281                    Some((_, OscSequence::START)) => {
282                        return Some((start, consume_end_exclusive::<OscSequence>(it, start)))
283                    }
284                    Some((_, DcsSequence::START)) => {
285                        return Some((start, consume_end_exclusive::<DcsSequence>(it, start)))
286                    }
287                    _ => {
288                        if let Some(end) = find_dfa_end_exclusive_after_entry(it) {
289                            return Some((start, end));
290                        }
291                        continue 'outer;
292                    }
293                }
294            }
295            '\u{9b}' => {
296                it.next();
297                if let Some(end) = find_dfa_end_exclusive_after_entry(it) {
298                    return Some((start, end));
299                }
300                continue 'outer;
301            }
302            _ => {
303                it.next();
304            }
305        }
306    }
307}
308
309/// Helper function to strip ansi codes.
310#[cfg(feature = "alloc")]
311pub fn strip_ansi_codes(s: &str) -> Cow<'_, str> {
312    let mut char_it = s.char_indices().peekable();
313    match find_ansi_code_exclusive(&mut char_it) {
314        Some(_) => {
315            let stripped: String = AnsiCodeIterator::new(s)
316                .filter_map(|(text, is_ansi)| if is_ansi { None } else { Some(text) })
317                .collect();
318            Cow::Owned(stripped)
319        }
320        None => Cow::Borrowed(s),
321    }
322}
323
324/// A wrapper struct that implements [`core::fmt::Display`], only displaying non-ansi parts.
325pub struct WithoutAnsi<'a> {
326    str: &'a str,
327}
328
329impl<'a> WithoutAnsi<'a> {
330    pub fn new(str: &'a str) -> Self {
331        Self { str }
332    }
333}
334
335impl Display for WithoutAnsi<'_> {
336    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
337        for (str, is_ansi) in AnsiCodeIterator::new(self.str) {
338            if !is_ansi {
339                f.write_str(str)?;
340            }
341        }
342        Ok(())
343    }
344}
345
346/// An iterator over ansi codes in a string.
347///
348/// This type can be used to scan over ansi codes in a string.
349/// It yields tuples in the form `(s, is_ansi)` where `s` is a slice of
350/// the original string and `is_ansi` indicates if the slice contains
351/// ansi codes or string values.
352pub struct AnsiCodeIterator<'a> {
353    s: &'a str,
354    pending_item: Option<(&'a str, bool)>,
355    last_idx: usize,
356    cur_idx: usize,
357    iter: Matches<'a>,
358}
359
360impl<'a> AnsiCodeIterator<'a> {
361    /// Creates a new ansi code iterator.
362    pub fn new(s: &'a str) -> AnsiCodeIterator<'a> {
363        AnsiCodeIterator {
364            s,
365            pending_item: None,
366            last_idx: 0,
367            cur_idx: 0,
368            iter: Matches::new(s),
369        }
370    }
371
372    /// Returns the string slice up to the current match.
373    pub fn current_slice(&self) -> &str {
374        &self.s[..self.cur_idx]
375    }
376
377    /// Returns the string slice from the current match to the end.
378    pub fn rest_slice(&self) -> &str {
379        &self.s[self.cur_idx..]
380    }
381}
382
383impl<'a> Iterator for AnsiCodeIterator<'a> {
384    type Item = (&'a str, bool);
385
386    fn next(&mut self) -> Option<(&'a str, bool)> {
387        if let Some(pending_item) = self.pending_item.take() {
388            self.cur_idx += pending_item.0.len();
389            Some(pending_item)
390        } else if let Some(m) = self.iter.next() {
391            let s = &self.s[self.last_idx..m.start];
392            self.last_idx = m.end;
393            if s.is_empty() {
394                self.cur_idx = m.end;
395                Some((m.as_str(), true))
396            } else {
397                self.cur_idx = m.start;
398                self.pending_item = Some((m.as_str(), true));
399                Some((s, false))
400            }
401        } else if self.last_idx < self.s.len() {
402            let rv = &self.s[self.last_idx..];
403            self.cur_idx = self.s.len();
404            self.last_idx = self.s.len();
405            Some((rv, false))
406        } else {
407            None
408        }
409    }
410}
411
412impl FusedIterator for AnsiCodeIterator<'_> {}
413
414#[cfg(test)]
415mod tests {
416    use super::*;
417
418    use core::fmt::Write;
419    use proptest::prelude::*;
420    use regex::Regex;
421    use std::sync::OnceLock;
422
423    // Keep the handwritten parser aligned with this reference regex, including OSC sequences and
424    // DCS passthrough sequences used by tmux.
425    fn strip_ansi_re() -> &'static Regex {
426        static RE: OnceLock<Regex> = OnceLock::new();
427
428        RE.get_or_init(|| {
429            Regex::new(concat!(
430                r"(?s)(?:",
431                r"\x1b\].*?(?:\x07|\x9c|\x1b\\|\z)|",
432                r"\x1bP(?:[^\x1b\x9c]|\x1b\x1b|\x1b[^\x1b\\])*?(?:\x9c|\x1b\\|\x1b\z|\z)|",
433                r"[\x1b\x9b]([()][012AB]|[\[()#;?]*(?:[0-9]{1,4}(?:;[0-9]{0,4})*)?",
434                r"[0-9A-PRZcf-nqry=><])",
435                r")",
436            ))
437            .unwrap()
438        })
439    }
440
441    impl<'a> PartialEq<Match<'a>> for regex::Match<'_> {
442        fn eq(&self, other: &Match<'a>) -> bool {
443            self.start() == other.start && self.end() == other.end
444        }
445    }
446
447    proptest! {
448        #[test]
449        fn dfa_matches_old_regex(s in r"([\x1b\x9b]?.*){0,5}") {
450            let old_matches: Vec<_> = strip_ansi_re().find_iter(&s).collect();
451            let new_matches: Vec<_> = Matches::new(&s).collect();
452            assert_eq!(old_matches, new_matches);
453        }
454    }
455
456    #[test]
457    fn dfa_matches_regex_on_small_strings() {
458        // To make sure the test runs in a reasonable time this is a slimmed down list of
459        // characters to reduce the groups that are only used with each other along with one
460        // arbitrarily chosen character not used in the regex (' ')
461        const POSSIBLE_BYTES: &[u8] = &[b' ', 0x1b, 0x9b, b'(', b'0', b'[', b';', b'3', b'C'];
462
463        fn check_all_strings_of_len(len: usize) {
464            _check_all_strings_of_len(len, &mut Vec::with_capacity(len));
465        }
466
467        fn _check_all_strings_of_len(len: usize, chunk: &mut Vec<u8>) {
468            if len == 0 {
469                if let Ok(s) = core::str::from_utf8(chunk) {
470                    let old_matches: Vec<_> = strip_ansi_re().find_iter(s).collect();
471                    let new_matches: Vec<_> = Matches::new(s).collect();
472                    assert_eq!(old_matches, new_matches);
473                }
474
475                return;
476            }
477
478            for b in POSSIBLE_BYTES {
479                chunk.push(*b);
480                _check_all_strings_of_len(len - 1, chunk);
481                chunk.pop();
482            }
483        }
484
485        for str_len in 0..=6 {
486            check_all_strings_of_len(str_len);
487        }
488    }
489
490    #[test]
491    fn complex_data() {
492        let s = std::fs::read_to_string(
493            std::path::Path::new("tests")
494                .join("data")
495                .join("sample_zellij_session.log"),
496        )
497        .unwrap();
498
499        let old_matches: Vec<_> = strip_ansi_re().find_iter(&s).collect();
500        let new_matches: Vec<_> = Matches::new(&s).collect();
501        assert_eq!(old_matches, new_matches);
502    }
503
504    #[test]
505    fn state_machine() {
506        let ansi_code = "\x1b)B";
507        let mut state = State::default();
508        assert!(!state.is_final());
509
510        for c in ansi_code.chars() {
511            state.transition(c);
512        }
513        assert!(state.is_final());
514
515        state.transition('A');
516        assert!(state.is_trapped());
517    }
518
519    #[test]
520    fn back_to_back_entry_char() {
521        let s = "\x1b\x1bf";
522        let matches: Vec<_> = Matches::new(s).map(|m| m.as_str()).collect();
523        assert_eq!(&["\x1bf"], matches.as_slice());
524    }
525
526    #[test]
527    fn early_paren_can_use_many_chars() {
528        let s = "\x1b(C";
529        let matches: Vec<_> = Matches::new(s).map(|m| m.as_str()).collect();
530        assert_eq!(&[s], matches.as_slice());
531    }
532
533    #[test]
534    fn long_run_of_digits() {
535        let s = "\u{1b}00000";
536        let matches: Vec<_> = Matches::new(s).map(|m| m.as_str()).collect();
537        assert_eq!(&[s], matches.as_slice());
538    }
539
540    #[test]
541    fn test_without_ansi() {
542        let str_with_ansi = "\x1b[1;97;41mError\x1b[0m";
543        let without_ansi = WithoutAnsi::new(str_with_ansi);
544        for _ in 0..2 {
545            let mut output = String::default();
546            write!(output, "{without_ansi}").unwrap();
547            assert_eq!(output, "Error");
548        }
549    }
550
551    #[test]
552    fn strip_osc8_hyperlink_st() {
553        let s = "\x1b]8;;file:///tmp/test\x1b\\hello\x1b]8;;\x1b\\";
554        assert_eq!(strip_ansi_codes(s).as_ref(), "hello");
555    }
556
557    #[test]
558    fn strip_osc8_hyperlink_bel() {
559        let s = "\x1b]8;;file:///tmp/test\x07hello\x1b]8;;\x07";
560        assert_eq!(strip_ansi_codes(s).as_ref(), "hello");
561    }
562
563    #[test]
564    fn strip_tmux_passthrough_dcs() {
565        let open = "\x1bPtmux;\x1b\x1b\x1b]8;;file:///tmp/test\x1b\x1b\\\x1b\\";
566        let close = "\x1bPtmux;\x1b\x1b\x1b]8;;\x1b\x1b\\\x1b\\";
567        assert_eq!(
568            strip_ansi_codes(&format!("{open}hello{close}")).as_ref(),
569            "hello"
570        );
571    }
572
573    #[test]
574    fn test_ansi_iter_re_vt100() {
575        let s = "\x1b(0lpq\x1b)Benglish";
576        let mut iter = AnsiCodeIterator::new(s);
577        assert_eq!(iter.next(), Some(("\x1b(0", true)));
578        assert_eq!(iter.next(), Some(("lpq", false)));
579        assert_eq!(iter.next(), Some(("\x1b)B", true)));
580        assert_eq!(iter.next(), Some(("english", false)));
581    }
582
583    #[test]
584    fn test_ansi_iter_re() {
585        use crate::style;
586        let s = format!("Hello {}!", style("World").red().force_styling(true));
587        let mut iter = AnsiCodeIterator::new(&s);
588        assert_eq!(iter.next(), Some(("Hello ", false)));
589        assert_eq!(iter.current_slice(), "Hello ");
590        assert_eq!(iter.rest_slice(), "\x1b[31mWorld\x1b[0m!");
591        assert_eq!(iter.next(), Some(("\x1b[31m", true)));
592        assert_eq!(iter.current_slice(), "Hello \x1b[31m");
593        assert_eq!(iter.rest_slice(), "World\x1b[0m!");
594        assert_eq!(iter.next(), Some(("World", false)));
595        assert_eq!(iter.current_slice(), "Hello \x1b[31mWorld");
596        assert_eq!(iter.rest_slice(), "\x1b[0m!");
597        assert_eq!(iter.next(), Some(("\x1b[0m", true)));
598        assert_eq!(iter.current_slice(), "Hello \x1b[31mWorld\x1b[0m");
599        assert_eq!(iter.rest_slice(), "!");
600        assert_eq!(iter.next(), Some(("!", false)));
601        assert_eq!(iter.current_slice(), "Hello \x1b[31mWorld\x1b[0m!");
602        assert_eq!(iter.rest_slice(), "");
603        assert_eq!(iter.next(), None);
604    }
605
606    #[test]
607    fn test_ansi_iter_re_on_multi() {
608        use crate::style;
609        let s = format!("{}", style("a").red().bold().force_styling(true));
610        let mut iter = AnsiCodeIterator::new(&s);
611        assert_eq!(iter.next(), Some(("\x1b[31m", true)));
612        assert_eq!(iter.current_slice(), "\x1b[31m");
613        assert_eq!(iter.rest_slice(), "\x1b[1ma\x1b[0m");
614        assert_eq!(iter.next(), Some(("\x1b[1m", true)));
615        assert_eq!(iter.current_slice(), "\x1b[31m\x1b[1m");
616        assert_eq!(iter.rest_slice(), "a\x1b[0m");
617        assert_eq!(iter.next(), Some(("a", false)));
618        assert_eq!(iter.current_slice(), "\x1b[31m\x1b[1ma");
619        assert_eq!(iter.rest_slice(), "\x1b[0m");
620        assert_eq!(iter.next(), Some(("\x1b[0m", true)));
621        assert_eq!(iter.current_slice(), "\x1b[31m\x1b[1ma\x1b[0m");
622        assert_eq!(iter.rest_slice(), "");
623        assert_eq!(iter.next(), None);
624    }
625}
console/ansi.rs

console/
ansi.rs
