fix: focusing into stacked panes from the left/right

5 files changed, 186 insertions, 7 deletions

diff --git a/zellij-server/src/panes/tiled_panes/stacked_panes.rs b/zellij-server/src/panes/tiled_panes/stacked_panes.rs
index a259ac9fe..472c9e009 100644
--- a/zellij-server/src/panes/tiled_panes/stacked_panes.rs
+++ b/zellij-server/src/panes/tiled_panes/stacked_panes.rs
@@ -44,6 +44,10 @@ impl <'a>StackedPanes <'a>{
         destination_pane.set_geom(destination_pane_geom);
         Ok(())
     }
+    pub fn flexible_pane_id_in_stack(&self, pane_id_in_stack: &PaneId) -> Option<PaneId> {
+        let all_stacked_pane_positions = self.positions_in_stack(pane_id_in_stack).ok()?;
+        all_stacked_pane_positions.iter().find(|(_pid, p)| p.rows.is_percent()).map(|(pid, _p)| *pid)
+    }
     fn pane_is_one_liner(&self, id: &PaneId) -> Result<bool> {
         let err_context = || format!("Cannot determin if pane is one liner or not");
         let panes = self.panes.borrow();
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 1706f1d39..8d4e3602f 100644
--- a/zellij-server/src/panes/tiled_panes/tiled_pane_grid.rs
+++ b/zellij-server/src/panes/tiled_panes/tiled_pane_grid.rs
@@ -867,7 +867,7 @@ impl<'a> TiledPaneGrid<'a> {
             .filter(|(_, p)| p.selectable())
             .map(|(p_id, p)| (*p_id, p))
             .collect();
-        let next_index = panes
+        let next_pane = panes
             .iter()
             .enumerate()
             .filter(|(_, (_, c))| {
@@ -875,9 +875,14 @@ impl<'a> TiledPaneGrid<'a> {
                     && c.horizontally_overlaps_with(Box::as_ref(current_pane))
             })
             .max_by_key(|(_, (_, c))| c.active_at())
-            .map(|(_, (pid, _))| pid)
-            .copied();
-        next_index
+            .map(|(_, (_, pane))| pane);
+        let next_pane_is_stacked = next_pane.map(|p| p.position_and_size().is_stacked).unwrap_or(false);
+        if next_pane_is_stacked {
+            if let Some(next_pane_id) = next_pane.map(|p| p.pid()) {
+                return StackedPanes::new(self.panes.clone()).flexible_pane_id_in_stack(&next_pane_id);
+            }
+        }
+        next_pane.map(|p| p.pid())
     }
     pub fn progress_stack_up_if_in_stack(&mut self, source_pane_id: &PaneId) -> Option<PaneId> {
         let destination_pane_id_in_stack = {
@@ -991,7 +996,7 @@ impl<'a> TiledPaneGrid<'a> {
             .filter(|(_, p)| p.selectable())
             .map(|(p_id, p)| (*p_id, p))
             .collect();
-        let next_index = panes
+        let next_pane = panes
             .iter()
             .enumerate()
             .filter(|(_, (_, c))| {
@@ -999,9 +1004,15 @@ impl<'a> TiledPaneGrid<'a> {
                     && c.horizontally_overlaps_with(Box::as_ref(current_pane))
             })
             .max_by_key(|(_, (_, c))| c.active_at())
-            .map(|(_, (pid, _))| pid)
+            .map(|(_, (_pid, pane))| pane)
             .copied();
-        next_index
+        let next_pane_is_stacked = next_pane.map(|p| p.position_and_size().is_stacked).unwrap_or(false);
+        if next_pane_is_stacked {
+            if let Some(next_pane_id) = next_pane.map(|p| p.pid()) {
+                return StackedPanes::new(self.panes.clone()).flexible_pane_id_in_stack(&next_pane_id);
+            }
+        }
+        next_pane.map(|p| p.pid())
     }
     fn horizontal_borders(&self, pane_ids: &[PaneId]) -> HashSet<usize> {
         pane_ids.iter().fold(HashSet::new(), |mut borders, p| {
diff --git a/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__move_focus_left_into_stacked_panes.snap b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__move_focus_left_into_stacked_panes.snap
new file mode 100644
index 000000000..03d7bf8f7
--- /dev/null
+++ b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__move_focus_left_into_stacked_panes.snap
@@ -0,0 +1,26 @@
+---
+source: zellij-server/src/tab/./unit/tab_integration_tests.rs
+assertion_line: 3345
+expression: snapshot
+---
+00 (C): ┌ Pane #2 ─────────────────────────────────────────────────┐┌ Pane #14 ─────────────────────────────────────────────────┐
+01 (C): ┌ Pane #3 ─────────────────────────────────────────────────┐│                                                           │
+02 (C): ┌ Pane #4 ─────────────────────────────────────────────────┐│                                                           │
+03 (C): ┌ Pane #5 ─────────────────────────────────────────────────┐│                                                           │
+04 (C): ┌ Pane #6 ─────────────────────────────────────────────────┐│                                                           │
+05 (C): ┌ Pane #7 ─────────────────────────────────────────────────┐│                                                           │
+06 (C): ┌ Pane #8 ─────────────────────────────────────────────────┐│                                                           │
+07 (C): ┌ Pane #9 ─────────────────────────────────────────────────┐│                                                           │
+08 (C): ┌ Pane #10 ────────────────────────────────────────────────┐│                                                           │
+09 (C): ┌ Pane #11 ────────────────────────────────────────────────┐└───────────────────────────────────────────────────────────┘
+10 (C): ┌ Pane #12 ────────────────────────────────────────────────┐┌ Pane #15 ─────────────────────────────────────────────────┐
+11 (C): ┌ Pane #13 ────────────────────────────────────────────────┐│                                                           │
+12 (C): │                                                          ││                                                           │
+13 (C): │                                                          ││                                                           │
+14 (C): │                                                          ││                                                           │
+15 (C): │                                                          ││                                                           │
+16 (C): │                                                          ││                                                           │
+17 (C): │                                                          ││                                                           │
+18 (C): │                                                          ││                                                           │
+19 (C): └──────────────────────────────────────────────────────────┘└───────────────────────────────────────────────────────────┘
+
diff --git a/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__move_focus_right_into_stacked_panes.snap b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__move_focus_right_into_stacked_panes.snap
new file mode 100644
index 000000000..ae4877347
--- /dev/null
+++ b/zellij-server/src/tab/unit/snapshots/zellij_server__tab__tab_integration_tests__move_focus_right_into_stacked_panes.snap
@@ -0,0 +1,26 @@
+---
+source: zellij-server/src/tab/./unit/tab_integration_tests.rs
+assertion_line: 3284
+expression: snapshot
+---
+00 (C): ┌ Pane #14 ────────────────────────────────────────────────┐┌ Pane #2 ──────────────────────────────────────────────────┐
+01 (C): │                                                          │┌ Pane #3 ──────────────────────────────────────────────────┐
+02 (C): │                                                          │┌ Pane #4 ──────────────────────────────────────────────────┐
+03 (C): │                                                          │┌ Pane #5 ──────────────────────────────────────────────────┐
+04 (C): │                                                          │┌ Pane #6 ──────────────────────────────────────────────────┐
+05 (C): │                                                          │┌ Pane #7 ──────────────────────────────────────────────────┐
+06 (C): │                                                          │┌ Pane #8 ──────────────────────────────────────────────────┐
+07 (C): │                                                          │┌ Pane #9 ──────────────────────────────────────────────────┐
+08 (C): │                                                          │┌ Pane #10 ─────────────────────────────────────────────────┐
+09 (C): └──────────────────────────────────────────────────────────┘┌ Pane #11 ─────────────────────────────────────────────────┐
+10 (C): ┌ Pane #15 ────────────────────────────────────────────────┐┌ Pane #12 ─────────────────────────────────────────────────┐
+11 (C): │                                                          │┌ Pane #13 ─────────────────────────────────────────────────┐
+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 4a3ddae72..af15aca55 100644
--- a/zellij-server/src/tab/unit/tab_integration_tests.rs
+++ b/zellij-server/src/tab/unit/tab_integration_tests.rs
@@ -3233,6 +3233,118 @@ fn move_focus_down_with_stacked_panes() {
 }
 
 #[test]
+fn move_focus_right_into_stacked_panes() {
+    // here we make sure that when we focus right into a stack,
+    // we will always focus on the "main" pane of the stack
+    // and not on one of its folds
+    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);
+    for i in 0..13 {
+        let new_pane_id = i + 2;
+        tab.new_pane(PaneId::Terminal(new_pane_id), None, None, Some(client_id))
+            .unwrap();
+    }
+    tab.horizontal_split(PaneId::Terminal(1), None, client_id).unwrap();
+
+    tab.move_focus_up(client_id);
+    tab.move_focus_right(client_id);
+    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, 12)),
+        "cursor coordinates moved to the main pane of the stack",
+    );
+
+    assert_snapshot!(snapshot);
+}
+
+#[test]
+fn move_focus_left_into_stacked_panes() {
+    // here we make sure that when we focus left into a stack,
+    // we will always focus on the "main" pane of the stack
+    // and not on one of its folds
+    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 { children stacked=true; }
+                        pane focus=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);
+    for i in 0..13 {
+        let new_pane_id = i + 2;
+        tab.new_pane(PaneId::Terminal(new_pane_id), None, None, Some(client_id))
+            .unwrap();
+    }
+    tab.horizontal_split(PaneId::Terminal(1), None, client_id).unwrap();
+
+    tab.move_focus_up(client_id);
+    tab.move_focus_left(client_id);
+    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((1, 12)),
+        "cursor coordinates moved to the main pane of the stack",
+    );
+
+    assert_snapshot!(snapshot);
+}
+
+#[test]
 fn close_main_stacked_pane() {
     let size = Size {
         cols: 121,