Improved ANSI passthrough (#1596)

Improve handling of ANSI passthrough. Fix ANSI passthrough for --wrap=never. Add test for ANSI passthrough.

5 files changed, 285 insertions, 62 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c7971bb4..eecdd130 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,8 @@
 
 ## Bugfixes
 
+- Fix for poor performance when ANSI escape sequences are piped to `bat`, see #1596 (@eth-p)
+- Fix for incorrect handling of ANSI escape sequences when using `--wrap=never`, see #1596 (@eth-p)
 - Python syntax highlighting no longer suffers from abysmal performance in specific scenarios. See #1688 (@keith-hall)
 - First line not shown in diff context. See #1891 (@divagant-martian)
 - Do not ignore syntaxes that handle file names with a `*.conf` extension. See #1703 (@cbolgiano)
diff --git a/src/lib.rs b/src/lib.rs
index 02e1fefc..37b1cd83 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -46,6 +46,7 @@ pub(crate) mod printer;
 pub mod style;
 pub(crate) mod syntax_mapping;
 mod terminal;
+mod vscreen;
 pub(crate) mod wrapping;
 
 pub use pretty_printer::{Input, PrettyPrinter};
diff --git a/src/printer.rs b/src/printer.rs
index 1ecbbd74..6b8cbc3a 100644
--- a/src/printer.rs
+++ b/src/printer.rs
@@ -30,6 +30,7 @@ use crate::input::OpenedInput;
 use crate::line_range::RangeCheckResult;
 use crate::preprocessor::{expand_tabs, replace_nonprintable};
 use crate::terminal::{as_terminal_escaped, to_ansi_color};
+use crate::vscreen::AnsiStyle;
 use crate::wrapping::WrappingMode;
 
 pub(crate) trait Printer {
@@ -118,7 +119,7 @@ pub(crate) struct InteractivePrinter<'a> {
     config: &'a Config<'a>,
     decorations: Vec<Box<dyn Decoration>>,
     panel_width: usize,
-    ansi_prefix_sgr: String,
+    ansi_style: AnsiStyle,
     content_type: Option<ContentType>,
     #[cfg(feature = "git")]
     pub line_changes: &'a Option<LineChanges>,
@@ -202,7 +203,7 @@ impl<'a> InteractivePrinter<'a> {
             config,
             decorations,
             content_type: input.reader.content_type,
-            ansi_prefix_sgr: String::new(),
+            ansi_style: AnsiStyle::new(),
             #[cfg(feature = "git")]
             line_changes,
             highlighter_from_set,
@@ -444,35 +445,51 @@ impl<'a> Printer for InteractivePrinter<'a> {
             let italics = self.config.use_italic_text;
 
             for &(style, region) in &regions {
-                let text = &*self.preprocess(region, &mut cursor_total);
-                let text_trimmed = text.trim_end_matches(|c| c == '\r' || c == '\n');
-                write!(
-                    handle,
-                    "{}",
-                    as_terminal_escaped(
-                        style,
-                        text_trimmed,
-                        true_color,
-                        colored_output,
-                        italics,
-                        background_color
-                    )
-                )?;
+                let ansi_iterator = AnsiCodeIterator::new(region);
+                for chunk in ansi_iterator {
+                    match chunk {
+                        // ANSI escape passthrough.
+                        (ansi, true) => {
+                            self.ansi_style.update(ansi);
+                            write!(handle, "{}", ansi)?;
+                        }
 
-                if text.len() != text_trimmed.len() {
-                    if let Some(background_color) = background_color {
-                        let ansi_style = Style {
-                            background: to_ansi_color(background_color, true_color),
-                            ..Default::default()
-                        };
-                        let width = if cursor_total <= cursor_max {
-                            cursor_max - cursor_total + 1
-                        } else {
-                            0
-                        };
-                        write!(handle, "{}", ansi_style.paint(" ".repeat(width)))?;
+                        // Regular text.
+                        (text, false) => {
+                            let text = &*self.preprocess(text, &mut cursor_total);
+                            let text_trimmed = text.trim_end_matches(|c| c == '\r' || c == '\n');
+
+                            write!(
+                                handle,
+                                "{}",
+                                as_terminal_escaped(
+                                    style,
+                                    &format!("{}{}", self.ansi_style, text_trimmed),
+                                    true_color,
+                                    colored_output,
+                                    italics,
+                                    background_color
+                                )
+                            )?;
+
+                            if text.len() != text_trimmed.len() {
+                                if let Some(background_color) = background_color {
+                                    let ansi_style = Style {
+                                        background: to_ansi_color(background_color, true_color),
+                                        ..Default::default()
+                                    };
+
+                                    let width = if cursor_total <= cursor_max {
+                                        cursor_max - cursor_total + 1
+                                    } else {
+                                        0
+                                    };
+                                    write!(handle, "{}", ansi_style.paint(" ".repeat(width)))?;
+                                }
+                                write!(handle, "{}", &text[text_trimmed.len()..])?;
+                            }
+                        }
                     }
-                    write!(handle, "{}", &text[text_trimmed.len()..])?;
                 }
             }
 
@@ -482,31 +499,12 @@ impl<'a> Printer for InteractivePrinter<'a> {
         } else {
             for &(style, region) in &regions {
                 let ansi_iterator = AnsiCodeIterator::new(region);
-                let mut ansi_prefix: String = String::new();
                 for chunk in ansi_iterator {
                     match chunk {
                         // ANSI escape passthrough.
-                        (text, true) => {
-                            let is_ansi_csi = text.starts_with("\x1B[");
-
-                            if is_ansi_csi && text.ends_with('m') {
-                                // It's an ANSI SGR sequence.
-                                // We should be mostly safe to just append these together.
-                                ansi_prefix.push_str(text);
-                                if text == "\x1B[0m" {
-                                    self.ansi_prefix_sgr = "\x1B[0m".to_owned();
-                                } else {
-                                    self.ansi_prefix_sgr.push_str(text);
-                                }
-                            } else if is_ansi_csi {
-                                // It's a regular CSI sequence.
-                                // We should be mostly safe to just append these together.
-                                ansi_prefix.push_str(text);
-                            } else {
-                                // It's probably a VT100 code.
-                                // Passing it through is the safest bet.
-                                write!(handle, "{}", text)?;
-                            }
+                        (ansi, true) => {
+                            self.ansi_style.update(ansi);
+                            write!(handle, "{}", ansi)?;
                         }
 
                         // Regular text.
@@ -556,10 +554,7 @@ impl<'a> Printer for InteractivePrinter<'a> {
                                         "{}\n{}",
                                         as_terminal_escaped(
                                             style,
-                                            &*format!(
-                                                "{}{}{}",
-                                                self.ansi_prefix_sgr, ansi_prefix, line_buf
-                                            ),
+                                            &*format!("{}{}", self.ansi_style, line_buf),
                                             self.config.true_color,
                                             self.config.colored_output,
                                             self.config.use_italic_text,
@@ -585,19 +580,13 @@ impl<'a> Printer for InteractivePrinter<'a> {
                                 "{}",
                                 as_terminal_escaped(
                                     style,
-                                    &*format!(
-                                        "{}{}{}",
-                                        self.ansi_prefix_sgr, ansi_prefix, line_buf
-                                    ),
+                                    &*format!("{}{}", self.ansi_style, line_buf),
                                     self.config.true_color,
                                     self.config.colored_output,
                                     self.config.use_italic_text,
                                     background_color
                                 )
                             )?;
-
-                            // Clear the ANSI prefix buffer.
-                            ansi_prefix.clear();
                         }
                     }
                 }
diff --git a/src/vscreen.rs b/src/vscreen.rs
new file mode 100644
index 00000000..ea5d4da6
--- /dev/null
+++ b/src/vscreen.rs
@@ -0,0 +1,212 @@
+use std::fmt::{Display, Formatter};
+
+// Wrapper to avoid unnecessary branching when input doesn't have ANSI escape sequences.
+pub struct AnsiStyle {
+    attributes: Option<Attributes>,
+}
+
+impl AnsiStyle {
+    pub fn new() -> Self {
+        AnsiStyle { attributes: None }
+    }
+
+    pub fn update(&mut self, sequence: &str) -> bool {
+        match &mut self.attributes {
+            Some(a) => a.update(sequence),
+            None => {
+                self.attributes = Some(Attributes::new());
+                self.attributes.as_mut().unwrap().update(sequence)
+            }
+        }
+    }
+}
+
+impl Display for AnsiStyle {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self.attributes {
+            Some(ref a) => a.fmt(f),
+            None => Ok(()),
+        }
+    }
+}
+
+struct Attributes {
+    foreground: String,
+    background: String,
+    underlined: String,
+
+    /// The character set to use.
+    /// REGEX: `\^[()][AB0-3]`
+    charset: String,
+
+    /// A buffer for unknown sequences.
+    unknown_buffer: String,
+
+    /// ON:  ^[1m
+    /// OFF: ^[22m
+    bold: String,
+
+    /// ON:  ^[2m
+    /// OFF: ^[22m
+    dim: String,
+
+    /// ON:  ^[4m
+    /// OFF: ^[24m
+    underline: String,
+
+    /// ON:  ^[3m
+    /// OFF: ^[23m
+    italic: String,
+
+    /// ON:  ^[9m
+    /// OFF: ^[29m
+    strike: String,
+}
+
+impl Attributes {
+    pub fn new() -> Self {
+        Attributes {
+            foreground: "".to_owned(),
+            background: "".to_owned(),
+            underlined: "".to_owned(),
+            charset: "".to_owned(),
+            unknown_buffer: "".to_owned(),
+            bold: "".to_owned(),
+            dim: "".to_owned(),
+            underline: "".to_owned(),
+            italic: "".to_owned(),
+            strike: "".to_owned(),
+        }
+    }
+
+    /// Update the attributes with an escape sequence.
+    /// Returns `false` if the sequence is unsupported.
+    pub fn update(&mut self, sequence: &str) -> bool {
+        let mut chars = sequence.char_indices().skip(1);
+
+        if let Some((_, t)) = chars.next() {
+            match t {
+                '(' => self.update_with_charset('(', chars.map(|(_, c)| c)),
+                ')' => self.update_with_charset(')', chars.map(|(_, c)| c)),
+                '[' => {
+                    if let Some((i, last)) = chars.last() {
+                        // SAFETY: Always starts with ^[ and ends with m.
+                        self.update_with_csi(last, &sequence[2..i])
+                    } else {
+                        false
+                    }
+                }
+                _ => self.update_with_unsupported(sequence),
+            }
+        } else {
+            false
+        }
+    }
+
+    fn sgr_reset(&mut self) {
+        self.foreground.clear();
+        self.background.clear();
+        self.underlined.clear();
+        self.bold.clear();
+        self.dim.clear();
+        self.underline.clear();
+        self.italic.clear();
+        self.strike.clear();
+    }
+
+    fn update_with_sgr(&mut self, parameters: &str) -> bool {
+        let mut iter = parameters
+            .split(';')
+            .map(|p| if p.is_empty() { "0" } else { p })
+            .map(|p| p.parse::<u16>())
+            .map(|p| p.unwrap_or(0)); // Treat errors as 0.
+
+        while let Some(p) = iter.next() {
+            match p {
+                0 => self.sgr_reset(),
+                1 => self.bold = format!("\x1B[{}m", parameters),
+                2 => self.dim = format!("\x1B[{}m", parameters),
+                3 => self.italic = format!("\x1B[{}m", parameters),
+                4 => self.underline = format!("\x1B[{}m", parameters),
+                23 => self.italic.clear(),
+                24 => self.underline.clear(),
+                22 => {
+                    self.bold.clear();
+                    self.dim.clear();
+                }
+                30..=39 => self.foreground = Self::parse_color(p, &mut iter),
+                40..=49 => self.background = Self::parse_color(p, &mut iter),
+                58..=59 => self.underlined = Self::parse_color(p, &mut iter),
+                90..=97 => self.foreground = Self::parse_color(p, &mut iter),
+                100..=107 => self.foreground = Self::parse_color(p, &mut iter),
+                _ => {
+                    // Unsupported SGR sequence.
+                    // Be compatible and pretend one just wasn't was provided.
+                }
+            }
+        }
+
+        true
+    }
+
+    fn update_with_csi(&mut self, finalizer: char, sequence: &str) -> bool {
+        if finalizer == 'm' {
+            self.update_with_sgr(sequence)
+        } else {
+            false
+        }
+    }
+
+    fn update_with_unsupported(&mut self, sequence: &str) -> bool {
+        self.unknown_buffer.push_str(sequence);
+        false
+    }
+
+    fn update_with_charset(&mut self, kind: char, set: impl Iterator<Item = char>) -> bool {
+        self.charset = format!("\x1B{}{}", kind, set.take(1).collect::<String>());
+        true
+    }
+
+    fn parse_color(color: u16, parameters: &mut dyn Iterator<Item = u16>) -> String {
+        match color % 10 {
+            8 => match parameters.next() {
+                Some(5) /* 256-color */ => format!("\x1B[{};5;{}m", color, join(";", 1, parameters)),
+                Some(2) /* 24-bit color */ => format!("\x1B[{};2;{}m", color, join(";", 3, parameters)),
+                Some(c) => format!("\x1B[{};{}m", color, c),
+                _ => "".to_owned(),
+            },
+            9 => "".to_owned(),
+            _ => format!("\x1B[{}m", color),
+        }
+    }
+}
+
+impl Display for Attributes {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "{}{}{}{}{}{}{}{}{}",
+            self.foreground,
+            self.background,
+            self.underlined,
+            self.charset,
+            self.bold,
+            self.dim,
+            self.underline,
+            self.italic,
+            self.strike,
+        )
+    }
+}
+
+fn join(
+    delimiter: &str,
+    limit: usize,
+    iterator: &mut dyn Iterator<Item = impl ToString>,
+) -> String {
+    iterator
+        .take(limit)
+        .map(|i| i.to_string())
+        .collect::<Vec<String>>()
+        .join(delimiter)
+}
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index 17e9dc4e..986d9952 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -1226,6 +1226,25 @@ fn grid_for_file_without_newline() {
         .stderr("");
 }
 
+// Ensure that ANSI passthrough is emitted properly for both wrapping and non-wrapping printer.
+#[test]
+fn ansi_passthrough_emit() {
+    for wrapping in &["never", "character"] {
+        bat()
+            .arg("--paging=never")
+            .arg("--color=never")
+            .arg("--terminal-width=80")
+            .arg(format!("--wrap={}", wrapping))
+            .arg("--decorations=always")
+            .arg("--style=plain")
+            .write_stdin("\x1B[33mColor\nColor \x1B[m\nPlain\n")
+            .assert()
+            .success()
+            .stdout("\x1B[33m\x1B[33mColor\n\x1B[33mColor \x1B[m\nPlain\n")
+            .stderr("");
+    }
+}
+
 #[test]
 fn ignored_suffix_arg() {
     bat()