fix(grid): glitchy resizes (#2182)

* fix(grid): glitchy resizes

* style(fmt): rustfmt

4 files changed, 63 insertions, 7 deletions

diff --git a/zellij-server/src/panes/grid.rs b/zellij-server/src/panes/grid.rs
index b71687ea2..a3b37722b 100644
--- a/zellij-server/src/panes/grid.rs
+++ b/zellij-server/src/panes/grid.rs
@@ -626,6 +626,27 @@ impl Grid {
         }
         cursor_index_in_canonical_line
     }
+    fn saved_cursor_index_in_canonical_line(&self) -> Option<usize> {
+        if let Some(saved_cursor_position) = self.saved_cursor_position.as_ref() {
+            let mut cursor_canonical_line_index = 0;
+            let mut cursor_index_in_canonical_line = 0;
+            for (i, line) in self.viewport.iter().enumerate() {
+                if line.is_canonical {
+                    cursor_canonical_line_index = i;
+                }
+                if i == saved_cursor_position.y {
+                    let line_wrap_position_in_line =
+                        saved_cursor_position.y - cursor_canonical_line_index;
+                    cursor_index_in_canonical_line =
+                        line_wrap_position_in_line + saved_cursor_position.x;
+                    break;
+                }
+            }
+            Some(cursor_index_in_canonical_line)
+        } else {
+            None
+        }
+    }
     fn canonical_line_y_coordinates(&self, canonical_line_index: usize) -> usize {
         let mut canonical_lines_traversed = 0;
         let mut y_coordinates = 0;
@@ -713,7 +734,7 @@ impl Grid {
         }
         found_something
     }
