expose some extra internal state to help reproduce line wrapping

adds `row_wrapped` and `cursor_state_formatted` to allow you to better
recreate the internal state of the cursor when using `rows_formatted`.
also make `rows_formatted` keep track of the wrapping state itself,
since there are some edge cases that aren't really able to easily be
tracked externally.

4 files changed, 169 insertions, 97 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c83ecba..a391fda 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,16 @@
 ### Added
 
 * `Screen::alternate_screen` to determine if the alternate screen is in use
+* `Screen::row_wrapped` to determine whether the row at the given index should
+  wrap its text
+* `Screen::cursor_state_formatted` to set the cursor position and hidden state
+  (including internal state like the one-past-the-end state which isn't visible
+  in the return value of `cursor_position`)
+
+### Fixed
+
+* `Screen::rows_formatted` now outputs correct escape codes in some edge cases
+  at the beginning of a row when the previous row was wrapped
 
 ## [0.12.0] - 2021-03-09
 
diff --git a/src/grid.rs b/src/grid.rs
index 3dd8312..aaf79d5 100644
--- a/src/grid.rs
+++ b/src/grid.rs
@@ -223,79 +223,7 @@ impl Grid {
             wrapping = row.wrapped();
         }
 
-        // writing a character to the last column of a row doesn't wrap the
-        // cursor immediately - it waits until the next character is actually
-        // drawn. it is only possible for the cursor to have this kind of
-        // position after drawing a character though, so if we end in this
-        // position, we need to redraw the character at the end of the row.
-        if prev_pos != self.pos && self.pos.col >= self.size.cols {
-            let mut pos = Pos {
-                row: self.pos.row,
-                col: self.size.cols - 1,
-            };
-            if self.visible_cell(pos).unwrap().is_wide_continuation() {
-                pos.col = self.size.cols - 2;
-            }
-            let cell = self.visible_cell(pos).unwrap();
-            if cell.has_contents() {
-                crate::term::MoveFromTo::new(prev_pos, pos)
-                    .write_buf(contents);
-                contents.extend(cell.contents().as_bytes());
-            } else {
-                // if the cell doesn't have contents, we can't have gotten
-                // here by drawing a character in the last column. this means
-                // that as far as i'm aware, we have to have reached here from
-                // a newline when we were already after the end of an earlier
-                // row. in the case where we are already after the end of an
-                // earlier row, we can just write a few newlines, otherwise we
-                // also need to do the same as above to get ourselves to after
-                // the end of a row.
-                let orig_row = pos.row;
-                let mut found = false;
-                for i in (0..orig_row).rev() {
-                    pos.row = i;
-                    pos.col = self.size.cols - 1;
-                    if self.visible_cell(pos).unwrap().is_wide_continuation()
-                    {
-                        pos.col = self.size.cols - 2;
-                    }
-                    let cell = self.visible_cell(pos).unwrap();
-                    if cell.has_contents() {
-                        if prev_pos.row != i || prev_pos.col < self.size.cols
-                        {
-                            crate::term::MoveFromTo::new(prev_pos, pos)
-                                .write_buf(contents);
-                            contents.extend(cell.contents().as_bytes());
-                        }
-                        contents.extend(
-                            "\n".repeat((orig_row - i) as usize).as_bytes(),
-                        );
-                        found = true;
-                        break;
-                    }
-                }
-
-                // this can happen if you get the cursor off the end of a row,
-                // and then do something to clear the end of the current row
-                // without moving the cursor (IL, DL, ED, EL, etc). we know
-                // there can't be something in the last column because we
-                // would have caught that above, so it should be safe to
-                // overwrite it.
-                if !found {
-                    pos.row = orig_row;
-                    crate::term::MoveFromTo::new(prev_pos, pos)
-                        .write_buf(contents);
-                    contents.push(b' ');
-                    crate::term::SaveCursor::default().write_buf(contents);
-                    crate::term::Backspace::default().write_buf(contents);
-                    crate::term::EraseChar::new(1).write_buf(contents);
-                    crate::term::RestoreCursor::default().write_buf(contents);
-                }
-            }
-        } else {
-            crate::term::MoveFromTo::new(prev_pos, self.pos)
-                .write_buf(contents);
-        }
+        self.write_cursor_position_formatted(contents, Some(prev_pos));
 
         prev_attrs
     }
