feat(wm): add workspace rules

This feature allows users to specify which monitor/workspace an
application's window, identified either by executable name or window
class name, should be assigned to.

A new fn, WindowManager.enforce_workspace_rules, is called whenever a
new rule is added, and periodically whenever an event is processed by
komorebi (just after orphan windows are repead, before the matching and
processing of the specific event).

Both class and exe identifiers are stored in the same HashMap for the
sake of simplicity, as I couldn't think of any situations where there
might be a clash between the two identifiers.

Did some light refactoring of window_manager.rs to make the new()
constructor a static method on the WindowManager struct.

Also fixed a bug in Workspace.new_container_for_window where the focused
index was not getting set correctly when the workspace had no
containers.

Cargo.lock

@@ -601,9 +601,9 @@ dependencies = [
 
 [[package]]
 name = "memchr"
-version = "2.4.0"
+version = "2.4.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b16bd47d9e329435e309c58469fe0791c2d0d1ba96ec0954152a5ae2b04387dc"
+checksum = "308cc39be01b73d0d18f82a0e7b2a3df85245f84af96fdddc5d202d27e47b86a"
 
 [[package]]
 name = "memoffset"
@@ -758,9 +758,9 @@ dependencies = [
 
 [[package]]
 name = "object"
-version = "0.26.0"
+version = "0.26.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c55827317fb4c08822499848a14237d2874d6f139828893017237e7ab93eb386"
+checksum = "ee2766204889d09937d00bfbb7fec56bb2a199e2ade963cab19185d8a6104c7c"
 dependencies = [
  "memchr",
 ]

README.md

@@ -118,7 +118,7 @@ Once `komorebi` is running, you can execute the `komorebi.sample.ahk` script to
 If you have AutoHotKey installed and a `komorebi.ahk` file in your home directory (run `$Env:UserProfile` at a
 PowerShell prompt to find your home directory), `komorebi` will automatically try to load it when starting.
 
-There is also tentative support for loading a AutoHotKey v2, if the file is named `komorebi.ahk2` and
+There is also tentative support for loading a AutoHotKey v2 files, if the file is named `komorebi.ahk2` and
 the `AutoHotKey64.exe` executable for AutoHotKey v2 is in your `Path`. If both `komorebi.ahk` and `komorebi.ahk2` files
 exist in your home directory, only `komorebi.ahk` will be loaded. An example of an AutoHotKey v2 configuration file
 for _komorebi_ can be found [here](https://gist.github.com/crosstyan/dafacc0778dabf693ce9236c57b201cd).
@@ -182,6 +182,7 @@ restore-windows              Restore all hidden windows (debugging command)
 reload-configuration         Reload ~/komorebi.ahk (if it exists)
 watch-configuration          Toggle the automatic reloading of ~/komorebi.ahk (if it exists)
 float-rule                   Add a rule to always float the specified application
+workspace-rule               Add a rule to associate an application with a workspace
 identify-tray-application    Identify an application that closes to the system tray
 focus-follows-mouse          Enable or disable focus follows mouse for the operating system
 help                         Print this message or the help of the given subcommand(s)
@@ -206,10 +207,9 @@ help                         Print this message or the help of the given subcomm
 - [x] BSP tree layout
 - [x] Flip BSP tree layout horizontally or vertically
 - [x] Equal-width, max-height column layout
-- [x] Floating rules based on exe name
+- [x] Floating rules based on exe name, window title and class
-- [x] Floating rules based on window title
+- [x] Workspace rules based on exe name and window class
-- [x] Floating rules based on window class
+- [x] Identify 'close/minimize to tray' applications by exe name and class
-- [x] Identify 'close/minimize to tray' applications
 - [x] Toggle floating windows
 - [x] Toggle monocle window
 - [x] Toggle native maximization

komorebi-core/src/lib.rs

@@ -58,6 +58,7 @@ pub enum SocketMessage {
     FloatClass(String),
     FloatExe(String),
     FloatTitle(String),
+    WorkspaceRule(ApplicationIdentifier, String, usize, usize),
     IdentifyTrayApplication(ApplicationIdentifier, String),
     State,
     FocusFollowsMouse(bool),

komorebi.sample.ahk

@@ -29,6 +29,10 @@ Run, komorebic.exe workspace-layout 0 1 columns, , Hide
 ; Set the floaty layout to not tile any windows
 Run, komorebic.exe workspace-tiling 0 4 disable, , Hide
 
+; Always show chat apps on the second workspace
+Run, komorebic.exe workspace-rule exe slack.exe 0 1, , Hide
+Run, komorebic.exe workspace-rule exe Discord.exe 0 1, , Hide
+
 ; Always float IntelliJ popups, matching on class
 Run, komorebic.exe float-rule class SunAwtDialog, , Hide
 ; Always float Control Panel, matching on title

komorebi/Cargo.toml

@@ -30,8 +30,8 @@ tracing = "0.1"
 tracing-appender = "0.1"
 tracing-subscriber = "0.2"
 uds_windows = "1"
-winvd = "0.0.20"
 which = "4"
+winvd = "0.0.20"
 
 [features]
 deadlock_detection = []

komorebi/src/main.rs

@@ -1,6 +1,7 @@
 #![warn(clippy::all, clippy::nursery, clippy::pedantic)]
 #![allow(clippy::missing_errors_doc)]
 
+use std::collections::HashMap;
 use std::process::Command;
 use std::sync::Arc;
 #[cfg(feature = "deadlock_detection")]
@@ -24,6 +25,7 @@ use which::which;
 
 use crate::process_command::listen_for_commands;
 use crate::process_event::listen_for_events;
+use crate::window_manager::WindowManager;
 use crate::window_manager_event::WindowManagerEvent;
 use crate::windows_api::WindowsApi;
 
@@ -66,6 +68,8 @@ lazy_static! {
         "firefox.exe".to_string(),
         "idea64.exe".to_string(),
     ]));
+    static ref WORKSPACE_RULES: Arc<Mutex<HashMap<String, (usize, usize)>>> =
+        Arc::new(Mutex::new(HashMap::new()));
 }
 
 fn setup() -> Result<(WorkerGuard, WorkerGuard)> {
@@ -217,7 +221,7 @@ fn main() -> Result<()> {
             let winevent_listener = winevent_listener::new(Arc::new(Mutex::new(outgoing)));
             winevent_listener.start();
 
-            let wm = Arc::new(Mutex::new(window_manager::new(Arc::new(Mutex::new(
+            let wm = Arc::new(Mutex::new(WindowManager::new(Arc::new(Mutex::new(
                 incoming,
             )))?));
 

komorebi/src/process_command.rs

@@ -21,6 +21,7 @@ use crate::FLOAT_EXES;
 use crate::FLOAT_TITLES;
 use crate::TRAY_AND_MULTI_WINDOW_CLASSES;
 use crate::TRAY_AND_MULTI_WINDOW_EXES;
+use crate::WORKSPACE_RULES;
 
 #[tracing::instrument]
 pub fn listen_for_commands(wm: Arc<Mutex<WindowManager>>) {
@@ -102,6 +103,19 @@ impl WindowManager {
                     float_titles.push(target);
                 }
             }
+            SocketMessage::WorkspaceRule(identifier, id, monitor_idx, workspace_idx) => {
+                match identifier {
+                    ApplicationIdentifier::Exe | ApplicationIdentifier::Class => {
+                        {
+                            let mut workspace_rules = WORKSPACE_RULES.lock();
+                            workspace_rules.insert(id, (monitor_idx, workspace_idx));
+                        }
+
+                        self.enforce_workspace_rules()?;
+                    }
+                    ApplicationIdentifier::Title => {}
+                }
+            }
             SocketMessage::AdjustContainerPadding(sizing, adjustment) => {
                 self.adjust_container_padding(sizing, adjustment)?;
             }

komorebi/src/process_event.rs

@@ -91,8 +91,12 @@ impl WindowManager {
             }
         }
 
+        self.enforce_workspace_rules()?;
+
         if matches!(event, WindowManagerEvent::MouseCapture(..)) {
-            tracing::trace!("only reaping orphans for mouse capture event");
+            tracing::trace!(
+                "only reaping orphans and enforcing workspace rules for mouse capture event"
+            );
             return Ok(());
         }
 

komorebi/src/window_manager.rs

@@ -35,6 +35,7 @@ use crate::FLOAT_TITLES;
 use crate::LAYERED_EXE_WHITELIST;
 use crate::TRAY_AND_MULTI_WINDOW_CLASSES;
 use crate::TRAY_AND_MULTI_WINDOW_EXES;
+use crate::WORKSPACE_RULES;
 
 #[derive(Debug)]
 pub struct WindowManager {
@@ -76,40 +77,64 @@ impl From<&mut WindowManager> for State {
 
 impl_ring_elements!(WindowManager, Monitor);
 
-#[tracing::instrument]
+#[derive(Debug, Clone, Copy)]
-pub fn new(incoming: Arc<Mutex<Receiver<WindowManagerEvent>>>) -> Result<WindowManager> {
+struct EnforceWorkspaceRuleOp {
-    let home = dirs::home_dir().context("there is no home directory")?;
+    hwnd: isize,
-    let mut socket = home;
+    origin_monitor_idx: usize,
-    socket.push("komorebi.sock");
+    origin_workspace_idx: usize,
-    let socket = socket.as_path();
+    target_monitor_idx: usize,
+    target_workspace_idx: usize,
+}
 
-    match std::fs::remove_file(&socket) {
+impl EnforceWorkspaceRuleOp {
-        Ok(_) => {}
+    const fn is_origin(&self, monitor_idx: usize, workspace_idx: usize) -> bool {
-        Err(error) => match error.kind() {
+        self.origin_monitor_idx == monitor_idx && self.origin_workspace_idx == workspace_idx
-            // Doing this because ::exists() doesn't work reliably on Windows via IntelliJ
+    }
-            ErrorKind::NotFound => {}
-            _ => {
-                return Err(error.into());
-            }
-        },
-    };
 
-    let listener = UnixListener::bind(&socket)?;
+    const fn is_target(&self, monitor_idx: usize, workspace_idx: usize) -> bool {
+        self.target_monitor_idx == monitor_idx && self.target_workspace_idx == workspace_idx
+    }
 
-    let virtual_desktop_id = winvd::helpers::get_current_desktop_number()
+    const fn is_enforced(&self) -> bool {
-        .expect("could not determine the current virtual desktop number");
+        (self.origin_monitor_idx == self.target_monitor_idx)
-
+            && (self.origin_workspace_idx == self.target_workspace_idx)
-    Ok(WindowManager {
+    }
-        monitors: Ring::default(),
-        incoming_events: incoming,
-        command_listener: listener,
-        is_paused: false,
-        hotwatch: Hotwatch::new()?,
-        virtual_desktop_id,
-    })
 }
 
 impl WindowManager {
+    #[tracing::instrument]
+    pub fn new(incoming: Arc<Mutex<Receiver<WindowManagerEvent>>>) -> Result<Self> {
+        let home = dirs::home_dir().context("there is no home directory")?;
+        let mut socket = home;
+        socket.push("komorebi.sock");
+        let socket = socket.as_path();
+
+        match std::fs::remove_file(&socket) {
+            Ok(_) => {}
+            Err(error) => match error.kind() {
+                // Doing this because ::exists() doesn't work reliably on Windows via IntelliJ
+                ErrorKind::NotFound => {}
+                _ => {
+                    return Err(error.into());
+                }
+            },
+        };
+
+        let listener = UnixListener::bind(&socket)?;
+
+        let virtual_desktop_id = winvd::helpers::get_current_desktop_number()
+            .expect("could not determine the current virtual desktop number");
+
+        Ok(Self {
+            monitors: Ring::default(),
+            incoming_events: incoming,
+            command_listener: listener,
+            is_paused: false,
+            hotwatch: Hotwatch::new()?,
+            virtual_desktop_id,
+        })
+    }
+
     #[tracing::instrument(skip(self))]
     pub fn init(&mut self) -> Result<()> {
         tracing::info!("initialising");
@@ -192,6 +217,123 @@ impl WindowManager {
         Ok(())
     }
 
+    #[tracing::instrument(skip(self))]
+    pub fn enforce_workspace_rules(&mut self) -> Result<()> {
+        let mut to_move = vec![];
+
+        let focused_monitor_idx = self.focused_monitor_idx();
+        let focused_workspace_idx = self
+            .monitors()
+            .get(focused_monitor_idx)
+            .context("there is no monitor with that index")?
+            .focused_workspace_idx();
+
+        let workspace_rules = WORKSPACE_RULES.lock();
+        // Go through all the monitors and workspaces
+        for (i, monitor) in self.monitors().iter().enumerate() {
+            for (j, workspace) in monitor.workspaces().iter().enumerate() {
+                // And all the visible windows (at the top of a container)
+                for window in workspace.visible_windows().into_iter().flatten() {
+                    // If the executable names or titles of any of those windows are in our rules map
+                    if let Some((monitor_idx, workspace_idx)) = workspace_rules.get(&window.exe()?)
+                    {
+                        tracing::info!(
+                            "{} should be on monitor {}, workspace {}",
+                            window.title()?,
+                            *monitor_idx,
+                            *workspace_idx
+                        );
+
+                        // Create an operation outline and save it for later in the fn
+                        to_move.push(EnforceWorkspaceRuleOp {
+                            hwnd: window.hwnd,
+                            origin_monitor_idx: i,
+                            origin_workspace_idx: j,
+                            target_monitor_idx: *monitor_idx,
+                            target_workspace_idx: *workspace_idx,
+                        });
+                    } else if let Some((monitor_idx, workspace_idx)) =
+                        workspace_rules.get(&window.title()?)
+                    {
+                        tracing::info!(
+                            "{} should be on monitor {}, workspace {}",
+                            window.title()?,
+                            *monitor_idx,
+                            *workspace_idx
+                        );
+
+                        to_move.push(EnforceWorkspaceRuleOp {
+                            hwnd: window.hwnd,
+                            origin_monitor_idx: i,
+                            origin_workspace_idx: j,
+                            target_monitor_idx: *monitor_idx,
+                            target_workspace_idx: *workspace_idx,
+                        });
+                    }
+                }
+            }
+        }
+
+        // Only retain operations where the target is not the current workspace
+        to_move.retain(|op| !op.is_target(focused_monitor_idx, focused_workspace_idx));
+        // Only retain operations where the rule has not already been enforced
+        to_move.retain(|op| !op.is_enforced());
+
+        let mut should_update_focused_workspace = false;
+
+        // Parse the operation and remove any windows that are not placed according to their rules
+        for op in &to_move {
+            let origin_workspace = self
+                .monitors_mut()
+                .get_mut(op.origin_monitor_idx)
+                .context("there is no monitor with that index")?
+                .workspaces_mut()
+                .get_mut(op.origin_workspace_idx)
+                .context("there is no workspace with that index")?;
+
+            // Hide the window we are about to remove if it is on the currently focused workspace
+            if op.is_origin(focused_monitor_idx, focused_workspace_idx) {
+                Window { hwnd: op.hwnd }.hide();
+                should_update_focused_workspace = true;
+            }
+
+            origin_workspace.remove_window(op.hwnd)?;
+        }
+
+        // Parse the operation again and associate those removed windows with the workspace that
+        // their rules have defined for them
+        for op in &to_move {
+            let target_monitor = self
+                .monitors_mut()
+                .get_mut(op.target_monitor_idx)
+                .context("there is no monitor with that index")?;
+
+            // The very first time this fn is called, the workspace might not even exist yet
+            if target_monitor
+                .workspaces()
+                .get(op.target_workspace_idx)
+                .is_none()
+            {
+                // If it doesn't, let's make sure it does for the next step
+                target_monitor.ensure_workspace_count(op.target_workspace_idx + 1);
+            }
+
+            let target_workspace = target_monitor
+                .workspaces_mut()
+                .get_mut(op.target_workspace_idx)
+                .context("there is no workspace with that index")?;
+
+            target_workspace.new_container_for_window(Window { hwnd: op.hwnd });
+        }
+
+        // Only re-tile the focused workspace if we need to
+        if should_update_focused_workspace {
+            self.update_focused_workspace(false)?;
+        }
+
+        Ok(())
+    }
+
     #[tracing::instrument(skip(self))]
     pub fn update_focused_workspace(&mut self, mouse_follows_focus: bool) -> Result<()> {
         tracing::info!("updating");

komorebi/src/workspace.rs

@@ -474,7 +474,11 @@ impl Workspace {
     }
 
     pub fn new_container_for_window(&mut self, window: Window) {
-        let next_idx = self.focused_container_idx() + 1;
+        let next_idx = if self.containers().is_empty() {
+            0
+        } else {
+            self.focused_container_idx() + 1
+        };
 
         let mut container = Container::default();
         container.add_window(window);
@@ -695,6 +699,15 @@ impl Workspace {
         }
     }
 
+    pub fn visible_windows(&self) -> Vec<Option<&Window>> {
+        let mut vec = vec![];
+        for container in self.containers() {
+            vec.push(container.focused_window());
+        }
+
+        vec
+    }
+
     pub fn visible_windows_mut(&mut self) -> Vec<Option<&mut Window>> {
         let mut vec = vec![];
         for container in self.containers_mut() {

komorebic/src/main.rs

@@ -165,6 +165,18 @@ struct ApplicationTarget {
     id: String,
 }
 
+#[derive(Clap)]
+struct WorkspaceRule {
+    #[clap(arg_enum)]
+    identifier: ApplicationIdentifier,
+    /// Identifier as a string
+    id: String,
+    /// Monitor index (zero-indexed)
+    monitor: usize,
+    /// Workspace index on the specified monitor (zero-indexed)
+    workspace: usize,
+}
+
 #[derive(Clap)]
 #[clap(author, about, version, setting = AppSettings::DeriveDisplayOrder)]
 struct Opts {
@@ -265,6 +277,9 @@ enum SubCommand {
     /// Add a rule to always float the specified application
     #[clap(setting = AppSettings::ArgRequiredElseHelp)]
     FloatRule(ApplicationTarget),
+    /// Add a rule to associate an application with a workspace
+    #[clap(setting = AppSettings::ArgRequiredElseHelp)]
+    WorkspaceRule(WorkspaceRule),
     /// Identify an application that closes to the system tray
     #[clap(setting = AppSettings::ArgRequiredElseHelp)]
     IdentifyTrayApplication(ApplicationTarget),
@@ -413,6 +428,12 @@ fn main() -> Result<()> {
                 send_message(&*SocketMessage::FloatTitle(arg.id).as_bytes()?)?;
             }
         },
+        SubCommand::WorkspaceRule(arg) => {
+            send_message(
+                &*SocketMessage::WorkspaceRule(arg.identifier, arg.id, arg.monitor, arg.workspace)
+                    .as_bytes()?,
+            )?;
+        }
         SubCommand::Stack(arg) => {
             send_message(&*SocketMessage::StackWindow(arg.operation_direction).as_bytes()?)?;
         }