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;
+}