layout: Preserve active workspace for removed outputs

1 files changed, 105 insertions, 3 deletions

diff --git a/src/layout/mod.rs b/src/layout/mod.rs
index fcac0ac8..97a9f622 100644
--- a/src/layout/mod.rs
+++ b/src/layout/mod.rs
@@ -30,6 +30,7 @@
 //! making the primary output their original output.
 
 use std::cmp::min;
+use std::collections::HashMap;
 use std::mem;
 use std::rc::Rc;
 use std::time::Duration;
@@ -205,6 +206,14 @@ pub struct Layout<W: LayoutElement> {
     /// This normally indicates that the layout has keyboard focus, but not always. E.g. when the
     /// screenshot UI is open, it keeps the layout drawing as active.
     is_active: bool,
+    /// Map from monitor name to id of its last active workspace.
+    ///
+    /// This data is stored upon monitor removal and is used to restore the active workspace when
+    /// the monitor is reconnected.
+    ///
+    /// The workspace id does not necessarily point to a valid workspace. If it doesn't, then it is
+    /// simply ignored.
+    last_active_workspace_id: HashMap<String, WorkspaceId>,
     /// Ongoing interactive move.
     interactive_move: Option<InteractiveMoveState<W>>,
     /// Configurable properties of the layout.
@@ -432,6 +441,7 @@ impl<W: LayoutElement> Layout<W> {
         Self {
             monitor_set: MonitorSet::NoOutputs { workspaces: vec![] },
             is_active: true,
+            last_active_workspace_id: HashMap::new(),
             interactive_move: None,
             options: Rc::new(options),
         }
@@ -449,6 +459,7 @@ impl<W: LayoutElement> Layout<W> {
         Self {
             monitor_set: MonitorSet::NoOutputs { workspaces },
             is_active: true,
+            last_active_workspace_id: HashMap::new(),
             interactive_move: None,
             options: opts,
         }
@@ -463,6 +474,9 @@ impl<W: LayoutElement> Layout<W> {
             } => {
                 let primary = &mut monitors[primary_idx];
 
+                let ws_id_to_activate = self.last_active_workspace_id.remove(&output.name());
+                let mut active_workspace_idx = None;
+
                 let mut stopped_primary_ws_switch = false;
 
                 let mut workspaces = vec![];
@@ -482,6 +496,10 @@ impl<W: LayoutElement> Layout<W> {
                         // another monitor. However, we will add an empty workspace in the end
                         // instead.
                         if ws.has_windows() || ws.name.is_some() {
+                            if Some(ws.id()) == ws_id_to_activate {
+                                active_workspace_idx = Some(workspaces.len());
+                            }
+
                             workspaces.push(ws);
                         }
 
@@ -499,6 +517,10 @@ impl<W: LayoutElement> Layout<W> {
 
                 workspaces.reverse();
 
+                if let Some(idx) = &mut active_workspace_idx {
+                    *idx = workspaces.len() - *idx - 1;
+                }
+
                 // Make sure there's always an empty workspace.
                 workspaces.push(Workspace::new(output.clone(), self.options.clone()));
 
@@ -506,7 +528,10 @@ impl<W: LayoutElement> Layout<W> {
                     ws.set_output(Some(output.clone()));
                 }
 
-                monitors.push(Monitor::new(output, workspaces, self.options.clone()));
+                let mut monitor = Monitor::new(output, workspaces, self.options.clone());
+                monitor.active_workspace_idx = active_workspace_idx.unwrap_or(0);
+                monitors.push(monitor);
+
                 MonitorSet::Normal {
                     monitors,
                     primary_idx,
@@ -517,11 +542,19 @@ impl<W: LayoutElement> Layout<W> {
                 // We know there are no empty workspaces there, so add one.
                 workspaces.push(Workspace::new(output.clone(), self.options.clone()));
 
-                for workspace in &mut workspaces {
+                let ws_id_to_activate = self.last_active_workspace_id.remove(&output.name());
+                let mut active_workspace_idx = 0;
+
+                for (i, workspace) in workspaces.iter_mut().enumerate() {
                     workspace.set_output(Some(output.clone()));
+
+                    if Some(workspace.id()) == ws_id_to_activate {
+                        active_workspace_idx = i;
+                    }
                 }
 
-                let monitor = Monitor::new(output, workspaces, self.options.clone());
+                let mut monitor = Monitor::new(output, workspaces, self.options.clone());
+                monitor.active_workspace_idx = active_workspace_idx;
 
                 MonitorSet::Normal {
                     monitors: vec![monitor],
@@ -544,6 +577,12 @@ impl<W: LayoutElement> Layout<W> {
                     .position(|mon| &mon.output == output)
                     .expect("trying to remove non-existing output");
                 let monitor = monitors.remove(idx);
+
+                self.last_active_workspace_id.insert(
+                    monitor.output_name().clone(),
+                    monitor.workspaces[monitor.active_workspace_idx].id(),
+                );
+
                 let mut workspaces = monitor.workspaces;
 
                 for ws in &mut workspaces {
@@ -5684,6 +5723,69 @@ mod tests {
         check_ops(&ops);
     }
 
+    #[test]
+    fn output_active_workspace_is_preserved() {
+        let ops = [
+            Op::AddOutput(1),
+            Op::AddWindow {
+                id: 1,
+                bbox: Rectangle::from_loc_and_size((0, 0), (100, 200)),
+                min_max_size: Default::default(),
+            },
+            Op::FocusWorkspaceDown,
+            Op::AddWindow {
+                id: 2,
+                bbox: Rectangle::from_loc_and_size((0, 0), (100, 200)),
+                min_max_size: Default::default(),
+            },
+            Op::RemoveOutput(1),
+            Op::AddOutput(1),
+        ];
+
+        let mut layout = Layout::default();
+        for op in ops {
+            op.apply(&mut layout);
+        }
+
+        let MonitorSet::Normal { monitors, .. } = layout.monitor_set else {
+            unreachable!()
+        };
+
+        assert_eq!(monitors[0].active_workspace_idx, 1);
+    }
+
+    #[test]
+    fn output_active_workspace_is_preserved_with_other_outputs() {
+        let ops = [
+            Op::AddOutput(1),
+            Op::AddOutput(2),
+            Op::AddWindow {
+                id: 1,
+                bbox: Rectangle::from_loc_and_size((0, 0), (100, 200)),
+                min_max_size: Default::default(),
+            },
+            Op::FocusWorkspaceDown,
+            Op::AddWindow {
+                id: 2,
+                bbox: Rectangle::from_loc_and_size((0, 0), (100, 200)),
+                min_max_size: Default::default(),
+            },
+            Op::RemoveOutput(1),
+            Op::AddOutput(1),
+        ];
+
+        let mut layout = Layout::default();
+        for op in ops {
+            op.apply(&mut layout);
+        }
+
+        let MonitorSet::Normal { monitors, .. } = layout.monitor_set else {
+            unreachable!()
+        };
+
+        assert_eq!(monitors[1].active_workspace_idx, 1);
+    }
+
     fn arbitrary_spacing() -> impl Strategy<Value = f64> {
         // Give equal weight to:
         // - 0: the element is disabled