-    fn force_change_size(&mut self, new_rows: usize, new_columns: usize) {
+    pub fn force_change_size(&mut self, new_rows: usize, new_columns: usize) {
         // this is an ugly hack - it's here because sometimes we need to change_size to the
         // existing size (eg. when resizing an alternative_grid to the current height/width) and
         // the change_size method is a no-op in that case. Should be fixed by making the
@@ -742,6 +763,7 @@ impl Grid {
             self.horizontal_tabstops = create_horizontal_tabstops(new_columns);
             let mut cursor_canonical_line_index = self.cursor_canonical_line_index();
             let cursor_index_in_canonical_line = self.cursor_index_in_canonical_line();
+            let saved_cursor_index_in_canonical_line = self.saved_cursor_index_in_canonical_line();
             let mut viewport_canonical_lines = vec![];
             for mut row in self.viewport.drain(..) {
                 if !row.is_canonical
@@ -817,9 +839,25 @@ impl Grid {
             self.viewport = new_viewport_rows;
 
             let mut new_cursor_y = self.canonical_line_y_coordinates(cursor_canonical_line_index);
+            let mut saved_cursor_y_coordinates =
+                if let Some(saved_cursor) = self.saved_cursor_position.as_ref() {
+                    Some(self.canonical_line_y_coordinates(saved_cursor.y))
+                } else {
+                    None
+                };
 
             let new_cursor_x = (cursor_index_in_canonical_line / new_columns)
                 + (cursor_index_in_canonical_line % new_columns);
+            let saved_cursor_x_coordinates = if let Some(saved_cursor_index_in_canonical_line) =
+                saved_cursor_index_in_canonical_line.as_ref()
+            {
+                Some(
+                    (*saved_cursor_index_in_canonical_line / new_columns)
+                        + (*saved_cursor_index_in_canonical_line % new_columns),
+                )
+            } else {
+                None
+            };
             let current_viewport_row_count = self.viewport.len();
             match current_viewport_row_count.cmp(&self.height) {
                 Ordering::Less => {
@@ -834,6 +872,9 @@ impl Grid {
                     );
                     let rows_pulled = self.viewport.len() - current_viewport_row_count;
                     new_cursor_y += rows_pulled;
+                    if let Some(saved_cursor_y_coordinates) = saved_cursor_y_coordinates.as_mut() {
+                        *saved_cursor_y_coordinates += rows_pulled;
+                    }
                 },
                 Ordering::Greater => {
                     let row_count_to_transfer = current_viewport_row_count - self.height;
@@ -842,6 +883,13 @@ impl Grid {
                     } else {
                         new_cursor_y -= row_count_to_transfer;
                     }
+                    if let Some(saved_cursor_y_coordinates) = saved_cursor_y_coordinates.as_mut() {
+                        if row_count_to_transfer > *saved_cursor_y_coordinates {
+                            *saved_cursor_y_coordinates = 0;
+                        } else {
+                            *saved_cursor_y_coordinates -= row_count_to_transfer;
+                        }
+                    }
                     transfer_rows_from_viewport_to_lines_above(
                         &mut self.viewport,
                         &mut self.lines_above,
@@ -855,8 +903,16 @@ impl Grid {
             self.cursor.y = new_cursor_y;
             self.cursor.x = new_cursor_x;
             if let Some(saved_cursor_position) = self.saved_cursor_position.as_mut() {
-                saved_cursor_position.y = new_cursor_y;
-                saved_cursor_position.x = new_cursor_x;
+                match (saved_cursor_x_coordinates, saved_cursor_y_coordinates) {
+                    (Some(saved_cursor_x_coordinates), Some(saved_cursor_y_coordinates)) => {
+                        saved_cursor_position.x = saved_cursor_x_coordinates;
+                        saved_cursor_position.y = saved_cursor_y_coordinates;
+                    },
+                    _ => {
+                        saved_cursor_position.x = new_cursor_x;
+                        saved_cursor_position.y = new_cursor_y;
+                    },
+                }
             };
         } else if new_columns != self.width && self.alternate_screen_state.is_some() {
             // in alternate screen just truncate exceeding width
diff --git a/zellij-server/src/panes/terminal_pane.rs b/zellij-server/src/panes/terminal_pane.rs
index e9fffb186..be14b74f5 100644
--- a/zellij-server/src/panes/terminal_pane.rs
+++ b/zellij-server/src/panes/terminal_pane.rs
@@ -794,7 +794,7 @@ impl TerminalPane {
     fn reflow_lines(&mut self) {
         let rows = self.get_content_rows();
         let cols = self.get_content_columns();
-        self.grid.change_size(rows, cols);
+        self.grid.force_change_size(rows, cols);
         if self.banner.is_some() {
             self.grid.reset_terminal_state();
             self.render_first_run_banner();
diff --git a/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__save_cursor_position_across_resizes.snap b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__save_cursor_position_across_resizes.snap
index 278b856a2..dbd6426a6 100644
--- a/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__save_cursor_position_across_resizes.snap
+++ b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__save_cursor_position_across_resizes.snap
@@ -1,10 +1,10 @@
 ---
 source: zellij-server/src/tab/./unit/tab_integration_tests.rs
-assertion_line: 1153
+assertion_line: 1952
 expression: snapshot
 ---
 00 (C): ┌ Pane #1 ─────────────────────────────────────────────────────────────────────────── SCROLL:  0/4 ┐
-01 (C): │                                       Let's save the cursor position here this overwrote me!tten │
+01 (C): │Let's save the cursor position here this overwrote me!tten                                        │
 02 (C): └──────────────────────────────────────────────────────────────────────────────────────────────────┘
 03 (C):                                                                                                     
 04 (C):                                                                                                     
diff --git a/zellij-server/src/tab/unit/tab_integration_tests.rs b/zellij-server/src/tab/unit/tab_integration_tests.rs
index 3defa1e6b..89e856444 100644
--- a/zellij-server/src/tab/unit/tab_integration_tests.rs
+++ b/zellij-server/src/tab/unit/tab_integration_tests.rs
@@ -1935,7 +1935,7 @@ fn save_cursor_position_across_resizes() {
 
     tab.handle_pty_bytes(
         1,
-        Vec::from("\n\nI am some text\nI am another line of text\nLet's save the cursor position here \u{1b}[sI should be ovewritten".as_bytes()),
+        Vec::from("\n\n\rI am some text\n\rI am another line of text\n\rLet's save the cursor position here \u{1b}[sI should be ovewritten".as_bytes()),
     ).unwrap();
     tab.resize_whole_tab(Size { cols: 100, rows: 3 }).unwrap();
     tab.handle_pty_bytes(1, Vec::from("\u{1b}[uthis overwrote me!".as_bytes()))