Allow workspace names to be changed dynamically (#904)

* Add un/set workspace name actions

* Add SetWorkspaceName reference to proptests

* Simplify unname_workspace

* Add ewaf version of set first workspace name test

* Simplify more

* Fix comment

* Make workspace in set-workspace-name a positional option

---------

Co-authored-by: Ivan Molodetskikh <yalterz@gmail.com>

3 files changed, 189 insertions, 20 deletions

diff --git a/src/input/mod.rs b/src/input/mod.rs
index 576bc2f5..2b0d1914 100644
--- a/src/input/mod.rs
+++ b/src/input/mod.rs
@@ -1120,6 +1120,18 @@ impl State {
                 // FIXME: granular
                 self.niri.queue_redraw_all();
             }
+            Action::SetWorkspaceName(name) => {
+                self.niri.layout.set_workspace_name(name, None);
+            }
+            Action::SetWorkspaceNameByRef { name, reference } => {
+                self.niri.layout.set_workspace_name(name, Some(reference));
+            }
+            Action::UnsetWorkspaceName => {
+                self.niri.layout.unset_workspace_name(None);
+            }
+            Action::UnsetWorkSpaceNameByRef(reference) => {
+                self.niri.layout.unset_workspace_name(Some(reference));
+            }
             Action::ConsumeWindowIntoColumn => {
                 self.niri.layout.consume_into_column();
                 // This does not cause immediate focus or window size change, so warping mouse to
diff --git a/src/layout/mod.rs b/src/layout/mod.rs
index 9df6d8e6..2c241221 100644
--- a/src/layout/mod.rs
+++ b/src/layout/mod.rs
@@ -38,7 +38,7 @@ use std::time::Duration;
 use monitor::MonitorAddWindowTarget;
 use niri_config::{
     CenterFocusedColumn, Config, CornerRadius, FloatOrInt, PresetSize, Struts,
-    Workspace as WorkspaceConfig,
+    Workspace as WorkspaceConfig, WorkspaceReference,
 };
 use niri_ipc::{PositionChange, SizeChange};
 use scrolling::{Column, ColumnWidth, InsertHint, InsertPosition};
@@ -1215,11 +1215,43 @@ impl<W: LayoutElement> Layout<W> {
         None
     }
 
+    pub fn find_workspace_by_ref(
+        &mut self,
+        reference: WorkspaceReference,
+    ) -> Option<&mut Workspace<W>> {
+        if let WorkspaceReference::Index(index) = reference {
+            self.active_monitor().and_then(|m| {
+                let index = index.saturating_sub(1) as usize;
+                m.workspaces.get_mut(index)
+            })
+        } else {
+            self.workspaces_mut().find(|ws| match &reference {
+                WorkspaceReference::Name(ref_name) => ws
+                    .name
+                    .as_ref()
+                    .map_or(false, |name| name.eq_ignore_ascii_case(ref_name)),
+                WorkspaceReference::Id(id) => ws.id().get() == *id,
+                WorkspaceReference::Index(_) => unreachable!(),
+            })
+        }
+    }
+
     pub fn unname_workspace(&mut self, workspace_name: &str) {
+        self.unname_workspace_by_ref(WorkspaceReference::Name(workspace_name.into()));
+    }
+
+    pub fn unname_workspace_by_ref(&mut self, reference: WorkspaceReference) {
+        let id = self.find_workspace_by_ref(reference).map(|ws| ws.id());
+        if let Some(id) = id {
+            self.unname_workspace_by_id(id);
+        }
+    }
+
+    pub fn unname_workspace_by_id(&mut self, id: WorkspaceId) {
         match &mut self.monitor_set {
             MonitorSet::Normal { monitors, .. } => {
                 for mon in monitors {
-                    if mon.unname_workspace(workspace_name) {
+                    if mon.unname_workspace(id) {
                         if mon.workspace_switch.is_none() {
                             mon.clean_up_workspaces();
                         }
@@ -1229,11 +1261,7 @@ impl<W: LayoutElement> Layout<W> {
             }
             MonitorSet::NoOutputs { workspaces } => {
                 for (idx, ws) in workspaces.iter_mut().enumerate() {
-                    if ws
-                        .name
-                        .as_ref()
-                        .map_or(false, |name| name.eq_ignore_ascii_case(workspace_name))
-                    {
+                    if ws.id() == id {
                         ws.unname();
 
                         // Clean up empty workspaces.
@@ -3759,6 +3787,70 @@ impl<W: LayoutElement> Layout<W> {
         monitor.move_workspace_up();
     }
 
+    pub fn set_workspace_name(&mut self, name: String, reference: Option<WorkspaceReference>) {
+        // ignore the request if the name is already used by another workspace
+        if self.find_workspace_by_name(&name).is_some() {
+            return;
+        }
+
+        let ws = if let Some(reference) = reference {
+            self.find_workspace_by_ref(reference)
+        } else {
+            self.active_workspace_mut()
+        };
+        let Some(ws) = ws else {
+            return;
+        };
+
+        ws.name.replace(name);
+
+        let wsid = ws.id();
+
+        // if `empty_workspace_above_first` is set and `ws` is the first
+        // worskpace on a monitor, another empty workspace needs to
+        // be added before.
+        // Conversely, if `ws` was the last workspace on a monitor, an
+        // empty workspace needs to be added after.
+
+        if let MonitorSet::Normal {
+            monitors,
+            active_monitor_idx,
+            ..
+        } = &mut self.monitor_set
+        {
+            let monitor = &mut monitors[*active_monitor_idx];
+            if self.options.empty_workspace_above_first
+                && monitor
+                    .workspaces
+                    .first()
+                    .is_some_and(|first| first.id() == wsid)
+            {
+                monitor.add_workspace_top();
+            }
+            if monitor
+                .workspaces
+                .last()
+                .is_some_and(|last| last.id() == wsid)
+            {
+                monitor.add_workspace_bottom();
+            }
+        }
+    }
+
+    pub fn unset_workspace_name(&mut self, reference: Option<WorkspaceReference>) {
+        let ws = if let Some(reference) = reference {
+            self.find_workspace_by_ref(reference)
+        } else {
+            self.active_workspace_mut()
+        };
+        let Some(ws) = ws else {
+            return;
+        };
+        let id = ws.id();
+
+        self.unname_workspace_by_id(id);
+    }
+
     pub fn start_open_animation_for_window(&mut self, window: &W::Id) {
         if let Some(InteractiveMoveState::Moving(move_)) = &self.interactive_move {
             if move_.tile.window().id() == window {
@@ -4084,7 +4176,7 @@ impl<W: LayoutElement> Default for MonitorSet<W> {
 mod tests {
     use std::cell::Cell;
 
-    use niri_config::{FloatOrInt, OutputName, WorkspaceName};
+    use niri_config::{FloatOrInt, OutputName, WorkspaceName, WorkspaceReference};
     use proptest::prelude::*;
     use proptest_derive::Arbitrary;
     use smithay::output::{Mode, PhysicalProperties, Subpixel};
@@ -4506,6 +4598,16 @@ mod tests {
         MoveColumnToWorkspace(#[proptest(strategy = "0..=4usize")] usize),
         MoveWorkspaceDown,
         MoveWorkspaceUp,
+        SetWorkspaceName {
+            #[proptest(strategy = "1..=5usize")]
+            new_ws_name: usize,
+            #[proptest(strategy = "proptest::option::of(1..=5usize)")]
+            ws_name: Option<usize>,
+        },
+        UnsetWorkspaceName {
+            #[proptest(strategy = "proptest::option::of(1..=5usize)")]
+            ws_name: Option<usize>,
+        },
         MoveWindowToOutput {
             #[proptest(strategy = "proptest::option::of(1..=5usize)")]
             window_id: Option<usize>,
@@ -4750,6 +4852,19 @@ mod tests {
                 Op::UnnameWorkspace { ws_name } => {
                     layout.unname_workspace(&format!("ws{ws_name}"));
                 }
+                Op::SetWorkspaceName {
+                    new_ws_name,
+                    ws_name,
+                } => {
+                    let ws_ref =
+                        ws_name.map(|ws_name| WorkspaceReference::Name(format!("ws{ws_name}")));
+                    layout.set_workspace_name(format!("ws{new_ws_name}"), ws_ref);
+                }
+                Op::UnsetWorkspaceName { ws_name } => {
+                    let ws_ref =
+                        ws_name.map(|ws_name| WorkspaceReference::Name(format!("ws{ws_name}")));
+                    layout.unset_workspace_name(ws_ref);
+                }
                 Op::AddWindow { mut params } => {
                     if layout.has_window(&params.id) {
                         return;
@@ -6815,6 +6930,53 @@ mod tests {
         check_ops(&ops);
     }
 
+    #[test]
+    fn set_first_workspace_name() {
+        let ops = [
+            Op::AddOutput(0),
+            Op::SetWorkspaceName {
+                new_ws_name: 0,
+                ws_name: None,
+            },
+        ];
+
+        check_ops(&ops);
+    }
+
+    #[test]
+    fn set_first_workspace_name_ewaf() {
+        let ops = [
+            Op::AddOutput(0),
+            Op::SetWorkspaceName {
+                new_ws_name: 0,
+                ws_name: None,
+            },
+        ];
+
+        let options = Options {
+            empty_workspace_above_first: true,
+            ..Default::default()
+        };
+        check_ops_with_options(options, &ops);
+    }
+
+    #[test]
+    fn set_last_workspace_name() {
+        let ops = [
+            Op::AddOutput(0),
+            Op::AddWindow {
+                params: TestWindowParams::new(0),
+            },
+            Op::FocusWorkspaceDown,
+            Op::SetWorkspaceName {
+                new_ws_name: 0,
+                ws_name: None,
+            },
+        ];
+
+        check_ops(&ops);
+    }
+
     fn parent_id_causes_loop(layout: &Layout<TestWindow>, id: usize, mut parent_id: usize) -> bool {
         if parent_id == id {
             return true;
diff --git a/src/layout/monitor.rs b/src/layout/monitor.rs
index 57ca7f53..70009d67 100644
--- a/src/layout/monitor.rs
+++ b/src/layout/monitor.rs
@@ -382,18 +382,13 @@ impl<W: LayoutElement> Monitor<W> {
         }
     }
 
-    pub fn unname_workspace(&mut self, workspace_name: &str) -> bool {
-        for ws in &mut self.workspaces {
-            if ws
-                .name
-                .as_ref()
-                .map_or(false, |name| name.eq_ignore_ascii_case(workspace_name))
-            {
-                ws.unname();
-                return true;
-            }
-        }
-        false
+    pub fn unname_workspace(&mut self, id: WorkspaceId) -> bool {
+        let Some(ws) = self.workspaces.iter_mut().find(|ws| ws.id() == id) else {
+            return false;
+        };
+
+        ws.unname();
+        true
     }
 
     pub fn move_left(&mut self) -> bool {