fix(stacked-panes): properly handle closing a pane near a stack

3 files changed, 91 insertions, 2 deletions

diff --git a/zellij-server/src/panes/tiled_panes/tiled_pane_grid.rs b/zellij-server/src/panes/tiled_panes/tiled_pane_grid.rs
index 1cde6d6fd..8c9e46e71 100644
--- a/zellij-server/src/panes/tiled_panes/tiled_pane_grid.rs
+++ b/zellij-server/src/panes/tiled_panes/tiled_pane_grid.rs
@@ -1032,8 +1032,14 @@ impl<'a> TiledPaneGrid<'a> {
         pane_ids.iter().fold(HashSet::new(), |mut borders, p| {
             let panes = self.panes.borrow();
             let pane = panes.get(p).unwrap();
-            borders.insert(pane.y());
-            borders.insert(pane.y() + pane.rows());
+            if pane.position_and_size().is_stacked {
+                let pane_geom = StackedPanes::new(self.panes.clone()).position_and_size_of_stack(&pane.pid()).unwrap();
+                borders.insert(pane_geom.y);
+                borders.insert(pane_geom.y + pane_geom.rows.as_usize());
+            } else {
+                borders.insert(pane.y());
+                borders.insert(pane.y() + pane.rows());
+            }
             borders
         })
     }
diff --git a/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__close_pane_near_stacked_panes.snap b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__close_pane_near_stacked_panes.snap
new file mode 100644
index 000000000..631815e7c
--- /dev/null
+++ b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__close_pane_near_stacked_panes.snap
@@ -0,0 +1,26 @@
+---
+source: zellij-server/src/tab/./unit/tab_integration_tests.rs
+assertion_line: 4536
+expression: snapshot
+---
+00 (C): ┌ Pane #1 ─────────────────────────────────────────────────┐┌ Pane #2 ──────────────────────────────────────────────────┐
+01 (C): │                                                          │┌ Pane #3 ──────────────────────────────────────────────────┐
+02 (C): │                                                          │┌ Pane #4 ──────────────────────────────────────────────────┐
+03 (C): │                                                          │┌ Pane #5 ──────────────────────────────────────────────────┐
+04 (C): │                                                          ││                                                           │
+05 (C): │                                                          ││                                                           │
+06 (C): │                                                          ││                                                           │
+07 (C): │                                                          ││                                                           │
+08 (C): │                                                          ││                                                           │
+09 (C): │                                                          ││                                                           │
+10 (C): │                                                          ││                                                           │
+11 (C): │                                                          ││                                                           │
+12 (C): │                                                          ││                                                           │
+13 (C): │                                                          ││                                                           │
+14 (C): │                                                          ││                                                           │
+15 (C): │                                                          ││                                                           │
+16 (C): │                                                          ││                                                           │
+17 (C): │                                                          ││                                                           │
+18 (C): │                                                          ││                                                           │
+19 (C): └──────────────────────────────────────────────────────────┘└───────────────────────────────────────────────────────────┘
+
diff --git a/zellij-server/src/tab/unit/tab_integration_tests.rs b/zellij-server/src/tab/unit/tab_integration_tests.rs
index 683e5bc56..fed42faba 100644
--- a/zellij-server/src/tab/unit/tab_integration_tests.rs
+++ b/zellij-server/src/tab/unit/tab_integration_tests.rs
@@ -4480,6 +4480,63 @@ fn close_stacked_pane_with_previously_focused_other_pane() {
 }
 
 #[test]
+fn close_pane_near_stacked_panes() {
+    let size = Size {
+        cols: 121,
+        rows: 20,
+    };
+    let client_id = 1;
+    let mut output = Output::default();
+    let swap_layouts = r#"
+        layout {
+            swap_tiled_layout {
+                tab {
+                    pane split_direction="vertical" {
+                        pane focus=true
+                        pane { children stacked=true; }
+                    }
+                }
+            }
+        }
+    "#;
+    let layout = Layout::from_kdl(swap_layouts, "file_name.kdl".into(), None).unwrap();
+    let swap_tiled_layouts = layout.swap_tiled_layouts.clone();
+    let swap_floating_layouts = layout.swap_floating_layouts.clone();
+    let mut tab = create_new_tab_with_swap_layouts(size, ModeInfo::default(), (swap_tiled_layouts, swap_floating_layouts));
+    let new_pane_id_1 = PaneId::Terminal(2);
+    let new_pane_id_2 = PaneId::Terminal(3);
+    let new_pane_id_3 = PaneId::Terminal(4);
+    let new_pane_id_4 = PaneId::Terminal(5);
+    let new_pane_id_5 = PaneId::Terminal(6);
+
+    tab.new_pane(new_pane_id_1, None, None, Some(client_id))
+        .unwrap();
+    tab.new_pane(new_pane_id_2, None, None, Some(client_id))
+        .unwrap();
+    tab.new_pane(new_pane_id_3, None, None, Some(client_id))
+        .unwrap();
+    tab.new_pane(new_pane_id_4, None, None, Some(client_id))
+        .unwrap();
+    tab.new_pane(new_pane_id_5, None, None, Some(client_id))
+        .unwrap();
+    tab.close_pane(PaneId::Terminal(6), false, None);
+    tab.render(&mut output, None).unwrap();
+    let (snapshot, cursor_coordinates) = take_snapshot_and_cursor_position(
+        output.serialize().unwrap().get(&client_id).unwrap(),
+        size.rows,
+        size.cols,
+        Palette::default(),
+    );
+    assert_eq!(
+        cursor_coordinates,
+        Some((61, 4)),
+        "cursor coordinates moved to the main pane of the stack",
+    );
+
+    assert_snapshot!(snapshot);
+}
+
+#[test]
 fn focus_next_pane_expands_stacked_panes() {
     let size = Size {
         cols: 121,