fm/modes/utils/
ansi.rs

1/// Copied from skim. See [ansi.rs](https://github.com/lotabout/skim/blob/master/src/ansi.rs)
2// Parse ANSI style code
3use std::default::Default;
4use std::mem;
5
6use beef::lean::Cow;
7use ratatui::style::{Color, Modifier, Style};
8use std::cmp::max;
9use vte::{Params, Perform};
10
11use crate::log_info;
12
13/// An ANSI Parser, will parse one line at a time.
14///
15/// It will cache the latest style used, that means if an style affect multiple
16/// lines, the parser will recognize it.
17#[derive(Default)]
18pub struct ANSIParser {
19    partial_str: String,
20    last_style: Style,
21
22    stripped: String,
23    stripped_char_count: usize,
24    fragments: Vec<(Style, (u32, u32))>, // [char_index_start, char_index_end)
25}
26
27impl Perform for ANSIParser {
28    fn print(&mut self, ch: char) {
29        self.partial_str.push(ch);
30    }
31
32    fn execute(&mut self, byte: u8) {
33        match byte {
34            // \b to delete character back
35            0x08 => {
36                self.partial_str.pop();
37            }
38            // put back \0 \r \n \t
39            0x00 | 0x0d | 0x0A | 0x09 => self.partial_str.push(byte as char),
40            // ignore all others
41            _ => log_info!("AnsiParser:execute ignored {:?}", byte),
42        }
43    }
44
45    fn hook(&mut self, params: &Params, _intermediates: &[u8], _ignore: bool, _action: char) {
46        log_info!("AnsiParser:hook ignored {:?}", params);
47    }
48
49    fn put(&mut self, byte: u8) {
50        log_info!("AnsiParser:put ignored {:?}", byte);
51    }
52
53    fn unhook(&mut self) {
54        log_info!("AnsiParser:unhook ignored");
55    }
56
57    fn osc_dispatch(&mut self, params: &[&[u8]], _bell_terminated: bool) {
58        log_info!("AnsiParser:osc ignored {:?}", params);
59    }
60
61    fn csi_dispatch(
62        &mut self,
63        params: &Params,
64        _intermediates: &[u8],
65        _ignore: bool,
66        action: char,
67    ) {
68        // https://en.wikipedia.org/wiki/ANSI_escape_code#SGR_(Select_Graphic_Rendition)_parameters
69        // Only care about graphic modes, ignore all others
70
71        if action != 'm' {
72            log_info!("ignore: params: {:?}, action : {:?}", params, action);
73            return;
74        }
75
76        // \[[m => means reset
77        let mut style = if params.is_empty() {
78            Style::default()
79        } else {
80            self.last_style
81        };
82
83        let mut iter = params.iter();
84        while let Some(code) = iter.next() {
85            match code[0] {
86                0 => style = Style::default(),
87                1 => style.add_modifier |= Modifier::BOLD,
88                2 => style.add_modifier |= !Modifier::BOLD,
89                4 => style.add_modifier |= Modifier::UNDERLINED,
90                5 => style.add_modifier |= Modifier::SLOW_BLINK,
91                7 => style.add_modifier |= Modifier::REVERSED,
92                num if (30..=37).contains(&num) => {
93                    style.fg = Some(Color::Indexed((num - 30) as u8));
94                }
95                38 => match iter.next() {
96                    Some(&[2]) => {
97                        // ESC[ 38;2;<r>;<g>;<b> m Select RGB foreground color
98                        let (r, g, b) = match (iter.next(), iter.next(), iter.next()) {
99                            (Some(r), Some(g), Some(b)) => (r[0] as u8, g[0] as u8, b[0] as u8),
100                            _ => {
101                                log_info!("ignore CSI {:?} m", params);
102                                continue;
103                            }
104                        };
105
106                        style.fg = Some(Color::Rgb(r, g, b));
107                    }
108                    Some(&[5]) => {
109                        // ESC[ 38;5;<n> m Select foreground color
110                        let color = match iter.next() {
111                            Some(color) => color[0] as u8,
112                            None => {
113                                log_info!("ignore CSI {:?} m", params);
114                                continue;
115                            }
116                        };
117
118                        style.fg = Some(Color::Indexed(color));
119                    }
120                    _ => {
121                        log_info!("error on parsing CSI {:?} m", params);
122                    }
123                },
124                39 => style.fg = Some(Color::Black),
125                num if (40..=47).contains(&num) => {
126                    style.bg = Some(Color::Indexed((num - 40) as u8));
127                }
128                48 => match iter.next() {
129                    Some(&[2]) => {
130                        // ESC[ 48;2;<r>;<g>;<b> m Select RGB background color
131                        let (r, g, b) = match (iter.next(), iter.next(), iter.next()) {
132                            (Some(r), Some(g), Some(b)) => (r[0] as u8, g[0] as u8, b[0] as u8),
133                            _ => {
134                                log_info!("ignore CSI {:?} m", params);
135                                continue;
136                            }
137                        };
138
139                        style.bg = Some(Color::Rgb(r, g, b));
140                    }
141                    Some(&[5]) => {
142                        // ESC[ 48;5;<n> m Select background color
143                        let color = match iter.next() {
144                            Some(color) => color[0] as u8,
145                            None => {
146                                log_info!("ignore CSI {:?} m", params);
147                                continue;
148                            }
149                        };
150
151                        style.bg = Some(Color::Indexed(color));
152                    }
153                    _ => {
154                        log_info!("ignore CSI {:?} m", params);
155                    }
156                },
157                49 => style.bg = Some(Color::Rgb(0, 0, 0)),
158                _ => {
159                    log_info!("ignore CSI {:?} m", params);
160                }
161            }
162        }
163
164        self.style_change(style);
165    }
166
167    fn esc_dispatch(&mut self, _intermediates: &[u8], _ignore: bool, _byte: u8) {
168        // ESC characters are replaced with \[
169        self.partial_str.push('"');
170        self.partial_str.push('[');
171    }
172}
173
174impl ANSIParser {
175    /// save the partial_str into fragments with current style
176    fn save_str(&mut self) {
177        if self.partial_str.is_empty() {
178            return;
179        }
180
181        let string = mem::take(&mut self.partial_str);
182        let string_char_count = string.chars().count();
183        self.fragments.push((
184            self.last_style,
185            (
186                self.stripped_char_count as u32,
187                (self.stripped_char_count + string_char_count) as u32,
188            ),
189        ));
190        self.stripped_char_count += string_char_count;
191        self.stripped.push_str(&string);
192    }
193
194    // accept a new style
195    fn style_change(&mut self, new_style: Style) {
196        if new_style == self.last_style {
197            return;
198        }
199
200        self.save_str();
201        self.last_style = new_style;
202    }
203
204    pub fn parse_ansi(&mut self, text: &str) -> AnsiString<'static> {
205        let mut statemachine = vte::Parser::new();
206
207        for byte in text.as_bytes() {
208            statemachine.advance(self, *byte);
209        }
210        self.save_str();
211
212        let stripped = mem::take(&mut self.stripped);
213        self.stripped_char_count = 0;
214        let fragments = mem::take(&mut self.fragments);
215        AnsiString::new_string(stripped, fragments)
216    }
217}
218
219/// A String that contains ANSI state (e.g. colors)
220///
221/// It is internally represented as Vec<(style, string)>
222#[derive(Clone, Debug)]
223pub struct AnsiString<'a> {
224    stripped: Cow<'a, str>,
225    // style: start, end
226    fragments: Option<Vec<(Style, (u32, u32))>>,
227}
228
229impl<'a> AnsiString<'a> {
230    pub fn new_empty() -> Self {
231        Self {
232            stripped: Cow::borrowed(""),
233            fragments: None,
234        }
235    }
236
237    fn new_raw_string(string: String) -> Self {
238        Self {
239            stripped: Cow::owned(string),
240            fragments: None,
241        }
242    }
243
244    fn new_raw_str(str_ref: &'a str) -> Self {
245        Self {
246            stripped: Cow::borrowed(str_ref),
247            fragments: None,
248        }
249    }
250
251    /// assume the fragments are ordered by (start, end) while end is exclusive
252    pub fn new_str(stripped: &'a str, fragments: Vec<(Style, (u32, u32))>) -> Self {
253        let fragments_empty =
254            fragments.is_empty() || (fragments.len() == 1 && fragments[0].0 == Style::default());
255        Self {
256            stripped: Cow::borrowed(stripped),
257            fragments: if fragments_empty {
258                None
259            } else {
260                Some(fragments)
261            },
262        }
263    }
264
265    /// assume the fragments are ordered by (start, end) while end is exclusive
266    pub fn new_string(stripped: String, fragments: Vec<(Style, (u32, u32))>) -> Self {
267        let fragments_empty =
268            fragments.is_empty() || (fragments.len() == 1 && fragments[0].0 == Style::default());
269        Self {
270            stripped: Cow::owned(stripped),
271            fragments: if fragments_empty {
272                None
273            } else {
274                Some(fragments)
275            },
276        }
277    }
278
279    pub fn parse(raw: &'a str) -> AnsiString<'static> {
280        ANSIParser::default().parse_ansi(raw)
281    }
282
283    #[inline]
284    pub fn is_empty(&self) -> bool {
285        self.stripped.is_empty()
286    }
287
288    #[inline]
289    pub fn into_inner(self) -> std::borrow::Cow<'a, str> {
290        std::borrow::Cow::Owned(self.stripped.into_owned())
291    }
292
293    pub fn iter(&'a self) -> Box<dyn Iterator<Item = (char, Style)> + 'a> {
294        if self.fragments.is_none() {
295            return Box::new(self.stripped.chars().map(|c| (c, Style::default())));
296        }
297
298        Box::new(AnsiStringIterator::new(
299            &self.stripped,
300            self.fragments.as_ref().unwrap(),
301        ))
302    }
303
304    pub fn has_styles(&self) -> bool {
305        self.fragments.is_some()
306    }
307
308    #[inline]
309    pub fn stripped(&self) -> &str {
310        &self.stripped
311    }
312
313    pub fn override_styles(&mut self, styles: Vec<(Style, (u32, u32))>) {
314        if styles.is_empty() {
315            // pass
316        } else if self.fragments.is_none() {
317            self.fragments = Some(styles);
318        } else {
319            let current_fragments = self.fragments.take().expect("unreachable");
320            let new_fragments = merge_fragments(&current_fragments, &styles);
321            self.fragments.replace(new_fragments);
322        }
323    }
324}
325
326impl<'a> From<&'a str> for AnsiString<'a> {
327    fn from(s: &'a str) -> AnsiString<'a> {
328        AnsiString::new_raw_str(s)
329    }
330}
331
332impl From<String> for AnsiString<'static> {
333    fn from(s: String) -> Self {
334        AnsiString::new_raw_string(s)
335    }
336}
337
338// (text, indices, highlight styleibute) -> AnsiString
339impl<'a> From<(&'a str, &'a [usize], Style)> for AnsiString<'a> {
340    fn from((text, indices, style): (&'a str, &'a [usize], Style)) -> Self {
341        let fragments = indices
342            .iter()
343            .map(|&idx| (style, (idx as u32, 1 + idx as u32)))
344            .collect();
345        AnsiString::new_str(text, fragments)
346    }
347}
348
349/// An iterator over all the (char, style) characters.
350pub struct AnsiStringIterator<'a> {
351    fragments: &'a [(Style, (u32, u32))],
352    fragment_idx: usize,
353    chars_iter: std::iter::Enumerate<std::str::Chars<'a>>,
354}
355
356impl<'a> AnsiStringIterator<'a> {
357    pub fn new(stripped: &'a str, fragments: &'a [(Style, (u32, u32))]) -> Self {
358        Self {
359            fragments,
360            fragment_idx: 0,
361            chars_iter: stripped.chars().enumerate(),
362        }
363    }
364}
365
366impl<'a> Iterator for AnsiStringIterator<'a> {
367    type Item = (char, Style);
368
369    fn next(&mut self) -> Option<Self::Item> {
370        match self.chars_iter.next() {
371            Some((char_idx, char)) => {
372                // update fragment_idx
373                loop {
374                    if self.fragment_idx >= self.fragments.len() {
375                        break;
376                    }
377
378                    let (_style, (_start, end)) = self.fragments[self.fragment_idx];
379                    if char_idx < (end as usize) {
380                        break;
381                    } else {
382                        self.fragment_idx += 1;
383                    }
384                }
385
386                let (style, (start, end)) = if self.fragment_idx >= self.fragments.len() {
387                    (Style::default(), (char_idx as u32, 1 + char_idx as u32))
388                } else {
389                    self.fragments[self.fragment_idx]
390                };
391
392                if (start as usize) <= char_idx && char_idx < (end as usize) {
393                    Some((char, style))
394                } else {
395                    Some((char, Style::default()))
396                }
397            }
398            None => None,
399        }
400    }
401}
402
403fn merge_fragments(
404    old: &[(Style, (u32, u32))],
405    new: &[(Style, (u32, u32))],
406) -> Vec<(Style, (u32, u32))> {
407    let mut ret = vec![];
408    let mut i = 0;
409    let mut j = 0;
410    let mut os = 0;
411
412    while i < old.len() && j < new.len() {
413        let (oa, (o_start, oe)) = old[i];
414        let (na, (ns, ne)) = new[j];
415        os = max(os, o_start);
416
417        if ns <= os && ne >= oe {
418            //   [--old--]   | [--old--]   |   [--old--] | [--old--]
419            // [----new----] | [---new---] | [---new---] | [--new--]
420            i += 1; // skip old
421        } else if ns <= os {
422            //           [--old--] |         [--old--] |   [--old--] |   [---old---]
423            // [--new--]           | [--new--]         | [--new--]   |   [--new--]
424            ret.push((na, (ns, ne)));
425            os = ne;
426            j += 1;
427        } else if ns >= oe {
428            // [--old--]         | [--old--]
429            //         [--new--] |           [--new--]
430            ret.push((oa, (os, oe)));
431            i += 1;
432        } else {
433            // [---old---] | [---old---] | [--old--]
434            //  [--new--]  |   [--new--] |      [--new--]
435            ret.push((oa, (os, ns)));
436            os = ns;
437        }
438    }
439
440    if i < old.len() {
441        for &(oa, (s, e)) in old[i..].iter() {
442            ret.push((oa, (max(os, s), e)))
443        }
444    }
445    if j < new.len() {
446        ret.extend_from_slice(&new[j..]);
447    }
448
449    ret
450}
451
452#[cfg(test)]
453mod tests {
454    use super::*;
455
456    #[test]
457    fn test_ansi_iterator() {
458        let input = "\x1B[48;2;5;10;15m\x1B[38;2;70;130;180mhi\x1B[0m";
459        let ansistring = ANSIParser::default().parse_ansi(input);
460        let mut it = ansistring.iter();
461        let style = Style {
462            fg: Some(Color::Rgb(70, 130, 180)),
463            bg: Some(Color::Rgb(5, 10, 15)),
464            ..Style::default()
465        };
466
467        assert_eq!(Some(('h', style)), it.next());
468        assert_eq!(Some(('i', style)), it.next());
469        assert_eq!(None, it.next());
470        assert_eq!(ansistring.stripped(), "hi");
471    }
472
473    #[test]
474    fn test_highlight_indices() {
475        let text = "abc";
476        let indices: Vec<usize> = vec![1];
477        let style = Style {
478            fg: Some(Color::Rgb(70, 130, 180)),
479            bg: Some(Color::Rgb(5, 10, 15)),
480            ..Style::default()
481        };
482
483        let ansistring = AnsiString::from((text, &indices as &[usize], style));
484        let mut it = ansistring.iter();
485
486        assert_eq!(Some(('a', Style::default())), it.next());
487        assert_eq!(Some(('b', style)), it.next());
488        assert_eq!(Some(('c', Style::default())), it.next());
489        assert_eq!(None, it.next());
490    }
491
492    #[test]
493    fn test_normal_string() {
494        let input = "ab";
495        let ansistring = ANSIParser::default().parse_ansi(input);
496
497        assert!(!ansistring.has_styles());
498
499        let mut it = ansistring.iter();
500        assert_eq!(Some(('a', Style::default())), it.next());
501        assert_eq!(Some(('b', Style::default())), it.next());
502        assert_eq!(None, it.next());
503
504        assert_eq!(ansistring.stripped(), "ab");
505    }
506
507    #[test]
508    fn test_multiple_styleibutes() {
509        let input = "\x1B[1;31mhi";
510        let ansistring = ANSIParser::default().parse_ansi(input);
511        let mut it = ansistring.iter();
512        let style = Style {
513            fg: Some(Color::Black),
514            add_modifier: Modifier::BOLD,
515            ..Style::default()
516        };
517
518        assert_eq!(Some(('h', style)), it.next());
519        assert_eq!(Some(('i', style)), it.next());
520        assert_eq!(None, it.next());
521        assert_eq!(ansistring.stripped(), "hi");
522    }
523
524    #[test]
525    fn test_reset() {
526        let input = "\x1B[35mA\x1B[mB";
527        let ansistring = ANSIParser::default().parse_ansi(input);
528        assert_eq!(ansistring.fragments.as_ref().map(|x| x.len()).unwrap(), 2);
529        assert_eq!(ansistring.stripped(), "AB");
530    }
531
532    #[test]
533    fn test_multi_bytes() {
534        let input = "中`\x1B[0m\x1B[1m\x1B[31mXYZ\x1B[0ms`";
535        let ansistring = ANSIParser::default().parse_ansi(input);
536        let mut it = ansistring.iter();
537        let default_style = Style::default();
538        let annotated = Style {
539            fg: Some(Color::Black),
540            add_modifier: Modifier::BOLD,
541            ..default_style
542        };
543
544        assert_eq!(Some(('中', default_style)), it.next());
545        assert_eq!(Some(('`', default_style)), it.next());
546        assert_eq!(Some(('X', annotated)), it.next());
547        assert_eq!(Some(('Y', annotated)), it.next());
548        assert_eq!(Some(('Z', annotated)), it.next());
549        assert_eq!(Some(('s', default_style)), it.next());
550        assert_eq!(Some(('`', default_style)), it.next());
551        assert_eq!(None, it.next());
552    }
553
554    #[test]
555    fn test_merge_fragments() {
556        let ao = Style::default();
557        let an = Style::default().bg(Color::Blue);
558
559        assert_eq!(
560            merge_fragments(&[(ao, (0, 1)), (ao, (1, 2))], &[]),
561            vec![(ao, (0, 1)), (ao, (1, 2))]
562        );
563
564        assert_eq!(
565            merge_fragments(&[], &[(an, (0, 1)), (an, (1, 2))]),
566            vec![(an, (0, 1)), (an, (1, 2))]
567        );
568
569        assert_eq!(
570            merge_fragments(
571                &[(ao, (1, 3)), (ao, (5, 6)), (ao, (9, 10))],
572                &[(an, (0, 1))]
573            ),
574            vec![(an, (0, 1)), (ao, (1, 3)), (ao, (5, 6)), (ao, (9, 10))]
575        );
576
577        assert_eq!(
578            merge_fragments(
579                &[(ao, (1, 3)), (ao, (5, 7)), (ao, (9, 11))],
580                &[(an, (0, 2))]
581            ),
582            vec![(an, (0, 2)), (ao, (2, 3)), (ao, (5, 7)), (ao, (9, 11))]
583        );
584
585        assert_eq!(
586            merge_fragments(
587                &[(ao, (1, 3)), (ao, (5, 7)), (ao, (9, 11))],
588                &[(an, (0, 3))]
589            ),
590            vec![(an, (0, 3)), (ao, (5, 7)), (ao, (9, 11))]
591        );
592
593        assert_eq!(
594            merge_fragments(
595                &[(ao, (1, 3)), (ao, (5, 7)), (ao, (9, 11))],
596                &[(an, (0, 6)), (an, (6, 7))]
597            ),
598            vec![(an, (0, 6)), (an, (6, 7)), (ao, (9, 11))]
599        );
600
601        assert_eq!(
602            merge_fragments(
603                &[(ao, (1, 3)), (ao, (5, 7)), (ao, (9, 11))],
604                &[(an, (1, 2))]
605            ),
606            vec![(an, (1, 2)), (ao, (2, 3)), (ao, (5, 7)), (ao, (9, 11))]
607        );
608
609        assert_eq!(
610            merge_fragments(
611                &[(ao, (1, 3)), (ao, (5, 7)), (ao, (9, 11))],
612                &[(an, (1, 3))]
613            ),
614            vec![(an, (1, 3)), (ao, (5, 7)), (ao, (9, 11))]
615        );
616
617        assert_eq!(
618            merge_fragments(
619                &[(ao, (1, 3)), (ao, (5, 7)), (ao, (9, 11))],
620                &[(an, (1, 4))]
621            ),
622            vec![(an, (1, 4)), (ao, (5, 7)), (ao, (9, 11))]
623        );
624
625        assert_eq!(
626            merge_fragments(
627                &[(ao, (1, 3)), (ao, (5, 7)), (ao, (9, 11))],
628                &[(an, (2, 3))]
629            ),
630            vec![(ao, (1, 2)), (an, (2, 3)), (ao, (5, 7)), (ao, (9, 11))]
631        );
632
633        assert_eq!(
634            merge_fragments(
635                &[(ao, (1, 3)), (ao, (5, 7)), (ao, (9, 11))],
636                &[(an, (2, 4))]
637            ),
638            vec![(ao, (1, 2)), (an, (2, 4)), (ao, (5, 7)), (ao, (9, 11))]
639        );
640
641        assert_eq!(
642            merge_fragments(
643                &[(ao, (1, 3)), (ao, (5, 7)), (ao, (9, 11))],
644                &[(an, (2, 6))]
645            ),
646            vec![(ao, (1, 2)), (an, (2, 6)), (ao, (6, 7)), (ao, (9, 11))]
647        );
648    }
649
650    #[test]
651    fn test_multi_byte_359() {
652        // https://github.com/lotabout/skim/issues/359
653        let highlight = Style::default().add_modifier(Modifier::BOLD);
654        let ansistring = AnsiString::new_str("ああa", vec![(highlight, (2, 3))]);
655        let mut it = ansistring.iter();
656        assert_eq!(Some(('あ', Style::default())), it.next());
657        assert_eq!(Some(('あ', Style::default())), it.next());
658        assert_eq!(Some(('a', highlight)), it.next());
659        assert_eq!(None, it.next());
660    }
661}
fm/modes/utils/ansi.rs

fm/modes/utils/
ansi.rs
