Restore workspaces to outputs based on priority

5 files changed, 130 insertions, 17 deletions

diff --git a/include/sway/tree/workspace.h b/include/sway/tree/workspace.h
index 81321fc8..c72a4ac0 100644
--- a/include/sway/tree/workspace.h
+++ b/include/sway/tree/workspace.h
@@ -9,6 +9,7 @@ struct sway_workspace {
         struct sway_container *swayc;
         struct sway_view *fullscreen;
         struct sway_container *floating;
+        list_t *output_priority;
 };
 
 extern char *prev_workspace_name;
@@ -33,4 +34,12 @@ bool workspace_is_visible(struct sway_container *ws);
 
 bool workspace_is_empty(struct sway_container *ws);
 
+void workspace_output_raise_priority(struct sway_container *workspace,
+                struct sway_container *old_output, struct sway_container *new_output);
+
+void workspace_output_add_priority(struct sway_container *workspace,
+                struct sway_container *output);
+
+struct sway_container *workspace_output_get_highest_available(
+                struct sway_container *ws, struct sway_container *exclude);
 #endif
diff --git a/sway/tree/container.c b/sway/tree/container.c
index ca993c41..f36820b3 100644
--- a/sway/tree/container.c
+++ b/sway/tree/container.c
@@ -213,6 +213,9 @@ static struct sway_container *container_workspace_destroy(
         sway_workspace->floating->parent = NULL;
         _container_destroy(sway_workspace->floating);
 
+        list_foreach(sway_workspace->output_priority, free);
+        list_free(sway_workspace->output_priority);
+
         free(sway_workspace);
 
         if (output) {
@@ -234,24 +237,33 @@ static struct sway_container *container_output_destroy(
                 // program
                 if (root_container.children->length > 1) {
                         // Move workspace from this output to another output
-                        struct sway_container *other_output =
+                        struct sway_container *fallback_output =
                                 root_container.children->items[0];
-                        if (other_output == output) {
-                                other_output = root_container.children->items[1];
+                        if (fallback_output == output) {
+                                fallback_output = root_container.children->items[1];
                         }
 
                         while (output->children->length) {
                                 struct sway_container *workspace = output->children->items[0];
+
+                                struct sway_container *new_output =
+                                        workspace_output_get_highest_available(workspace, output);
+                                if (!new_output) {
+                                        new_output = fallback_output;
+                                        workspace_output_add_priority(workspace, new_output);
+                                }
+
                                 container_remove_child(workspace);
-                                if (workspace->children->length > 0) {
-                                        container_add_child(other_output, workspace);
+                                if (!workspace_is_empty(workspace)) {
+                                        container_add_child(new_output, workspace);
                                         ipc_event_workspace(workspace, NULL, "move");
                                 } else {
                                         container_workspace_destroy(workspace);
                                 }
+
+                                container_sort_workspaces(new_output);
+                                arrange_output(new_output);
                         }
-                        container_sort_workspaces(other_output);
-                        arrange_output(other_output);
                 }
         }
 
@@ -433,6 +445,9 @@ void container_descendants(struct sway_container *root,
                         func(item, data);
                 }
                 container_descendants(item, type, func, data);
+                if (i < root->children->length && root->children->items[i] != item) {
+                        --i;
+                }
         }
 }
 
diff --git a/sway/tree/layout.c b/sway/tree/layout.c
index b54dc2fe..fc17d8c3 100644
--- a/sway/tree/layout.c
+++ b/sway/tree/layout.c
@@ -183,6 +183,7 @@ void container_move_to(struct sway_container *container,
                 }
                 container_sort_workspaces(new_parent);
                 seat_set_focus(seat, new_parent);
+                workspace_output_raise_priority(container, old_parent, new_parent);
         }
         container_notify_subtree_changed(old_parent);
         container_notify_subtree_changed(new_parent);
diff --git a/sway/tree/output.c b/sway/tree/output.c
index 8823eba0..1ab8bed2 100644
--- a/sway/tree/output.c
+++ b/sway/tree/output.c
@@ -1,11 +1,35 @@
 #define _POSIX_C_SOURCE 200809L
 #include <string.h>
 #include <strings.h>
+#include "sway/ipc-server.h"
 #include "sway/output.h"
+#include "sway/tree/arrange.h"
 #include "sway/tree/output.h"
 #include "sway/tree/workspace.h"
 #include "log.h"
 
+static void restore_workspace(struct sway_container *ws, void *output) {
+        if (ws->parent == output) {
+                return;
+        }
+
+        struct sway_container *highest = workspace_output_get_highest_available(
+                        ws, NULL);
+        if (!highest) {
+                return;
+        }
+
+        if (highest == output) {
+                struct sway_container *other = container_remove_child(ws);
+                container_add_child(output, ws);
+                ipc_event_workspace(ws, NULL, "move");
+
+                container_sort_workspaces(output);
+                arrange_output(output);
+                arrange_output(other);
+        }
+}
+
 struct sway_container *output_create(
                 struct sway_output *sway_output) {
         const char *name = sway_output->wlr_output->name;
@@ -56,19 +80,24 @@ struct sway_container *output_create(
         output->width = size.width;
         output->height = size.height;
 
-        // Create workspace
-        char *ws_name = workspace_next_name(output->name);
-        wlr_log(L_DEBUG, "Creating default workspace %s", ws_name);
-        struct sway_container *ws = workspace_create(output, ws_name);
-        // Set each seat's focus if not already set
-        struct sway_seat *seat = NULL;
-        wl_list_for_each(seat, &input_manager->seats, link) {
-                if (!seat->has_focus) {
-                        seat_set_focus(seat, ws);
+        container_descendants(&root_container, C_WORKSPACE, restore_workspace,
+                        output);
+
+        if (!output->children->length) {
+                // Create workspace
+                char *ws_name = workspace_next_name(output->name);
+                wlr_log(L_DEBUG, "Creating default workspace %s", ws_name);
+                struct sway_container *ws = workspace_create(output, ws_name);
+                // Set each seat's focus if not already set
+                struct sway_seat *seat = NULL;
+                wl_list_for_each(seat, &input_manager->seats, link) {
+                        if (!seat->has_focus) {
+                                seat_set_focus(seat, ws);
+                        }
                 }
+                free(ws_name);
         }
 
-        free(ws_name);
         container_create_notify(output);
         return output;
 }
diff --git a/sway/tree/workspace.c b/sway/tree/workspace.c
index 3c42e259..9ba210fd 100644
--- a/sway/tree/workspace.c
+++ b/sway/tree/workspace.c
@@ -68,7 +68,9 @@ struct sway_container *workspace_create(struct sway_container *output,
         swayws->floating = container_create(C_CONTAINER);
         swayws->floating->parent = swayws->swayc;
         swayws->floating->layout = L_FLOATING;
+        swayws->output_priority = create_list();
         workspace->sway_workspace = swayws;
+        workspace_output_add_priority(workspace, output);
 
         container_add_child(output, workspace);
         container_sort_workspaces(output);
@@ -454,3 +456,60 @@ bool workspace_is_empty(struct sway_container *ws) {
         }
         return true;
 }
+
+static int find_output(const void *id1, const void *id2) {
+        return strcmp(id1, id2) ? 0 : 1;
+}
+
+void workspace_output_raise_priority(struct sway_container *workspace,
+                struct sway_container *old_output, struct sway_container *output) {
+        struct sway_workspace *ws = workspace->sway_workspace;
+
+        int old_index = list_seq_find(ws->output_priority, find_output,
+                        old_output->name);
+        if (old_index < 0) {
+                return;
+        }
+
+        int new_index = list_seq_find(ws->output_priority, find_output,
+                        output->name);
+        if (new_index < 0) {
+                list_insert(ws->output_priority, old_index, strdup(output->name));
+        } else if (new_index > old_index) {
+                char *name = ws->output_priority->items[new_index];
+                list_del(ws->output_priority, new_index);
+                list_insert(ws->output_priority, old_index, name);
+        }
+}
+
+void workspace_output_add_priority(struct sway_container *workspace,
+                struct sway_container *output) {
+        int index = list_seq_find(workspace->sway_workspace->output_priority,
+                        find_output, output->name);
+        if (index < 0) {
+                list_add(workspace->sway_workspace->output_priority,
+                                strdup(output->name));
+        }
+}
+
+static bool _output_by_name(struct sway_container *output, void *data) {
+        return output->type == C_OUTPUT && strcasecmp(output->name, data) == 0;
+}
+
+struct sway_container *workspace_output_get_highest_available(
+                struct sway_container *ws, struct sway_container *exclude) {
+        for (int i = 0; i < ws->sway_workspace->output_priority->length; i++) {
+                char *name = ws->sway_workspace->output_priority->items[i];
+                if (exclude && strcasecmp(name, exclude->name) == 0) {
+                        continue;
+                }
+
+                struct sway_container *output = container_find(&root_container,
+                                _output_by_name, name);
+                if (output) {
+                        return output;
+                }
+        }
+
+        return NULL;
+}