From 3dee249558c7e3f7b4b0fc46e459c996068514f9 Mon Sep 17 00:00:00 2001
From: Victoria Brekenfeld <victoria@system76.com>
Date: Mon, 20 Nov 2023 21:19:47 +0100
Subject: [PATCH] shell: Add keybindings to migrate workspaces

---
 src/config/key_bindings.rs |  4 +++
 src/input/mod.rs           | 50 +++++++++++++++++++++++++++++
 src/shell/mod.rs           | 66 +++++++++++++++++++++++++++++++++++---
 3 files changed, 115 insertions(+), 5 deletions(-)

diff --git a/src/config/key_bindings.rs b/src/config/key_bindings.rs
index 2580771c..f57e137f 100644
--- a/src/config/key_bindings.rs
+++ b/src/config/key_bindings.rs
@@ -167,6 +167,10 @@ pub enum Action {
     MoveToOutput(Direction),
     SendToOutput(Direction),
 
+    MigrateWorkspaceToNextOutput,
+    MigrateWorkspaceToPreviousOutput,
+    MigrateWorkspaceToOutput(Direction),
+
     Focus(FocusDirection),
     Move(Direction),
 
diff --git a/src/input/mod.rs b/src/input/mod.rs
index d4f7d532..e9ccdbe8 100644
--- a/src/input/mod.rs
+++ b/src/input/mod.rs
@@ -1633,6 +1633,56 @@ impl State {
                     }
                 }
             }
+            Action::MigrateWorkspaceToNextOutput => {
+                let current_output = seat.active_output();
+                let active = self.common.shell.active_space(&current_output).handle;
+                let next_output = self
+                    .common
+                    .shell
+                    .outputs()
+                    .skip_while(|o| *o != &current_output)
+                    .skip(1)
+                    .next()
+                    .cloned();
+                if let Some(next_output) = next_output {
+                    self.common
+                        .shell
+                        .migrate_workspace(&current_output, &next_output, &active);
+                }
+            }
+            Action::MigrateWorkspaceToPreviousOutput => {
+                let current_output = seat.active_output();
+                let active = self.common.shell.active_space(&current_output).handle;
+                let prev_output = self
+                    .common
+                    .shell
+                    .outputs()
+                    .rev()
+                    .skip_while(|o| *o != &current_output)
+                    .skip(1)
+                    .next()
+                    .cloned();
+                if let Some(prev_output) = prev_output {
+                    self.common
+                        .shell
+                        .migrate_workspace(&current_output, &prev_output, &active);
+                }
+            }
+            Action::MigrateWorkspaceToOutput(direction) => {
+                let current_output = seat.active_output();
+                let active = self.common.shell.active_space(&current_output).handle;
+                let next_output = self
+                    .common
+                    .shell
+                    .next_output(&current_output, direction)
+                    .cloned();
+
+                if let Some(next_output) = next_output {
+                    self.common
+                        .shell
+                        .migrate_workspace(&current_output, &next_output, &active);
+                }
+            }
             Action::Focus(focus) => {
                 let current_output = seat.active_output();
                 let overview = self.common.shell.overview_mode().0;
diff --git a/src/shell/mod.rs b/src/shell/mod.rs
index 1a629e99..58c6b6f5 100644
--- a/src/shell/mod.rs
+++ b/src/shell/mod.rs
@@ -633,7 +633,7 @@ impl Workspaces {
                             toplevel_info_state,
                         );
                     } else {
-                    // update workspace protocol state
+                        // update workspace protocol state
                         move_workspace_to_group(
                             &mut workspace,
                             &workspace_group,
@@ -641,10 +641,10 @@ impl Workspaces {
                             toplevel_info_state,
                         );
 
-                    // update mapping
-                    workspace.set_output(&new_output, toplevel_info_state);
-                    workspace.refresh(xdg_activation_state);
-                    new_set.workspaces.push(workspace);
+                        // update mapping
+                        workspace.set_output(&new_output, toplevel_info_state);
+                        workspace.refresh(xdg_activation_state);
+                        new_set.workspaces.push(workspace);
                     }
                 }
                 if self.mode == WorkspaceMode::OutputBound {
@@ -665,6 +665,46 @@ impl Workspaces {
         }
     }
 
+    fn migrate_workspace(
+        &mut self,
+        from: &Output,
+        to: &Output,
+        handle: &WorkspaceHandle,
+        workspace_state: &mut WorkspaceUpdateGuard<'_, State>,
+        toplevel_info_state: &mut ToplevelInfoState<State, CosmicSurface>,
+        xdg_activation_state: &XdgActivationState,
+    ) {
+        if !self.sets.contains_key(to) {
+            return;
+        }
+
+        if let Some(mut workspace) = self.sets.get_mut(from).and_then(|set| {
+            let pos = set.workspaces.iter().position(|w| &w.handle == handle)?;
+            Some(set.workspaces.remove(pos))
+        }) {
+            let new_set = self.sets.get_mut(to).unwrap();
+            match self.amount {
+                WorkspaceAmount::Dynamic => {
+                    move_workspace_to_group(
+                        &mut workspace,
+                        &new_set.group,
+                        workspace_state,
+                        toplevel_info_state,
+                    );
+                    workspace.set_output(to, toplevel_info_state);
+                    workspace.refresh(xdg_activation_state);
+                    new_set.workspaces.insert(new_set.active + 1, workspace)
+                }
+                WorkspaceAmount::Static(_) => merge_workspaces(
+                    workspace,
+                    &mut new_set.workspaces[new_set.active],
+                    workspace_state,
+                    toplevel_info_state,
+                ),
+            };
+        }
+    }
+
     pub fn update_config(
         &mut self,
         config: &Config,
@@ -1013,6 +1053,22 @@ impl Shell {
         self.refresh(); // cleans up excess of workspaces and empty workspaces
     }
 
+    pub fn migrate_workspace(&mut self, from: &Output, to: &Output, handle: &WorkspaceHandle) {
+        if from == to {
+            return;
+        }
+
+        self.workspaces.migrate_workspace(
+            from,
+            to,
+            handle,
+            &mut self.workspace_state.update(),
+            &mut self.toplevel_info_state,
+            &self.xdg_activation_state,
+        );
+        self.refresh(); // fixes index of moved workspace
+    }
+
     pub fn update_config(&mut self, config: &Config) {
         let mut workspace_state = self.workspace_state.update();
         let toplevel_info_state = &mut self.toplevel_info_state;