contents_formatted and contents_diff should also restore the cursor

since the cursor is a visible part of the terminal

5 files changed, 68 insertions, 22 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
diff --git a/tests/escape.rs b/tests/escape.rs
index 01ac574..fcd4e6b 100644
--- a/tests/escape.rs
+++ b/tests/escape.rs
@@ -64,7 +64,7 @@ fn ris() {
     assert_eq!(parser.screen().contents(), "foo");
     assert_eq!(
         parser.screen().contents_formatted(),
-        b"f\x1b[31;47;1;3;4moo"
+        b"\x1b[?25lf\x1b[31;47;1;3;4moo\x1b[21;21H"
     );
 
     assert_eq!(parser.screen().title(), "window title");
@@ -172,13 +172,13 @@ fn decsc() {
     assert_eq!(parser.screen().cursor_position(), (0, 3));
     assert_eq!(
         parser.screen().contents_formatted(),
-        b"\x1b[32mbar\r\n\r\n\r\n\r\n\x1b[31mfoo"
+        b"\x1b[32mbar\r\n\r\n\r\n\r\n\x1b[31mfoo\x1b[1;4H"
     );
 
     parser.process(b"\x1b8\x1b[Hz");
     assert_eq!(parser.screen().cursor_position(), (4, 1));
     assert_eq!(
         parser.screen().contents_formatted(),
-        b"\x1b[32mbar\r\n\r\n\r\n\r\n\x1b[31mzoo"
+        b"\x1b[32mbar\r\n\r\n\r\n\r\n\x1b[31mzoo\x1b[5;2H"
     );
 }
diff --git a/tests/window_contents.rs b/tests/window_contents.rs
index 235c05e..ed15517 100644
--- a/tests/window_contents.rs
+++ b/tests/window_contents.rs
@@ -22,7 +22,7 @@ fn formatted() {
     assert!(!parser.screen().cell(0, 5).unwrap().bold());
     assert_eq!(
         parser.screen().contents_formatted(),
-        b"foo\x1b[33;1;7mb\x1b[mar"
+        b"foo\x1b[33;1;7mb\x1b[mar\x1b[1;5H"
     );
 
     parser.process(b"\x1b[1;5H\x1b[22;42ma");
@@ -33,7 +33,7 @@ fn formatted() {
     assert!(!parser.screen().cell(0, 5).unwrap().bold());
     assert_eq!(
         parser.screen().contents_formatted(),
-        b"foo\x1b[33;1;7mb\x1b[42;22ma\x1b[mr"
+        b"foo\x1b[33;1;7mb\x1b[42;22ma\x1b[mr\x1b[1;6H"
     );
 
     parser.process(b"\x1b[1;6H\x1b[35mr\r\nquux");
@@ -64,7 +64,7 @@ fn empty_cells() {
     assert_eq!(parser.screen().contents(), "foo   bar");
     assert_eq!(
         parser.screen().contents_formatted(),
-        b"\x1b[31mfoo\x1b[m\x1b[C\x1b[C\x1b[32m bar"
+        b"\x1b[31mfoo\x1b[m\x1b[C\x1b[C\x1b[32m bar\x1b[1;4H"
     );
 }
 
@@ -349,7 +349,7 @@ fn diff() {
     let screen5 = parser.screen().clone();
     assert_eq!(
         screen5.contents_diff(&screen4),
-        b"\x1b[m\x1b[1;8H\x1b[X\x1b[C"
+        b"\x1b[m\x1b[1;8H\x1b[X\x1b[C\x1b[1;8H"
     );
     compare_diff(
         &screen4,
@@ -385,7 +385,6 @@ fn diff_crawl(i: usize) {
 
     let mut all_frames: Vec<u8> = vec![];
     for two_screens in screens.windows(2) {
-        eprintln!("loop");
         match two_screens {
             [(prev_frame, prev_screen), (_, screen)] => {
                 all_frames.extend(prev_frame);
@@ -419,6 +418,11 @@ fn compare_diff(
     let mut parser = vt100::Parser::new(rows, cols);
     parser.process(&contents);
     compare_cells(parser.screen(), &prev_screen);
+    assert_eq!(parser.screen().hide_cursor(), prev_screen.hide_cursor());
+    assert_eq!(
+        parser.screen().cursor_position(),
+        prev_screen.cursor_position()
+    );
     assert_eq!(
         parser.screen().contents_formatted(),
         prev_screen.contents_formatted()
@@ -440,6 +444,8 @@ fn compare_diff(
             .unwrap();
     }
     compare_cells(parser.screen(), &screen);
+    assert_eq!(parser.screen().hide_cursor(), screen.hide_cursor());
+    assert_eq!(parser.screen().cursor_position(), screen.cursor_position());
     assert_eq!(
         parser.screen().contents_formatted(),
         screen.contents_formatted()