1 files changed, 38 insertions, 11 deletions

diff --git a/sway/commands/unmark.c b/sway/commands/unmark.c
index a7d39432..ea2a5709 100644
--- a/sway/commands/unmark.c
+++ b/sway/commands/unmark.c
@@ -7,26 +7,53 @@
 #include "log.h"
 #include "stringop.h"
 
+static void remove_all_marks_iterator(struct sway_container *con, void *data) {
+	if (con->type == C_VIEW) {
+		view_clear_marks(con->sway_view);
+	}
+}
+
+// unmark                  Remove all marks from all views
+// unmark foo              Remove single mark from whichever view has it
+// [criteria] unmark       Remove all marks from matched view
+// [criteria] unmark foo   Remove single mark from matched view
+
 struct cmd_results *cmd_unmark(int argc, char **argv) {
-	if (argc == 0) {
-		// Remove all marks from the current container
+	// Determine the view
+	struct sway_view *view = NULL;
+	if (config->handler_context.using_criteria) {
 		struct sway_container *container =
 			config->handler_context.current_container;
 		if (container->type != C_VIEW) {
 			return cmd_results_new(CMD_INVALID, "unmark",
 					"Only views can have marks");
 		}
-		view_clear_marks(container->sway_view);
-	} else {
-		// Remove a single mark from whichever container has it
-		char *mark = join_args(argv, argc);
-		if (!view_find_and_unmark(mark)) {
-			free(mark);
-			return cmd_results_new(CMD_INVALID, "unmark",
-					"No view exists with that mark");
+		view = container->sway_view;
+	}
+
+	// Determine the mark
+	char *mark = NULL;
+	if (argc > 0) {
+		mark = join_args(argv, argc);
+	}
+
+	if (view && mark) {
+		// Remove the mark from the given view
+		if (view_has_mark(view, mark)) {
+			view_find_and_unmark(mark);
 		}
-		free(mark);
+	} else if (view && !mark) {
+		// Clear all marks from the given view
+		view_clear_marks(view);
+	} else if (!view && mark) {
+		// Remove mark from whichever view has it
+		view_find_and_unmark(mark);
+	} else {
+		// Remove all marks from all views
+		container_for_each_descendant_dfs(&root_container,
+				remove_all_marks_iterator, NULL);
 	}
+	free(mark);
 
 	return cmd_results_new(CMD_SUCCESS, NULL, NULL);
 }