contents_formatted and contents_diff should also restore the cursor

since the cursor is a visible part of the terminal

3 files changed, 54 insertions, 14 deletions

diff --git a/src/grid.rs b/src/grid.rs
index fb0f7d7..218247e 100644
--- a/src/grid.rs
+++ b/src/grid.rs
@@ -139,9 +139,13 @@ impl Grid {
     pub fn contents_formatted(&self) -> Vec<u8> {
         let mut contents = vec![];
         let mut prev_attrs = crate::attrs::Attrs::default();
+        let mut final_col = 0;
         for row in self.rows() {
-            let (mut new_contents, new_attrs) =
+            let (mut new_contents, new_attrs, new_col) =
                 row.contents_formatted(0, self.size.cols, prev_attrs);
+            if !new_contents.is_empty() {
+                final_col = new_col;
+            }
             contents.append(&mut new_contents);
             if !row.wrapped() {
                 contents.extend(b"\r\n");
@@ -149,26 +153,49 @@ impl Grid {
             prev_attrs = new_attrs;
         }
 
+        let mut final_row = self.size.rows;
         while contents.ends_with(b"\r\n") {
             contents.truncate(contents.len() - 2);
+            final_row -= 1;
         }
+
+        if final_row != self.pos.row || final_col != self.pos.col {
+            contents.extend(
+                format!("\x1b[{};{}H", self.pos.row + 1, self.pos.col + 1)
+                    .as_bytes(),
+            );
+        }
+
         contents
     }
 
     pub fn contents_diff(&self, prev: &Self) -> Vec<u8> {
         let mut contents = b"\x1b[m".to_vec();
         let mut prev_attrs = crate::attrs::Attrs::default();
+        let mut final_row = prev.pos.row;
+        let mut final_col = prev.pos.col;
         for (idx, (row, prev_row)) in self.rows().zip(prev.rows()).enumerate()
         {
-            let (mut new_contents, new_attrs) = row.contents_diff(
+            let (mut new_contents, new_attrs, new_col) = row.contents_diff(
                 idx.try_into().unwrap(),
                 prev_row,
                 prev_attrs,
             );
+            if !new_contents.is_empty() {
+                final_row = idx.try_into().unwrap();
+                final_col = new_col;
+            }
             contents.append(&mut new_contents);
             prev_attrs = new_attrs;
         }
 
+        if self.pos.row != final_row || self.pos.col != final_col {
+            contents.extend(
+                format!("\x1b[{};{}H", self.pos.row + 1, self.pos.col + 1)
+                    .as_bytes(),
+            );
+        }
+
         contents
     }
 
diff --git a/src/row.rs b/src/row.rs
index 5f2f775..cf47374 100644
--- a/src/row.rs
+++ b/src/row.rs
@@ -93,16 +93,13 @@ impl Row {
         start: u16,
         width: u16,
         attrs: crate::attrs::Attrs,
-    ) -> (Vec<u8>, crate::attrs::Attrs) {
+    ) -> (Vec<u8>, crate::attrs::Attrs, u16) {
         let mut prev_was_wide = false;
         let mut contents = vec![];
         let mut prev_attrs = attrs;
 
-        for cell in self
-            .cells()
-            .skip(start as usize)
-            .take(width.min(self.content_width(start)) as usize)
-        {
+        let cols = width.min(self.content_width(start));
+        for cell in self.cells().skip(start as usize).take(cols as usize) {
             if prev_was_wide {
                 prev_was_wide = false;
                 continue;
@@ -123,7 +120,7 @@ impl Row {
             prev_was_wide = cell.is_wide();
         }
 
-        (contents, prev_attrs)
+        (contents, prev_attrs, cols)
     }
 
     pub fn contents_diff(
@@ -131,10 +128,11 @@ impl Row {
         row_idx: u16,
         prev: &Self,
         attrs: crate::attrs::Attrs,
-    ) -> (Vec<u8>, crate::attrs::Attrs) {
+    ) -> (Vec<u8>, crate::attrs::Attrs, u16) {
         let mut needs_move = true;
         let mut contents = vec![];
         let mut prev_attrs = attrs;
+        let mut final_col = 0;
         for (idx, (cell, prev_cell)) in
             self.cells().zip(prev.cells()).enumerate()
         {
@@ -160,10 +158,11 @@ impl Row {
                 } else {
                     b"\x1b[X\x1b[C"
                 });
+                final_col = idx + 1;
             }
         }
 
-        (contents, prev_attrs)
+        (contents, prev_attrs, final_col.try_into().unwrap())
     }
 
     fn content_width(&self, start: u16) -> u16 {
diff --git a/src/screen.rs b/src/screen.rs
index 7d6fca6..5a3a957 100644
--- a/src/screen.rs
+++ b/src/screen.rs
@@ -155,7 +155,12 @@ impl Screen {
     /// terminal modes (such as application keypad mode or alternate screen
     /// mode) will not be included here.
     pub fn contents_formatted(&self) -> Vec<u8> {
-        self.grid().contents_formatted()
+        let mut grid_contents = vec![];
+        if self.hide_cursor() {
+            grid_contents.extend(b"\x1b[?25l");
+        }
+        grid_contents.append(&mut self.grid().contents_formatted());
+        grid_contents
     }
 
     /// Returns the formatted contents of the terminal by row, restricted to
@@ -174,7 +179,7 @@ impl Screen {
         width: u16,
     ) -> impl Iterator<Item = Vec<u8>> + '_ {
         self.grid().rows().map(move |row| {
-            let (contents, _) = row.contents_formatted(
+            let (contents, ..) = row.contents_formatted(
                 start,
                 width,
                 crate::attrs::Attrs::default(),
@@ -184,7 +189,16 @@ impl Screen {
     }
 
     pub fn contents_diff(&self, prev: &Self) -> Vec<u8> {
-        self.grid().contents_diff(prev.grid())
+        let mut grid_contents = vec![];
+        if self.hide_cursor() != prev.hide_cursor() {
+            grid_contents.extend(if self.hide_cursor() {
+                b"\x1b[?25l"
+            } else {
+                b"\x1b[?25h"
+            });
+        }
+        grid_contents.append(&mut self.grid().contents_diff(prev.grid()));
+        grid_contents
     }
 
     /// Returns the `Cell` object at the given location in the terminal, if it