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

4 files changed, 133 insertions, 3 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)?;