Refactoring ansi/iterator (#1191)

Signed-off-by: Maxim Zhiburt <zhiburt@gmail.com>

1 files changed, 63 insertions, 65 deletions

diff --git a/src/ansi/iterator.rs b/src/ansi/iterator.rs
index f9952162..6318fd51 100644
--- a/src/ansi/iterator.rs
+++ b/src/ansi/iterator.rs
@@ -24,6 +24,7 @@ pub struct AnsiElementIterator<'a> {
     pos: usize,
 }
 
+#[derive(Default)]
 struct Performer {
     // Becomes non-None when the parser finishes parsing an ANSI sequence.
     // This is never Element::Text.
@@ -42,6 +43,21 @@ pub enum Element {
     Text(usize, usize),
 }
 
+impl Element {
+    fn set_range(&mut self, start: usize, end: usize) {
+        let (from, to) = match self {
+            Element::Sgr(_, from, to) => (from, to),
+            Element::Csi(from, to) => (from, to),
+            Element::Esc(from, to) => (from, to),
+            Element::Osc(from, to) => (from, to),
+            Element::Text(from, to) => (from, to),
+        };
+
+        *from = start;
+        *to = end;
+    }
+}
+
 impl<'a> AnsiElementIterator<'a> {
     pub fn new(s: &'a str) -> Self {
         Self {
@@ -54,17 +70,12 @@ impl<'a> AnsiElementIterator<'a> {
         }
     }
 
-    #[allow(dead_code)]
-    pub fn dbg(s: &str) {
-        for el in AnsiElementIterator::new(s) {
-            match el {
-                Element::Sgr(_, i, j) => println!("SGR({}, {}, {:?})", i, j, &s[i..j]),
-                Element::Csi(i, j) => println!("CSI({}, {}, {:?})", i, j, &s[i..j]),
-                Element::Esc(i, j) => println!("ESC({}, {}, {:?})", i, j, &s[i..j]),
-                Element::Osc(i, j) => println!("OSC({}, {}, {:?})", i, j, &s[i..j]),
-                Element::Text(i, j) => println!("Text({}, {}, {:?})", i, j, &s[i..j]),
-            }
-        }
+    fn advance_vte(&mut self, byte: u8) {
+        let mut performer = Performer::default();
+        self.machine.advance(&mut performer, byte);
+        self.element = performer.element;
+        self.text_length += performer.text_length;
+        self.pos += 1;
     }
 }
 
@@ -72,53 +83,39 @@ impl<'a> Iterator for AnsiElementIterator<'a> {
     type Item = Element;
 
     fn next(&mut self) -> Option<Element> {
-        loop {
-            // If the last element emitted was text, then there may be a non-text element waiting
-            // to be emitted. In that case we do not consume a new byte.
-            let byte = if self.element.is_some() {
-                None
-            } else {
-                self.bytes.next()
-            };
-            if byte.is_some() || self.element.is_some() {
-                if let Some(byte) = byte {
-                    let mut performer = Performer {
-                        element: None,
-                        text_length: 0,
-                    };
-                    self.machine.advance(&mut performer, byte);
-                    self.element = performer.element;
-                    self.text_length += performer.text_length;
-                    self.pos += 1;
-                }
-                if self.element.is_some() {
-                    // There is a non-text element waiting to be emitted, but it may have preceding
-                    // text, which must be emitted first.
-                    if self.text_length > 0 {
-                        let start = self.start;
-                        self.start += self.text_length;
-                        self.text_length = 0;
-                        return Some(Element::Text(start, self.start));
-                    }
-                    let start = self.start;
-                    self.start = self.pos;
-                    let element = match self.element.as_ref().unwrap() {
-                        Element::Sgr(style, _, _) => Element::Sgr(*style, start, self.pos),
-                        Element::Csi(_, _) => Element::Csi(start, self.pos),
-                        Element::Esc(_, _) => Element::Esc(start, self.pos),
-                        Element::Osc(_, _) => Element::Osc(start, self.pos),
-                        Element::Text(_, _) => unreachable!(),
-                    };
-                    self.element = None;
-                    return Some(element);
-                }
-            } else if self.text_length > 0 {
+        // If the last element emitted was text, then there may be a non-text element waiting
+        // to be emitted. In that case we do not consume a new byte.
+        while self.element.is_none() {
+            match self.bytes.next() {
+                Some(b) => self.advance_vte(b),
+                None => break,
+            }
+        }
+
+        if let Some(mut element) = self.element.take() {
+            // There is a non-text element waiting to be emitted, but it may have preceding
+            // text, which must be emitted first.
+            if self.text_length > 0 {
+                let start = self.start;
+                self.start += self.text_length;
                 self.text_length = 0;
-                return Some(Element::Text(self.start, self.pos));
-            } else {
-                return None;
+                self.element = Some(element);
+                return Some(Element::Text(start, self.start));
             }
+
+            let start = self.start;
+            self.start = self.pos;
+            element.set_range(start, self.pos);
+
+            return Some(element);
         }
+
+        if self.text_length > 0 {
+            self.text_length = 0;
+            return Some(Element::Text(self.start, self.pos));
+        }
+
+        None
     }
 }
 
@@ -129,20 +126,21 @@ impl vte::Perform for Performer {
             return;
         }
 
-        if let ('m', None) = (c, intermediates.first()) {
+        let is_sgr = c == 'm' && intermediates.first().is_none();
+        let element = if is_sgr {
             if params.is_empty() {
                 // Attr::Reset
                 // Probably doesn't need to be handled: https://github.com/dandavison/delta/pull/431#discussion_r536883568
+                None
             } else {
-                self.element = Some(Element::Sgr(
-                    ansi_term_style_from_sgr_parameters(&mut params.iter()),
-                    0,
-                    0,
-                ));
+                let style = ansi_term_style_from_sgr_parameters(&mut params.iter());
+                Some(Element::Sgr(style, 0, 0))
             }
         } else {
-            self.element = Some(Element::Csi(0, 0));
-        }
+            Some(Element::Csi(0, 0))
+        };
+
+        self.element = element;
     }
 
     fn print(&mut self, c: char) {
@@ -311,7 +309,7 @@ mod tests {
                     style,
                     style::Style::from_git_str(git_style_string).ansi_term_style
                 )),
-                _ => assert!(false),
+                _ => unreachable!(),
             }
 
             // Second element should be text: "+"
@@ -331,7 +329,7 @@ mod tests {
                     style,
                     style::Style::from_git_str(git_style_string).ansi_term_style
                 )),
-                _ => assert!(false),
+                _ => unreachable!(),
             }
 
             // Fifth element should be text: "text"