Added actions to allow focusing up or down as normal but to wrap to the column to the left or right if there is no window above or below

3 files changed, 124 insertions, 0 deletions

diff --git a/src/input/mod.rs b/src/input/mod.rs
index 0d2074a4..8bf94d15 100644
--- a/src/input/mod.rs
+++ b/src/input/mod.rs
@@ -563,6 +563,30 @@ impl State {
                 // FIXME: granular
                 self.niri.queue_redraw_all();
             }
+            Action::FocusWindowDownOrColumnLeft => {
+                self.niri.layout.focus_down_or_left();
+                self.maybe_warp_cursor_to_focus();
+                // FIXME: granular
+                self.niri.queue_redraw_all();
+            }
+            Action::FocusWindowDownOrColumnRight => {
+                self.niri.layout.focus_down_or_right();
+                self.maybe_warp_cursor_to_focus();
+                // FIXME: granular
+                self.niri.queue_redraw_all();
+            }
+            Action::FocusWindowUpOrColumnLeft => {
+                self.niri.layout.focus_up_or_left();
+                self.maybe_warp_cursor_to_focus();
+                // FIXME: granular
+                self.niri.queue_redraw_all();
+            }
+            Action::FocusWindowUpOrColumnRight => {
+                self.niri.layout.focus_up_or_right();
+                self.maybe_warp_cursor_to_focus();
+                // FIXME: granular
+                self.niri.queue_redraw_all();
+            }
             Action::FocusWindowOrWorkspaceDown => {
                 self.niri.layout.focus_window_or_workspace_down();
                 self.maybe_warp_cursor_to_focus();
diff --git a/src/layout/mod.rs b/src/layout/mod.rs
index 02eeda3b..6e8fa292 100644
--- a/src/layout/mod.rs
+++ b/src/layout/mod.rs
@@ -1249,6 +1249,34 @@ impl<W: LayoutElement> Layout<W> {
         monitor.focus_up();
     }
 
+    pub fn focus_down_or_left(&mut self) {
+        let Some(monitor) = self.active_monitor() else {
+            return;
+        };
+        monitor.focus_down_or_left();
+    }
+
+    pub fn focus_down_or_right(&mut self) {
+        let Some(monitor) = self.active_monitor() else {
+            return;
+        };
+        monitor.focus_down_or_right();
+    }
+
+    pub fn focus_up_or_left(&mut self) {
+        let Some(monitor) = self.active_monitor() else {
+            return;
+        };
+        monitor.focus_up_or_left();
+    }
+
+    pub fn focus_up_or_right(&mut self) {
+        let Some(monitor) = self.active_monitor() else {
+            return;
+        };
+        monitor.focus_up_or_right();
+    }
+
     pub fn focus_window_or_workspace_down(&mut self) {
         let Some(monitor) = self.active_monitor() else {
             return;
@@ -2603,6 +2631,10 @@ mod tests {
         FocusColumnLast,
         FocusWindowDown,
         FocusWindowUp,
+        FocusWindowDownOrColumnLeft,
+        FocusWindowDownOrColumnRight,
+        FocusWindowUpOrColumnLeft,
+        FocusWindowUpOrColumnRight,
         FocusWindowOrWorkspaceDown,
         FocusWindowOrWorkspaceUp,
         MoveColumnLeft,
@@ -2887,6 +2919,10 @@ mod tests {
                 Op::FocusColumnLast => layout.focus_column_last(),
                 Op::FocusWindowDown => layout.focus_down(),
                 Op::FocusWindowUp => layout.focus_up(),
+                Op::FocusWindowDownOrColumnLeft => layout.focus_down_or_left(),
+                Op::FocusWindowDownOrColumnRight => layout.focus_down_or_right(),
+                Op::FocusWindowUpOrColumnLeft => layout.focus_up_or_left(),
+                Op::FocusWindowUpOrColumnRight => layout.focus_up_or_right(),
                 Op::FocusWindowOrWorkspaceDown => layout.focus_window_or_workspace_down(),
                 Op::FocusWindowOrWorkspaceUp => layout.focus_window_or_workspace_up(),
                 Op::MoveColumnLeft => layout.move_left(),
@@ -3098,8 +3134,12 @@ mod tests {
             Op::FocusColumnLeft,
             Op::FocusColumnRight,
             Op::FocusWindowUp,
+            Op::FocusWindowUpOrColumnLeft,
+            Op::FocusWindowUpOrColumnRight,
             Op::FocusWindowOrWorkspaceUp,
             Op::FocusWindowDown,
+            Op::FocusWindowDownOrColumnLeft,
+            Op::FocusWindowDownOrColumnRight,
             Op::FocusWindowOrWorkspaceDown,
             Op::MoveColumnLeft,
             Op::MoveColumnRight,
@@ -3251,8 +3291,12 @@ mod tests {
             Op::FocusColumnLeft,
             Op::FocusColumnRight,
             Op::FocusWindowUp,
+            Op::FocusWindowUpOrColumnLeft,
+            Op::FocusWindowUpOrColumnRight,
             Op::FocusWindowOrWorkspaceUp,
             Op::FocusWindowDown,
+            Op::FocusWindowDownOrColumnLeft,
+            Op::FocusWindowDownOrColumnRight,
             Op::FocusWindowOrWorkspaceDown,
             Op::MoveColumnLeft,
             Op::MoveColumnRight,
diff --git a/src/layout/monitor.rs b/src/layout/monitor.rs
index c56146a4..b9cd431a 100644
--- a/src/layout/monitor.rs
+++ b/src/layout/monitor.rs
@@ -328,6 +328,62 @@ impl<W: LayoutElement> Monitor<W> {
         self.active_workspace().focus_up();
     }
 
+    pub fn focus_down_or_left(&mut self) {
+        let workspace = self.active_workspace();
+        if !workspace.columns.is_empty() {
+            let column = &workspace.columns[workspace.active_column_idx];
+            let curr_idx = column.active_tile_idx;
+            let new_idx = min(column.active_tile_idx + 1, column.tiles.len() - 1);
+            if curr_idx == new_idx {
+                self.focus_left();
+            } else {
+                workspace.focus_down();
+            }
+        }
+    }
+
+    pub fn focus_down_or_right(&mut self) {
+        let workspace = self.active_workspace();
+        if !workspace.columns.is_empty() {
+            let column = &workspace.columns[workspace.active_column_idx];
+            let curr_idx = column.active_tile_idx;
+            let new_idx = min(column.active_tile_idx + 1, column.tiles.len() - 1);
+            if curr_idx == new_idx {
+                self.focus_right();
+            } else {
+                workspace.focus_down();
+            }
+        }
+    }
+
+    pub fn focus_up_or_left(&mut self) {
+        let workspace = self.active_workspace();
+        if !workspace.columns.is_empty() {
+            let curr_idx = workspace.columns[workspace.active_column_idx].active_tile_idx;
+            let new_idx = curr_idx.saturating_sub(1);
+            if curr_idx == new_idx {
+                self.focus_left();
+            } else {
+                workspace.focus_up();
+            }
+        }
+    }
+
+    pub fn focus_up_or_right(&mut self) {
+        let workspace = self.active_workspace();
+        if workspace.columns.is_empty() {
+            self.switch_workspace_up();
+        } else {
+            let curr_idx = workspace.columns[workspace.active_column_idx].active_tile_idx;
+            let new_idx = curr_idx.saturating_sub(1);
+            if curr_idx == new_idx {
+                self.focus_left();
+            } else {
+                workspace.focus_up();
+            }
+        }
+    }
+
     pub fn focus_window_or_workspace_down(&mut self) {
         let workspace = self.active_workspace();
         if workspace.columns.is_empty() {