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.

2 files changed, 79 insertions, 92 deletions

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.