Move workspace gesture into monitor & fix missing workspace cleanup

2 files changed, 97 insertions, 48 deletions

diff --git a/src/layout/mod.rs b/src/layout/mod.rs
index 6b32aafa..490f5327 100644
--- a/src/layout/mod.rs
+++ b/src/layout/mod.rs
@@ -49,10 +49,9 @@ use smithay::utils::{Logical, Point, Rectangle, Scale, Size, Transform};
 use smithay::wayland::compositor::{send_surface_state, with_states};
 use smithay::wayland::shell::xdg::SurfaceCachedState;
 
+use self::monitor::Monitor;
 pub use self::monitor::MonitorRenderElement;
-use self::monitor::{Monitor, WorkspaceSwitch, WorkspaceSwitchGesture};
 use self::workspace::{compute_working_area, Column, ColumnWidth, OutputId, Workspace};
-use crate::animation::Animation;
 use crate::niri::WindowOffscreenId;
 use crate::niri_render_elements;
 use crate::render_helpers::renderer::NiriRenderer;
@@ -1222,6 +1221,8 @@ impl<W: LayoutElement> Layout<W> {
 
     #[cfg(test)]
     fn verify_invariants(&self) {
+        use crate::layout::monitor::WorkspaceSwitch;
+
         let (monitors, &primary_idx, &active_monitor_idx) = match &self.monitor_set {
             MonitorSet::Normal {
                 monitors,
@@ -1594,24 +1595,11 @@ impl<W: LayoutElement> Layout<W> {
         for monitor in monitors {
             // Cancel the gesture on other outputs.
             if &monitor.output != output {
-                if let Some(WorkspaceSwitch::Gesture(_)) = monitor.workspace_switch {
-                    monitor.workspace_switch = None;
-                }
+                monitor.workspace_switch_gesture_end(true);
                 continue;
             }
 
-            let center_idx = monitor.active_workspace_idx;
-            let current_idx = monitor
-                .workspace_switch
-                .as_ref()
-                .map(|s| s.current_idx())
-                .unwrap_or(center_idx as f64);
-
-            let gesture = WorkspaceSwitchGesture {
-                center_idx,
-                current_idx,
-            };
-            monitor.workspace_switch = Some(WorkspaceSwitch::Gesture(gesture));
+            monitor.workspace_switch_gesture_begin();
         }
     }
 
@@ -1622,20 +1610,12 @@ impl<W: LayoutElement> Layout<W> {
         };
 
         for monitor in monitors {
-            if let Some(WorkspaceSwitch::Gesture(gesture)) = &mut monitor.workspace_switch {
-                // Normalize like GNOME Shell's workspace switching.
-                let delta_y = delta_y / 400.;
-
-                let min = gesture.center_idx.saturating_sub(1) as f64;
-                let max = (gesture.center_idx + 1).min(monitor.workspaces.len() - 1) as f64;
-                let new_idx = (gesture.current_idx + delta_y).clamp(min, max);
-
-                if gesture.current_idx == new_idx {
+            if let Some(refresh) = monitor.workspace_switch_gesture_update(delta_y) {
+                if refresh {
+                    return Some(Some(monitor.output.clone()));
+                } else {
                     return Some(None);
                 }
-
-                gesture.current_idx = new_idx;
-                return Some(Some(monitor.output.clone()));
             }
         }
 
@@ -1649,25 +1629,7 @@ impl<W: LayoutElement> Layout<W> {
         };
 
         for monitor in monitors {
-            if let Some(WorkspaceSwitch::Gesture(gesture)) = &mut monitor.workspace_switch {
-                if cancelled {
-                    monitor.workspace_switch = None;
-                    return Some(monitor.output.clone());
-                }
-
-                // FIXME: keep track of gesture velocity and use it to compute the final point and
-                // to animate to it.
-                let current_idx = gesture.current_idx;
-                let idx = current_idx.round() as usize;
-
-                monitor.active_workspace_idx = idx;
-                monitor.workspace_switch = Some(WorkspaceSwitch::Animation(Animation::new(
-                    current_idx,
-                    idx as f64,
-                    self.options.animations.workspace_switch,
-                    niri_config::Animation::default_workspace_switch(),
-                )));
-
+            if monitor.workspace_switch_gesture_end(cancelled) {
                 return Some(monitor.output.clone());
             }
         }
@@ -2072,6 +2034,17 @@ mod tests {
             delta: f64,
         },
         ViewOffsetGestureEnd,
+        WorkspaceSwitchGestureBegin {
+            #[proptest(strategy = "1..=5usize")]
+            output_idx: usize,
+        },
+        WorkspaceSwitchGestureUpdate {
+            #[proptest(strategy = "-400f64..400f64")]
+            delta: f64,
+        },
+        WorkspaceSwitchGestureEnd {
+            cancelled: bool,
+        },
     }
 
     impl Op {
@@ -2322,6 +2295,20 @@ mod tests {
                     // We don't handle cancels in this gesture.
                     layout.view_offset_gesture_end(false);
                 }
+                Op::WorkspaceSwitchGestureBegin { output_idx: id } => {
+                    let name = format!("output{id}");
+                    let Some(output) = layout.outputs().find(|o| o.name() == name).cloned() else {
+                        return;
+                    };
+
+                    layout.workspace_switch_gesture_begin(&output);
+                }
+                Op::WorkspaceSwitchGestureUpdate { delta } => {
+                    layout.workspace_switch_gesture_update(delta);
+                }
+                Op::WorkspaceSwitchGestureEnd { cancelled } => {
+                    layout.workspace_switch_gesture_end(cancelled);
+                }
             }
         }
     }
diff --git a/src/layout/monitor.rs b/src/layout/monitor.rs
index df167b89..70a6981d 100644
--- a/src/layout/monitor.rs
+++ b/src/layout/monitor.rs
@@ -691,4 +691,66 @@ impl<W: LayoutElement> Monitor<W> {
             }
         }
     }
+
+    pub fn workspace_switch_gesture_begin(&mut self) {
+        let center_idx = self.active_workspace_idx;
+        let current_idx = self
+            .workspace_switch
+            .as_ref()
+            .map(|s| s.current_idx())
+            .unwrap_or(center_idx as f64);
+
+        let gesture = WorkspaceSwitchGesture {
+            center_idx,
+            current_idx,
+        };
+        self.workspace_switch = Some(WorkspaceSwitch::Gesture(gesture));
+    }
+
+    pub fn workspace_switch_gesture_update(&mut self, delta_y: f64) -> Option<bool> {
+        let Some(WorkspaceSwitch::Gesture(gesture)) = &mut self.workspace_switch else {
+            return None;
+        };
+
+        // Normalize like GNOME Shell's workspace switching.
+        let delta_y = delta_y / 400.;
+
+        let min = gesture.center_idx.saturating_sub(1) as f64;
+        let max = (gesture.center_idx + 1).min(self.workspaces.len() - 1) as f64;
+        let new_idx = (gesture.current_idx + delta_y).clamp(min, max);
+
+        if gesture.current_idx == new_idx {
+            return Some(false);
+        }
+
+        gesture.current_idx = new_idx;
+        Some(true)
+    }
+
+    pub fn workspace_switch_gesture_end(&mut self, cancelled: bool) -> bool {
+        let Some(WorkspaceSwitch::Gesture(gesture)) = &mut self.workspace_switch else {
+            return false;
+        };
+
+        if cancelled {
+            self.workspace_switch = None;
+            self.clean_up_workspaces();
+            return true;
+        }
+
+        // FIXME: keep track of gesture velocity and use it to compute the final point and to
+        // animate to it.
+        let current_idx = gesture.current_idx;
+        let idx = current_idx.round() as usize;
+
+        self.active_workspace_idx = idx;
+        self.workspace_switch = Some(WorkspaceSwitch::Animation(Animation::new(
+            current_idx,
+            idx as f64,
+            self.options.animations.workspace_switch,
+            niri_config::Animation::default_workspace_switch(),
+        )));
+
+        true
+    }
 }