Restore workspaces to outputs based on priority

1 files changed, 22 insertions, 7 deletions

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