shell: Try to migrate workspaces back to their original outputs

src/shell/mod.rs

@@ -534,8 +534,38 @@ impl Workspaces {
         workspace_state.add_group_output(&set.group, &output);
 
         self.sets.insert(output.clone(), set);
-        for workspace in &mut self.sets.get_mut(output).unwrap().workspaces {
-            workspace.set_output(output, toplevel_info_state);
+        let mut moved_workspaces = Vec::new();
+        for set in self.sets.values_mut() {
+            let (preferrs, doesnt) = set
+                .workspaces
+                .drain(..)
+                .partition(|w| w.preferrs_output(output));
+            moved_workspaces.extend(preferrs);
+            set.workspaces = doesnt;
+        }
+        {
+            let set = self.sets.get_mut(output).unwrap();
+            for workspace in &mut moved_workspaces {
+                workspace_state.remove_workspace(workspace.handle);
+                let old_workspace_handle = workspace.handle;
+                workspace.handle = workspace_state.create_workspace(&set.group).unwrap();
+                workspace_state.set_workspace_capabilities(
+                    &workspace.handle,
+                    [WorkspaceCapabilities::Activate].into_iter(),
+                );
+                for window in workspace.mapped() {
+                    for (surface, _) in window.windows() {
+                        toplevel_info_state
+                            .toplevel_leave_workspace(&surface, &old_workspace_handle);
+                        toplevel_info_state.toplevel_enter_workspace(&surface, &workspace.handle);
+                    }
+                }
+            }
+            set.workspaces.extend(moved_workspaces);
+            for workspace in &mut set.workspaces {
+                workspace.set_output(output, toplevel_info_state);
+                workspace.refresh(xdg_activation_state);
+            }
         }
     }
 
@@ -951,6 +981,7 @@ impl Shell {
             &mut self.workspace_state.update(),
             &mut self.toplevel_info_state,
         );
+        self.refresh(); // fixes indicies of any moved workspaces
     }
 
     pub fn remove_output(&mut self, output: &Output, seats: impl Iterator<Item = Seat<State>>) {

src/shell/workspace.rs

@@ -217,6 +217,7 @@ impl Workspace {
     ) -> Workspace {
         let tiling_layer = TilingLayout::new(theme, &output);
         let floating_layer = FloatingLayout::new(&output);
+        let output_name = output.name();
 
         Workspace {
             output,
@@ -228,7 +229,11 @@ impl Workspace {
             focus_stack: FocusStacks::default(),
             pending_buffers: Vec::new(),
             screencopy_sessions: Vec::new(),
-            output_stack: VecDeque::new(),
+            output_stack: {
+                let mut queue = VecDeque::new();
+                queue.push_back(output_name);
+                queue
+            },
             pending_tokens: HashSet::new(),
             backdrop_id: Id::new(),
             dirty: AtomicBool::new(false),
@@ -343,9 +348,23 @@ impl Workspace {
                 toplevel_info.toplevel_enter_output(&surface, output);
             }
         }
+        let output_name = output.name();
+        if let Some(pos) = self
+            .output_stack
+            .iter()
+            .position(|name| name == &output_name)
+        {
+            self.output_stack.truncate(pos + 1);
+        } else {
+            self.output_stack.push_back(output.name());
+        }
         self.output = output.clone();
     }
 
+    pub fn preferrs_output(&self, output: &Output) -> bool {
+        self.output_stack.contains(&output.name())
+    }
+
     pub fn unmap(&mut self, mapped: &CosmicMapped) -> Option<ManagedState> {
         let was_fullscreen = self
             .fullscreen