@@ -327,12 +255,22 @@ impl Grid {
             wrapping = row.wrapped();
         }
 
+        self.write_cursor_position_formatted(contents, Some(prev_pos));
+
+        prev_attrs
+    }
+
+    pub fn write_cursor_position_formatted(
+        &self,
+        contents: &mut Vec<u8>,
+        prev_pos: Option<Pos>,
+    ) {
         // writing a character to the last column of a row doesn't wrap the
         // cursor immediately - it waits until the next character is actually
         // drawn. it is only possible for the cursor to have this kind of
         // position after drawing a character though, so if we end in this
         // position, we need to redraw the character at the end of the row.
-        if prev_pos != self.pos && self.pos.col >= self.size.cols {
+        if prev_pos != Some(self.pos) && self.pos.col >= self.size.cols {
             let mut pos = Pos {
                 row: self.pos.row,
                 col: self.size.cols - 1,
@@ -342,8 +280,12 @@ impl Grid {
             }
             let cell = self.visible_cell(pos).unwrap();
             if cell.has_contents() {
-                crate::term::MoveFromTo::new(prev_pos, pos)
-                    .write_buf(contents);
+                if let Some(prev_pos) = prev_pos {
+                    crate::term::MoveFromTo::new(prev_pos, pos)
+                        .write_buf(contents);
+                } else {
+                    crate::term::MoveTo::new(pos).write_buf(contents);
+                }
                 contents.extend(cell.contents().as_bytes());
             } else {
                 // if the cell doesn't have contents, we can't have gotten
@@ -365,10 +307,16 @@ impl Grid {
                     }
                     let cell = self.visible_cell(pos).unwrap();
                     if cell.has_contents() {
-                        if prev_pos.row != i || prev_pos.col < self.size.cols
-                        {
-                            crate::term::MoveFromTo::new(prev_pos, pos)
-                                .write_buf(contents);
+                        if let Some(prev_pos) = prev_pos {
+                            if prev_pos.row != i
+                                || prev_pos.col < self.size.cols
+                            {
+                                crate::term::MoveFromTo::new(prev_pos, pos)
+                                    .write_buf(contents);
+                                contents.extend(cell.contents().as_bytes());
+                            }
+                        } else {
+                            crate::term::MoveTo::new(pos).write_buf(contents);
                             contents.extend(cell.contents().as_bytes());
                         }
                         contents.extend(
@@ -387,8 +335,12 @@ impl Grid {
                 // overwrite it.
                 if !found {
                     pos.row = orig_row;
-                    crate::term::MoveFromTo::new(prev_pos, pos)
-                        .write_buf(contents);
+                    if let Some(prev_pos) = prev_pos {
+                        crate::term::MoveFromTo::new(prev_pos, pos)
+                            .write_buf(contents);
+                    } else {
+                        crate::term::MoveTo::new(pos).write_buf(contents);
+                    }
                     contents.push(b' ');
                     crate::term::SaveCursor::default().write_buf(contents);
                     crate::term::Backspace::default().write_buf(contents);
@@ -396,12 +348,12 @@ impl Grid {
                     crate::term::RestoreCursor::default().write_buf(contents);
                 }
             }
-        } else {
+        } else if let Some(prev_pos) = prev_pos {
             crate::term::MoveFromTo::new(prev_pos, self.pos)
                 .write_buf(contents);
+        } else {
+            crate::term::MoveTo::new(self.pos).write_buf(contents);
         }
-
-        prev_attrs
     }
 
     pub fn erase_all(&mut self, attrs: crate::attrs::Attrs) {
diff --git a/src/screen.rs b/src/screen.rs
index 83677ee..c41184c 100644
--- a/src/screen.rs
+++ b/src/screen.rs
@@ -285,6 +285,7 @@ impl Screen {
         start: u16,
         width: u16,
     ) -> impl Iterator<Item = Vec<u8>> + '_ {
+        let mut wrapping = false;
         self.grid().visible_rows().enumerate().map(move |(i, row)| {
             let i = i.try_into().unwrap();
             let mut contents = vec![];
@@ -293,10 +294,11 @@ impl Screen {
                 start,
                 width,
                 i,
-                false,
+                wrapping,
                 crate::grid::Pos { row: i, col: start },
                 crate::attrs::Attrs::default(),
             );
+            wrapping = row.wrapped();
             contents
         })
     }
@@ -532,6 +534,40 @@ impl Screen {
         );
     }
 
+    /// Returns the current cursor position of the terminal.
+    ///
+    /// The return value will be (row, col).
+    #[must_use]
+    pub fn cursor_position(&self) -> (u16, u16) {
+        let pos = self.grid().pos();
+        (pos.row, pos.col)
+    }
+
+    /// Returns terminal escape sequences sufficient to set the current
+    /// cursor state of the terminal.
+    ///
+    /// This is not typically necessary, since `contents_formatted` will leave
+    /// the cursor in the correct state, but this can be useful in the case of
+    /// drawing additional things on top of a terminal output, since you will
+    /// need to restore the terminal state without the terminal contents
+    /// necessarily being the same.
+    ///
+    /// Note that this is more complicated than it sounds, because the cursor
+    /// position's state includes more than just the data returned by
+    /// `cursor_position`, since moving the cursor to the next row during text
+    /// wrapping is delayed until a character is written.
+    #[must_use]
+    pub fn cursor_state_formatted(&self) -> Vec<u8> {
+        let mut contents = vec![];
+        self.write_cursor_state_formatted(&mut contents);
+        contents
+    }
+
+    fn write_cursor_state_formatted(&self, contents: &mut Vec<u8>) {
+        crate::term::HideCursor::new(self.hide_cursor()).write_buf(contents);
+        self.grid().write_cursor_position_formatted(contents, None);
+    }
+
     /// Returns the `Cell` object at the given location in the terminal, if it
     /// exists.
     #[must_use]
@@ -539,13 +575,12 @@ impl Screen {
         self.grid().visible_cell(crate::grid::Pos { row, col })
     }
 
-    /// Returns the current cursor position of the terminal.
-    ///
-    /// The return value will be (row, col).
+    /// Returns whether the text in row `row` should wrap to the next line.
     #[must_use]
-    pub fn cursor_position(&self) -> (u16, u16) {
-        let pos = self.grid().pos();
-        (pos.row, pos.col)
+    pub fn row_wrapped(&self, row: u16) -> bool {
+        self.grid()
+            .visible_row(crate::grid::Pos { row, col: 0 })
+            .map_or(false, crate::row::Row::wrapped)
     }
 
     /// Returns the terminal's window title.
diff --git a/tests/helpers/mod.rs b/tests/helpers/mod.rs
index daf18dd..17f9a60 100644
--- a/tests/helpers/mod.rs
+++ b/tests/helpers/mod.rs
@@ -1,3 +1,5 @@
+use std::convert::TryInto as _;
+
 mod fixtures;
 pub use fixtures::fixture;
 pub use fixtures::FixtureScreen;
@@ -21,18 +23,55 @@ macro_rules! ok {
     };
 }
 
+#[derive(Eq, PartialEq)]
+struct Bytes<'a>(&'a [u8]);
+
+impl<'a> std::fmt::Debug for Bytes<'a> {
+    fn fmt(
+        &self,
+        f: &mut std::fmt::Formatter<'_>,
+    ) -> Result<(), std::fmt::Error> {
+        f.write_str("b\"")?;
+        for c in self.0 {
+            match c {
+                10 => f.write_str("\\n")?,
+                13 => f.write_str("\\r")?,
+                92 => f.write_str("\\\\")?,
+                32..=126 => f.write_str(&char::from(*c).to_string())?,
+                _ => f.write_fmt(format_args!("\\x{:02x}", c))?,
+            }
+        }
+        f.write_str("\"")?;
+        Ok(())
+    }
+}
+
 pub fn compare_screens(
     got: &vt100::Screen,
     expected: &vt100::Screen,
 ) -> bool {
+    let (rows, cols) = got.size();
+
     is!(got.contents(), expected.contents());
-    is!(got.contents_formatted(), expected.contents_formatted());
     is!(
-        got.contents_diff(vt100::Parser::default().screen()),
-        expected.contents_diff(vt100::Parser::default().screen())
+        Bytes(&got.contents_formatted()),
+        Bytes(&expected.contents_formatted())
+    );
+    for (got_row, expected_row) in
+        got.rows(0, cols).zip(expected.rows(0, cols))
+    {
+        is!(got_row, expected_row);
+    }
+    for (got_row, expected_row) in got
+        .rows_formatted(0, cols)
+        .zip(expected.rows_formatted(0, cols))
+    {
+        is!(Bytes(&got_row), Bytes(&expected_row));
+    }
+    is!(
+        Bytes(&got.contents_diff(vt100::Parser::default().screen())),
+        Bytes(&expected.contents_diff(vt100::Parser::default().screen()))
     );
-
-    let (rows, cols) = got.size();
 
     for row in 0..rows {
         for col in 0..cols {
@@ -80,6 +119,12 @@ pub fn contents_formatted_reproduces_state(input: &[u8]) -> bool {
     contents_formatted_reproduces_screen(parser.screen())
 }
 
+pub fn rows_formatted_reproduces_state(input: &[u8]) -> bool {
+    let mut parser = vt100::Parser::default();
+    parser.process(input);
+    rows_formatted_reproduces_screen(parser.screen())
+}
+
 pub fn contents_formatted_reproduces_screen(screen: &vt100::Screen) -> bool {
     let empty_screen = vt100::Parser::default().screen().clone();
 
@@ -95,10 +140,39 @@ pub fn contents_formatted_reproduces_screen(screen: &vt100::Screen) -> bool {
     compare_screens(&got_screen, screen)
 }
 
+pub fn rows_formatted_reproduces_screen(screen: &vt100::Screen) -> bool {
+    let empty_screen = vt100::Parser::default().screen().clone();
+
+    let mut new_input = vec![];
+    let mut wrapped = false;
+    for (idx, row) in screen.rows_formatted(0, 80).enumerate() {
+        new_input.extend(b"\x1b[m");
+        if !wrapped {
+            new_input.extend(format!("\x1b[{}H", idx + 1).as_bytes());
+        }
+        new_input.extend(row);
+        wrapped = screen.row_wrapped(idx.try_into().unwrap());
+    }
+    new_input.extend(screen.cursor_state_formatted());
+    new_input.extend(screen.attributes_formatted());
+    new_input.extend(screen.input_mode_formatted());
+    new_input.extend(screen.title_formatted());
+    new_input.extend(screen.bells_diff(&empty_screen));
+    let mut new_parser = vt100::Parser::default();
+    new_parser.process(&new_input);
+    let got_screen = new_parser.screen().clone();
+
+    compare_screens(&got_screen, screen)
+}
+
 fn assert_contents_formatted_reproduces_state(input: &[u8]) {
     assert!(contents_formatted_reproduces_state(input));
 }
 
+fn assert_rows_formatted_reproduces_state(input: &[u8]) {
+    assert!(rows_formatted_reproduces_state(input));
+}
+
 #[allow(dead_code)]
 pub fn contents_diff_reproduces_state(input: &[u8]) -> bool {
     contents_diff_reproduces_state_from(input, &[])
@@ -167,6 +241,7 @@ pub fn assert_reproduces_state_from(input: &[u8], prev_input: &[u8]) {
     let full_input: Vec<_> =
         prev_input.iter().chain(input.iter()).copied().collect();
     assert_contents_formatted_reproduces_state(&full_input);
+    assert_rows_formatted_reproduces_state(&full_input);
     assert_contents_diff_reproduces_state_from(input, prev_input);
 }