fix(screen): focus pane on screen edge when moving pane focus offtab (#2293)

* fix(screen): focus pane on proper edge when switching tabs through pane switch

* style(fmt): rustfmt

7 files changed, 201 insertions, 23 deletions

diff --git a/zellij-server/src/panes/floating_panes/floating_pane_grid.rs b/zellij-server/src/panes/floating_panes/floating_pane_grid.rs
index afeecc8d2..b8472dfc5 100644
--- a/zellij-server/src/panes/floating_panes/floating_pane_grid.rs
+++ b/zellij-server/src/panes/floating_panes/floating_pane_grid.rs
@@ -2,7 +2,7 @@ use crate::tab::{MIN_TERMINAL_HEIGHT, MIN_TERMINAL_WIDTH};
 use crate::{panes::PaneId, tab::Pane};
 use std::cmp::Ordering;
 use std::collections::HashMap;
-use zellij_utils::data::ResizeStrategy;
+use zellij_utils::data::{Direction, ResizeStrategy};
 use zellij_utils::errors::prelude::*;
 use zellij_utils::pane_size::{Dimension, PaneGeom, Size, Viewport};
 
@@ -625,6 +625,50 @@ impl<'a> FloatingPaneGrid<'a> {
             .copied();
         next_index
     }
+    pub fn pane_id_on_edge(&self, direction: Direction) -> Option<PaneId> {
+        let panes = self.panes.borrow();
+        let panes: Vec<(PaneId, &&mut Box<dyn Pane>)> = panes
+            .iter()
+            .filter(|(_, p)| p.selectable())
+            .map(|(p_id, p)| (*p_id, p))
+            .collect();
+        let next_index = panes
+            .iter()
+            .enumerate()
+            .max_by(|(_, (_, a)), (_, (_, b))| match direction {
+                Direction::Left => {
+                    let x_comparison = a.x().cmp(&b.x());
+                    match x_comparison {
+                        Ordering::Equal => a.y().cmp(&b.y()),
+                        _ => x_comparison,
+                    }
+                },
+                Direction::Right => {
+                    let x_comparison = b.x().cmp(&a.x());
+                    match x_comparison {
+                        Ordering::Equal => a.y().cmp(&b.y()),
+                        _ => x_comparison,
+                    }
+                },
+                Direction::Up => {
+                    let y_comparison = a.y().cmp(&b.y());
+                    match y_comparison {
+                        Ordering::Equal => a.x().cmp(&b.x()),
+                        _ => y_comparison,
+                    }
+                },
+                Direction::Down => {
+                    let y_comparison = b.y().cmp(&a.y());
+                    match y_comparison {
+                        Ordering::Equal => b.x().cmp(&a.x()),
+                        _ => y_comparison,
+                    }
+                },
+            })
+            .map(|(_, (pid, _))| pid)
+            .copied();
+        next_index
+    }
     pub fn next_selectable_pane_id_below(&self, current_pane_id: &PaneId) -> Option<PaneId> {
         let panes = self.panes.borrow();
         let current_pane = panes.get(current_pane_id)?;
diff --git a/zellij-server/src/panes/floating_panes/mod.rs b/zellij-server/src/panes/floating_panes/mod.rs
index 765f0196d..84e85431f 100644
--- a/zellij-server/src/panes/floating_panes/mod.rs
+++ b/zellij-server/src/panes/floating_panes/mod.rs
@@ -529,6 +529,19 @@ impl FloatingPanes {
         }
         Ok(false)
     }
+    pub fn focus_pane_on_edge(&mut self, direction: Direction, client_id: ClientId) {
+        let display_area = *self.display_area.borrow();
+        let viewport = *self.viewport.borrow();
+        let mut floating_pane_grid = FloatingPaneGrid::new(
+            &mut self.panes,
+            &mut self.desired_pane_positions,
+            display_area,
+            viewport,
+        );
+        let pane_id = floating_pane_grid.pane_id_on_edge(direction).unwrap();
+        self.focus_pane(pane_id, client_id);
+        self.set_force_render();
+    }
 
     pub fn move_active_pane_down(&mut self, client_id: ClientId) {
         let display_area = *self.display_area.borrow();
diff --git a/zellij-server/src/panes/tiled_panes/mod.rs b/zellij-server/src/panes/tiled_panes/mod.rs
index ac0f8c635..d4eb66b22 100644
--- a/zellij-server/src/panes/tiled_panes/mod.rs
+++ b/zellij-server/src/panes/tiled_panes/mod.rs
@@ -18,7 +18,7 @@ use crate::{
 };
 use stacked_panes::StackedPanes;
 use zellij_utils::{
-    data::{ModeInfo, ResizeStrategy, Style},
+    data::{Direction, ModeInfo, ResizeStrategy, Style},
     errors::prelude::*,
     input::{command::RunCommand, layout::SplitDirection},
     pane_size::{Offset, PaneGeom, Size, SizeInPixels, Viewport},
@@ -862,6 +862,35 @@ impl TiledPanes {
             (size_in_pixels.height as f64 / size_in_pixels.width as f64).round() as usize
         })
     }
+    pub fn focus_pane_on_edge(&mut self, direction: Direction, client_id: ClientId) {
+        let pane_grid = TiledPaneGrid::new(
+            &mut self.panes,
+            &self.panes_to_hide,
+            *self.display_area.borrow(),
+            *self.viewport.borrow(),
+        );
+        let next_index = pane_grid.pane_id_on_edge(direction).unwrap();
+        // render previously active pane so that its frame does not remain actively
+        // colored
+        let previously_active_pane = self
+            .panes
+            .get_mut(self.active_panes.get(&client_id).unwrap())
+            .unwrap();
+
+        previously_active_pane.set_should_render(true);
+        // we render the full viewport to remove any ui elements that might have been
+        // there before (eg. another user's cursor)
+        previously_active_pane.render_full_viewport();
+
+        let next_active_pane = self.panes.get_mut(&next_index).unwrap();
+        next_active_pane.set_should_render(true);
+        // we render the full viewport to remove any ui elements that might have been
+        // there before (eg. another user's cursor)
+        next_active_pane.render_full_viewport();
+
+        self.focus_pane(next_index, client_id);
+        self.set_pane_active_at(next_index);
+    }
     pub fn move_focus_left(&mut self, client_id: ClientId) -> bool {
         match self.get_active_pane_id(client_id) {
             Some(active_pane_id) => {
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 8fc40d2b4..5ee5cf3ca 100644
--- a/zellij-server/src/panes/tiled_panes/tiled_pane_grid.rs
+++ b/zellij-server/src/panes/tiled_panes/tiled_pane_grid.rs
@@ -3,7 +3,7 @@ use super::pane_resizer::PaneResizer;
 use super::stacked_panes::StackedPanes;
 use crate::tab::{MIN_TERMINAL_HEIGHT, MIN_TERMINAL_WIDTH};
 use crate::{panes::PaneId, tab::Pane};
-use std::cmp::Reverse;
+use std::cmp::{Ordering, Reverse};
 use std::collections::{HashMap, HashSet};
 use zellij_utils::data::{Direction, Resize, ResizeStrategy};
 use zellij_utils::{
@@ -1025,6 +1025,50 @@ impl<'a> TiledPaneGrid<'a> {
             .copied();
         next_index
     }
+    pub fn pane_id_on_edge(&self, direction: Direction) -> Option<PaneId> {
+        let panes = self.panes.borrow();
+        let panes: Vec<(PaneId, &&mut Box<dyn Pane>)> = panes
+            .iter()
+            .filter(|(_, p)| p.selectable())
+            .map(|(p_id, p)| (*p_id, p))
+            .collect();
+        let next_index = panes
+            .iter()
+            .enumerate()
+            .max_by(|(_, (_, a)), (_, (_, b))| match direction {
+                Direction::Left => {
+                    let x_comparison = a.x().cmp(&b.x());
+                    match x_comparison {
+                        Ordering::Equal => b.y().cmp(&a.y()),
+                        _ => x_comparison,
+                    }
+                },
+                Direction::Right => {
+                    let x_comparison = b.x().cmp(&a.x());
+                    match x_comparison {
+                        Ordering::Equal => b.y().cmp(&a.y()),
+                        _ => x_comparison,
+                    }
+                },
+                Direction::Up => {
+                    let y_comparison = a.y().cmp(&b.y());
+                    match y_comparison {
+                        Ordering::Equal => a.x().cmp(&b.x()),
+                        _ => y_comparison,
+                    }
+                },
+                Direction::Down => {
+                    let y_comparison = b.y().cmp(&a.y());
+                    match y_comparison {
+                        Ordering::Equal => b.x().cmp(&a.x()),
+                        _ => y_comparison,
+                    }
+                },
+            })
+            .map(|(_, (pid, _))| pid)
+            .copied();
+        next_index
+    }
     pub fn next_selectable_pane_id_above(&self, current_pane_id: &PaneId) -> Option<PaneId> {
         let panes = self.panes.borrow();
         let current_pane = panes.get(current_pane_id)?;
diff --git a/zellij-server/src/screen.rs b/zellij-server/src/screen.rs
index a2fa3e506..5a4223142 100644
--- a/zellij-server/src/screen.rs
+++ b/zellij-server/src/screen.rs
@@ -597,7 +597,12 @@ impl Screen {
     }
 
     /// A helper function to switch to a new tab at specified position.
-    fn switch_active_tab(&mut self, new_tab_pos: usize, client_id: ClientId) -> Result<()> {
+    fn switch_active_tab(
+        &mut self,
+        new_tab_pos: usize,
+        should_change_pane_focus: Option<Direction>,
+        client_id: ClientId,
+    ) -> Result<()> {
         let err_context = || {
             format!(
             "Failed to switch to active tab at position {new_tab_pos} for client id: {client_id:?}"
@@ -621,6 +626,15 @@ impl Screen {
                             self.connected_clients.borrow().iter().copied().collect();
                         for client_id in all_connected_clients {
                             self.update_client_tab_focus(client_id, new_tab_index);
+                            match (
+                                should_change_pane_focus,
+                                self.get_indexed_tab_mut(new_tab_index),
+                            ) {
+                                (Some(direction), Some(new_tab)) => {
+                                    new_tab.focus_pane_on_edge(direction, client_id);
+                                },
+                                _ => {},
+                            }
                         }
                     } else {
                         self.move_clients_between_tabs(
@@ -629,6 +643,15 @@ impl Screen {
                             Some(vec![client_id]),
                         )
                         .with_context(err_context)?;
+                        match (
+                            should_change_pane_focus,
+                            self.get_indexed_tab_mut(new_tab_index),
+                        ) {
+                            (Some(direction), Some(new_tab)) => {
+                                new_tab.focus_pane_on_edge(direction, client_id);
+                            },
+                            _ => {},
+                        }
                         self.update_client_tab_focus(client_id, new_tab_index);
                     }
 
@@ -656,7 +679,7 @@ impl Screen {
     fn switch_active_tab_name(&mut self, name: String, client_id: ClientId) -> Result<bool> {
         match self.tabs.values().find(|t| t.name == name) {
             Some(new_tab) => {
-                self.switch_active_tab(new_tab.position, client_id)?;
+                self.switch_active_tab(new_tab.position, None, client_id)?;
                 Ok(true)
             },
             None => Ok(false),
@@ -664,7 +687,11 @@ impl Screen {
     }
 
     /// Sets this [`Screen`]'s active [`Tab`] to the next tab.
-    pub fn switch_tab_next(&mut self, client_id: ClientId) -> Result<()> {
+    pub fn switch_tab_next(
+        &mut self,
+        should_change_pane_focus: Option<Direction>,
+        client_id: ClientId,
+    ) -> Result<()> {
         let err_context = || format!("failed to switch to next tab for client {client_id}");
 
         let client_id = if self.get_active_tab(client_id).is_ok() {
@@ -678,7 +705,11 @@ impl Screen {
                 Ok(active_tab) => {
                     let active_tab_pos = active_tab.position;
                     let new_tab_pos = (active_tab_pos + 1) % self.tabs.len();
-                    return self.switch_active_tab(new_tab_pos, client_id);
+                    return self.switch_active_tab(
+                        new_tab_pos,
+                        should_change_pane_focus,
+                        client_id,
+                    );
                 },
                 Err(err) => Err::<(), _>(err).with_context(err_context).non_fatal(),
             }
@@ -687,7 +718,11 @@ impl Screen {
     }
 
     /// Sets this [`Screen`]'s active [`Tab`] to the previous tab.
-    pub fn switch_tab_prev(&mut self, client_id: ClientId) -> Result<()> {
+    pub fn switch_tab_prev(
+        &mut self,
+        should_change_pane_focus: Option<Direction>,
+        client_id: ClientId,
+    ) -> Result<()> {
         let err_context = || format!("failed to switch to previous tab for client {client_id}");
 
         let client_id = if self.get_active_tab(client_id).is_ok() {
@@ -706,7 +741,11 @@ impl Screen {
                         active_tab_pos - 1
                     };
 
-                    return self.switch_active_tab(new_tab_pos, client_id);
+                    return self.switch_active_tab(
+                        new_tab_pos,
+                        should_change_pane_focus,
+                        client_id,
+                    );
                 },
                 Err(err) => Err::<(), _>(err).with_context(err_context).non_fatal(),
             }
@@ -715,7 +754,7 @@ impl Screen {
     }
 
     pub fn go_to_tab(&mut self, tab_index: usize, client_id: ClientId) -> Result<()> {
-        self.switch_active_tab(tab_index.saturating_sub(1), client_id)
+        self.switch_active_tab(tab_index.saturating_sub(1), None, client_id)
     }
 
     pub fn go_to_tab_name(&mut self, name: String, client_id: ClientId) -> Result<bool> {
@@ -1320,7 +1359,7 @@ impl Screen {
                         .move_focus_left(client_id)
                         .and_then(|success| {
                             if !success {
-                                self.switch_tab_prev(client_id)
+                                self.switch_tab_prev(Some(Direction::Left), client_id)
                                     .context("failed to move focus to previous tab")
                             } else {
                                 Ok(())
@@ -1354,7 +1393,7 @@ impl Screen {
                         .move_focus_right(client_id)
                         .and_then(|success| {
                             if !success {
-                                self.switch_tab_next(client_id)
+                                self.switch_tab_next(Some(Direction::Right), client_id)
                                     .context("failed to move focus to next tab")
                             } else {
                                 Ok(())
@@ -2004,12 +2043,12 @@ pub(crate) fn screen_thread_main(
                 screen.unblock_input()?;
             },
             ScreenInstruction::SwitchTabNext(client_id) => {
-                screen.switch_tab_next(client_id)?;
+                screen.switch_tab_next(None, client_id)?;
                 screen.unblock_input()?;
                 screen.render()?;
             },
             ScreenInstruction::SwitchTabPrev(client_id) => {
-                screen.switch_tab_prev(client_id)?;
+                screen.switch_tab_prev(None, client_id)?;
                 screen.unblock_input()?;
                 screen.render()?;
             },
diff --git a/zellij-server/src/tab/mod.rs b/zellij-server/src/tab/mod.rs
index 4134594b6..521e1dbbd 100644
--- a/zellij-server/src/tab/mod.rs
+++ b/zellij-server/src/tab/mod.rs
@@ -1915,6 +1915,15 @@ impl Tab {
         }
         self.tiled_panes.focus_previous_pane(client_id);
     }
+    pub fn focus_pane_on_edge(&mut self, direction: Direction, client_id: ClientId) {
+        let err_context = || format!("failed to move focus left for client {}", client_id);
+
+        if self.floating_panes.panes_are_visible() {
+            self.floating_panes.focus_pane_on_edge(direction, client_id);
+        } else if self.has_selectable_panes() && !self.tiled_panes.fullscreen_is_active() {
+            self.tiled_panes.focus_pane_on_edge(direction, client_id);
+        }
+    }
     // returns a boolean that indicates whether the focus moved
     pub fn move_focus_left(&mut self, client_id: ClientId) -> Result<bool> {
         let err_context = || format!("failed to move focus left for client {}", client_id);
diff --git a/zellij-server/src/unit/screen_tests.rs b/zellij-server/src/unit/screen_tests.rs
index 7b97ca504..516c86309 100644
--- a/zellij-server/src/unit/screen_tests.rs
+++ b/zellij-server/src/unit/screen_tests.rs
@@ -529,7 +529,7 @@ pub fn switch_to_prev_tab() {
 
     new_tab(&mut screen, 1, 1);
     new_tab(&mut screen, 2, 2);
-    screen.switch_tab_prev(1).expect("TEST");
+    screen.switch_tab_prev(None, 1).expect("TEST");
 
     assert_eq!(
         screen.get_active_tab(1).unwrap().position,
@@ -548,8 +548,8 @@ pub fn switch_to_next_tab() {
 
     new_tab(&mut screen, 1, 1);
     new_tab(&mut screen, 2, 2);
-    screen.switch_tab_prev(1).expect("TEST");
-    screen.switch_tab_next(1).expect("TEST");
+    screen.switch_tab_prev(None, 1).expect("TEST");
+    screen.switch_tab_next(None, 1).expect("TEST");
 
     assert_eq!(
         screen.get_active_tab(1).unwrap().position,
@@ -623,7 +623,7 @@ pub fn close_the_middle_tab() {
     new_tab(&mut screen, 1, 1);
     new_tab(&mut screen, 2, 2);
     new_tab(&mut screen, 3, 3);
-    screen.switch_tab_prev(1).expect("TEST");
+    screen.switch_tab_prev(None, 1).expect("TEST");
     screen.close_tab(1).expect("TEST");
 
     assert_eq!(screen.tabs.len(), 2, "Two tabs left");
@@ -645,7 +645,7 @@ fn move_focus_left_at_left_screen_edge_changes_tab() {
     new_tab(&mut screen, 1, 1);
     new_tab(&mut screen, 2, 2);
     new_tab(&mut screen, 3, 3);
-    screen.switch_tab_prev(1).expect("TEST");
+    screen.switch_tab_prev(None, 1).expect("TEST");
     screen.move_focus_left_or_previous_tab(1).expect("TEST");
 
     assert_eq!(
@@ -666,7 +666,7 @@ fn move_focus_right_at_right_screen_edge_changes_tab() {
     new_tab(&mut screen, 1, 1);
     new_tab(&mut screen, 2, 2);
     new_tab(&mut screen, 3, 3);
-    screen.switch_tab_prev(1).expect("TEST");
+    screen.switch_tab_prev(None, 1).expect("TEST");
     screen.move_focus_right_or_next_tab(1).expect("TEST");
 
     assert_eq!(
@@ -802,7 +802,7 @@ pub fn toggle_to_previous_tab_delete() {
         "Active tab toggler to previous tab"
     );
 
-    screen.switch_tab_prev(1).expect("TEST");
+    screen.switch_tab_prev(None, 1).expect("TEST");
     assert_eq!(
         screen.tab_history.get(&1).unwrap(),
         &[0, 1, 3],
@@ -813,7 +813,7 @@ pub fn toggle_to_previous_tab_delete() {
         2,
         "Active tab toggler to previous tab"
     );
-    screen.switch_tab_prev(1).expect("TEST");
+    screen.switch_tab_prev(None, 1).expect("TEST");
     assert_eq!(
         screen.tab_history.get(&1).unwrap(),
         &[0, 3, 2],
@@ -868,7 +868,7 @@ fn switch_to_tab_with_fullscreen() {
     }
     new_tab(&mut screen, 2, 2);
 
-    screen.switch_tab_prev(1).expect("TEST");
+    screen.switch_tab_prev(None, 1).expect("TEST");
 
     assert_eq!(
         screen.get_active_tab(1).unwrap().position,