Replace AnsiCodeIterator in printer.rs

This uses the new EscapeSequenceIterator, saving us a preprocessing step
for each line.

2 files changed, 62 insertions, 101 deletions

diff --git a/src/printer.rs b/src/printer.rs
index 45fd5336..6d495777 100644
--- a/src/printer.rs
+++ b/src/printer.rs
@@ -7,8 +7,6 @@ use nu_ansi_term::Style;
 
 use bytesize::ByteSize;
 
-use console::AnsiCodeIterator;
-
 use syntect::easy::HighlightLines;
 use syntect::highlighting::Color;
 use syntect::highlighting::Theme;
@@ -33,9 +31,23 @@ use crate::line_range::RangeCheckResult;
 use crate::preprocessor::{expand_tabs, replace_nonprintable};
 use crate::style::StyleComponent;
 use crate::terminal::{as_terminal_escaped, to_ansi_color};
-use crate::vscreen::{strip_problematic_sequences, AnsiStyle};
+use crate::vscreen::{AnsiStyle, EscapeSequence, EscapeSequenceIterator};
 use crate::wrapping::WrappingMode;
 
+const ANSI_UNDERLINE_ENABLE: EscapeSequence = EscapeSequence::CSI {
+    raw_sequence: "\x1B[4m",
+    parameters: "4",
+    intermediates: "",
+    final_byte: "m",
+};
+
+const ANSI_UNDERLINE_DISABLE: EscapeSequence = EscapeSequence::CSI {
+    raw_sequence: "\x1B[24m",
+    parameters: "24",
+    intermediates: "",
+    final_byte: "m",
+};
+
 pub enum OutputHandle<'a> {
     IoWrite(&'a mut dyn io::Write),
     FmtWrite(&'a mut dyn fmt::Write),
@@ -554,7 +566,7 @@ impl<'a> Printer for InteractivePrinter<'a> {
             self.config.highlighted_lines.0.check(line_number) == RangeCheckResult::InRange;
 
         if highlight_this_line && self.config.theme == "ansi" {
-            self.ansi_style.update("^[4m");
+            self.ansi_style.update(ANSI_UNDERLINE_ENABLE);
         }
 
         let background_color = self
@@ -581,18 +593,11 @@ impl<'a> Printer for InteractivePrinter<'a> {
             let italics = self.config.use_italic_text;
 
             for &(style, region) in &regions {
-                let text = strip_problematic_sequences(region);
-                let ansi_iterator = AnsiCodeIterator::new(&text);
+                let ansi_iterator = EscapeSequenceIterator::new(region);
                 for chunk in ansi_iterator {
                     match chunk {
-                        // ANSI escape passthrough.
-                        (ansi, true) => {
-                            self.ansi_style.update(ansi);
-                            write!(handle, "{}", ansi)?;
-                        }
-
                         // Regular text.
-                        (text, false) => {
+                        EscapeSequence::Text(text) => {
                             let text = &*self.preprocess(text, &mut cursor_total);
                             let text_trimmed = text.trim_end_matches(|c| c == '\r' || c == '\n');
 
@@ -626,6 +631,12 @@ impl<'a> Printer for InteractivePrinter<'a> {
                                 write!(handle, "{}", &text[text_trimmed.len()..])?;
                             }
                         }
+
+                        // ANSI escape passthrough.
+                        _ => {
+                            write!(handle, "{}", chunk.raw())?;
+                            self.ansi_style.update(chunk);
+                        }
                     }
                 }
             }
@@ -635,18 +646,11 @@ impl<'a> Printer for InteractivePrinter<'a> {
             }
         } else {
             for &(style, region) in &regions {
-                let text = strip_problematic_sequences(region);
-                let ansi_iterator = AnsiCodeIterator::new(&text);
+                let ansi_iterator = EscapeSequenceIterator::new(region);
                 for chunk in ansi_iterator {
                     match chunk {
-                        // ANSI escape passthrough.
-                        (ansi, true) => {
-                            self.ansi_style.update(ansi);
-                            write!(handle, "{}", ansi)?;
-                        }
-
                         // Regular text.
-                        (text, false) => {
+                        EscapeSequence::Text(text) => {
                             let text = self.preprocess(
                                 text.trim_end_matches(|c| c == '\r' || c == '\n'),
                                 &mut cursor_total,
@@ -726,6 +730,12 @@ impl<'a> Printer for InteractivePrinter<'a> {
                                 )
                             )?;
                         }
+
+                        // ANSI escape passthrough.
+                        _ => {
+                            write!(handle, "{}", chunk.raw())?;
+                            self.ansi_style.update(chunk);
+                        }
                     }
                 }
             }
@@ -746,8 +756,8 @@ impl<'a> Printer for InteractivePrinter<'a> {
         }
 
         if highlight_this_line && self.config.theme == "ansi" {
-            self.ansi_style.update("^[24m");
-            write!(handle, "\x1B[24m")?;
+            write!(handle, "{}", ANSI_UNDERLINE_DISABLE.raw())?;
+            self.ansi_style.update(ANSI_UNDERLINE_DISABLE);
         }
 
         Ok(())
diff --git a/src/vscreen.rs b/src/vscreen.rs
index ce7188d7..ea1f02b4 100644
--- a/src/vscreen.rs
+++ b/src/vscreen.rs
@@ -14,7 +14,7 @@ impl AnsiStyle {
         AnsiStyle { attributes: None }
     }
 
-    pub fn update(&mut self, sequence: &str) -> bool {
+    pub fn update(&mut self, sequence: EscapeSequence) -> bool {
         match &mut self.attributes {
             Some(a) => a.update(sequence),
             None => {
@@ -85,26 +85,36 @@ impl Attributes {
 
     /// 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
+    pub fn update(&mut self, sequence: EscapeSequence) -> bool {
+        use EscapeSequence::*;
+        match sequence {
+            Text(_) => return false,
+            Unknown(_) => { /* defer to update_with_unsupported */ }
+            OSC { .. } => return false,
+            CSI {
+                final_byte,
+                parameters,
+                ..
+            } => {
+                match final_byte {
+                    "m" => return self.update_with_sgr(parameters),
+                    _ => {
+                        // NOTE(eth-p): We might want to ignore these, since they involve cursor or buffer manipulation.
+                        /* defer to update_with_unsupported */
                     }
                 }
-                _ => self.update_with_unsupported(sequence),
             }
-        } else {
-            false
+            NF { nf_sequence, .. } => {
+                let mut iter = nf_sequence.chars();
+                match iter.next() {
+                    Some('(') => return self.update_with_charset('(', iter),
+                    Some(')') => return self.update_with_charset(')', iter),
+                    _ => { /* defer to update_with_unsupported */ }
+                }
+            }
         }
+
+        self.update_with_unsupported(sequence.raw())
     }
 
     fn sgr_reset(&mut self) {
@@ -153,14 +163,6 @@ impl Attributes {
         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
@@ -520,49 +522,6 @@ impl<'a> Iterator for EscapeSequenceIterator<'a> {
     }
 }
 
-/// Strips problematic ANSI escape sequences from a string.
-///
-/// Ideally, this will be replaced with something that uses [[Attributes]] to create a table of char offsets
-/// -> absolute styles and style deltas. Something like that would let us simplify the printer (and support
-/// re-printing OSC hyperlink commands).
-pub fn strip_problematic_sequences(text: &str) -> String {
-    use EscapeSequenceOffsets::*;
-
-    let mut buffer = String::with_capacity(text.len());
-    for seq in EscapeSequenceOffsetsIterator::new(text) {
-        match seq {
-            Text { start, end } => buffer.push_str(&text[start..end]),
-            Unknown { start, end } => buffer.push_str(&text[start..end]),
-
-            NF {
-                start_sequence: start,
-                start: _,
-                end,
-            } => buffer.push_str(&text[start..end]),
-
-            CSI {
-                start_sequence: start,
-                start_parameters: _,
-                start_intermediates: _,
-                start_final_byte: _,
-                end,
-            } => buffer.push_str(&text[start..end]),
-
-            OSC {
-                start_sequence: _,
-                start_command: _,
-                start_terminator: _,
-                end: _,
-            } => {
-                // TODO(eth-p): Support re-printing hyperlinks.
-                // In the meantime, strip these.
-            }
-        }
-    }
-
-    buffer
-}
-
 /// A parsed ANSI/VT100 escape sequence.
 #[derive(Debug, PartialEq)]
 pub enum EscapeSequence<'a> {
@@ -601,7 +560,7 @@ impl<'a> EscapeSequence<'a> {
 #[cfg(test)]
 mod tests {
     use crate::vscreen::{
-        strip_problematic_sequences, EscapeSequence, EscapeSequenceIterator, EscapeSequenceOffsets,
+        EscapeSequence, EscapeSequenceIterator, EscapeSequenceOffsets,
         EscapeSequenceOffsetsIterator,
     };
 
@@ -828,14 +787,6 @@ mod tests {
     }
 
     #[test]
-    fn test_strip_problematic_sequences() {
-        assert_eq!(
-            strip_problematic_sequences("text\x1B[33m\x1B]OSC\x1B\\\x1B(0"),
-            "text\x1B[33m\x1B(0"
-        );
-    }
-
-    #[test]
     fn test_escape_sequence_iterator_iterates() {
         let mut iter = EscapeSequenceIterator::new("text\x1B[33m\x1B]OSC\x07\x1B]OSC\x1B\\\x1B(0");
         assert_eq!(iter.next(), Some(EscapeSequence::Text("text")));