feat: refactor + qol

cli/main.rs

@@ -17,20 +17,20 @@ enum Commands {
     Add { tag: u8 },
     Remove { tag: u8 },
     Toggle { tag: u8 },
-    Enable { tag: u8 },
-    Disable { tag: u8 },
-    ToggleWs { tag: u8 },
+    EnableTag { tag: u8 },
+    DisableTag { tag: u8 },
+    ToggleTag { tag: u8 },
 }
 
 impl From<Commands> for niri_tag::TagCmd {
     fn from(value: Commands) -> Self {
         match value {
-            Commands::Add { tag } => TagCmd::Add(tag),
-            Commands::Remove { tag } => TagCmd::Remove(tag),
-            Commands::Enable { tag } => TagCmd::Enable(tag),
-            Commands::Disable { tag } => TagCmd::Disable(tag),
-            Commands::Toggle { tag } => TagCmd::Toggle(tag),
-            Commands::ToggleWs { tag } => TagCmd::ToggleWs(tag),
+            Commands::Add { tag } => TagCmd::AddTagToWin(tag),
+            Commands::Remove { tag } => TagCmd::RemoveTagFromWin(tag),
+            Commands::Toggle { tag } => TagCmd::ToggleTagOnWin(tag),
+            Commands::EnableTag { tag } => TagCmd::EnableTag(tag),
+            Commands::DisableTag { tag } => TagCmd::DisableTag(tag),
+            Commands::ToggleTag { tag } => TagCmd::ToggleTag(tag),
         }
     }
 }
@@ -39,14 +39,13 @@ fn main() -> Result<()> {
     let cli = Cli::parse();
     use Commands::*;
 
-    println!("{:?}", cli.cmd);
     match cli.cmd {
         Add { tag } if tag > 0 => (),
         Remove { tag } if tag > 0 => (),
-        Enable { tag } if tag > 0 => (),
-        Disable { tag } if tag > 0 => (),
         Toggle { tag } if tag > 0 => (),
-        ToggleWs { tag } if tag > 0 => (),
+        EnableTag { tag } if tag > 0 => (),
+        DisableTag { tag } if tag > 0 => (),
+        ToggleTag { tag } if tag > 0 => (),
         _ => return Err(anyhow!("Can't change tag 0!")),
     };
 

daemon/main.rs

@@ -2,7 +2,7 @@ mod listeners;
 mod manager;
 mod socket;
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 
 fn main() -> Result<()> {
     // let systemd know we're ready
@@ -20,6 +20,11 @@ fn main() -> Result<()> {
     let (ipc_tx, ipc_rx) = smol::channel::unbounded();
     smol::spawn(listeners::ipc_listener(ipc_tx)).detach();
     // begin managing niri tags
-    let niri_tag = manager::NiriTag::new();
-    smol::block_on(niri_tag.manage_tags(event_rx, ipc_rx))
+    smol::block_on(async {
+        let niri_tag = manager::NiriTag::new()
+            .await
+            .context("Initialising niri tag manager")
+            .unwrap();
+        niri_tag.manage_tags(event_rx, ipc_rx).await
+    })
 }

daemon/manager.rs

@@ -1,7 +1,7 @@
 use crate::socket::{create_niri_socket, query, tell};
-use anyhow::{Result, anyhow};
+use anyhow::{Context, Result, anyhow};
 use niri_ipc::{
-    Action, Event, Reply, Request, Response, Window, WorkspaceReferenceArg,
+    Action, Event, Reply, Request, Response, Window, Workspace, WorkspaceReferenceArg,
     state::{EventStreamState, EventStreamStatePart},
 };
 use niri_tag::TagCmd;
@@ -16,46 +16,266 @@ use std::collections::HashMap;
 pub struct NiriTag {
     tags: HashMap<u8, bool>,
     windows: HashMap<u64, u8>,
+    state: EventStreamState,
+    socket: BufReader<UnixStream>,
+}
+
+enum TagAction {
+    ChangeWindow(u64),
+    ChangeTag(u8),
 }
 
 impl NiriTag {
-    pub fn new() -> Self {
-        Self {
+    pub async fn new() -> Result<Self> {
+        Ok(Self {
             tags: HashMap::new(),
             windows: HashMap::new(),
+            state: EventStreamState::default(),
+            socket: create_niri_socket().await?,
+        })
+    }
+
+    async fn do_action(&mut self, action: TagAction) -> Result<()> {
+        use TagAction::*;
+        let same_output =
+            |wsid: u64, candidates: &HashMap<&u64, &Workspace>| -> Result<Workspace> {
+                candidates
+                    .values()
+                    .filter_map(|ws| {
+                        let output = ws.output.clone()?;
+                        let win_output = self
+                            .state
+                            .workspaces
+                            .workspaces
+                            .get(&wsid)?
+                            .output
+                            .clone()?;
+                        (win_output == output).then_some(ws)
+                    })
+                    .last()
+                    .context(anyhow!(
+                        "No inactive workspaces on output of workspace {} found",
+                        wsid
+                    ))
+                    .copied()
+                    .cloned()
+            };
+        let (active, inactive): (HashMap<_, _>, HashMap<_, _>) = self
+            .state
+            .workspaces
+            .workspaces
+            .iter()
+            .partition(|(_, ws)| ws.is_active);
+        match action {
+            ChangeWindow(wid) => {
+                let current_tag = *self.windows.entry(wid).or_insert(0);
+                let tag_visible = *self.tags.entry(current_tag).or_insert(true);
+                let win = self
+                    .state
+                    .windows
+                    .windows
+                    .get(&wid)
+                    .ok_or(anyhow!("Failed to retrieve window {} from niri state", wid))?;
+                let wsid: u64 = win
+                    .workspace_id
+                    .ok_or(anyhow!("Retrieving workspace id of a changed window"))?;
+                let win_visible = active.contains_key(&wsid);
+                match (win_visible, tag_visible) {
+                    (true, false) => {
+                        let inactive_same_output = same_output(wsid, &inactive)?;
+                        tell(
+                            &mut self.socket,
+                            Request::Action(Action::MoveWindowToWorkspace {
+                                window_id: Some(wid),
+                                reference: WorkspaceReferenceArg::Id(inactive_same_output.id),
+                                focus: false,
+                            }),
+                        )
+                        .await
+                    }
+                    (false, true) => {
+                        let active_same_output = same_output(wsid, &active)?;
+                        tell(
+                            &mut self.socket,
+                            Request::Action(Action::MoveWindowToWorkspace {
+                                window_id: Some(wid),
+                                reference: WorkspaceReferenceArg::Id(active_same_output.id),
+                                focus: true,
+                            }),
+                        )
+                        .await
+                    }
+                    _ => Ok(()),
+                }
+            }
+            ChangeTag(tag) => {
+                tracing::debug!("Changing tag {}", tag);
+                let tag_visible = *self.tags.entry(tag).or_insert(true);
+                tracing::debug!("Windows: {:?}", self.windows);
+                let affected_windows: Vec<u64> = self
+                    .windows
+                    .iter()
+                    .filter(|(_, t)| tag == **t)
+                    .map(|(wid, _)| *wid)
+                    .collect();
+                tracing::debug!(
+                    "{} affected windows of tag {}: {:?}",
+                    affected_windows.len(),
+                    tag,
+                    affected_windows
+                );
+                let focus = affected_windows.last().cloned();
+                for wid in affected_windows {
+                    tracing::debug!("Changing affected window {}", wid);
+                    if let Some(win) = self.state.windows.windows.get(&wid) {
+                        let wsid = win.workspace_id.unwrap();
+                        match same_output(wsid, if tag_visible { &active } else { &inactive }) {
+                            Ok(status_same_output) => {
+                                if let Err(e) = tell(
+                                    &mut self.socket,
+                                    Request::Action(Action::MoveWindowToWorkspace {
+                                        window_id: Some(wid),
+                                        reference: WorkspaceReferenceArg::Id(status_same_output.id),
+                                        focus: false,
+                                    }),
+                                )
+                                .await
+                                {
+                                    tracing::error!(
+                                        "Failed to move window {} to workspace {}: {}",
+                                        wid,
+                                        status_same_output.id,
+                                        e
+                                    );
+                                }
+                            }
+                            Err(e) => tracing::error!(
+                                "Failed to get workspace on same output as {}: {}",
+                                wsid,
+                                e
+                            ),
+                        }
+                    } else {
+                        tracing::warn!("Failed to get wid {} from niri state", wid);
+                        continue;
+                    }
+                }
+                if let Some(focus) = focus {
+                    if tag_visible {
+                        tell(
+                            &mut self.socket,
+                            Request::Action(Action::FocusWindow { id: focus }),
+                        )
+                        .await?;
+                    }
+                }
+                Ok(())
+            }
+        }
+    }
+
+    async fn get_focused_window(&mut self) -> Result<Window> {
+        let q = query(&mut self.socket, Request::FocusedWindow).await?;
+        if let Reply::Ok(Response::FocusedWindow(win)) = q {
+            if let Some(win) = win {
+                Ok(win)
+            } else {
+                Err(anyhow!("No focused window to operate on"))
+            }
+        } else {
+            Err(anyhow!(
+                "Invalid response from Niri when requesting FocusedWindow: {}",
+                if q.is_err() {
+                    q.unwrap_err()
+                } else {
+                    serde_json::to_string(&q.unwrap())?
+                }
+            ))
+        }
+    }
+
+    async fn handle_recvd(&mut self, recvd: Receivable) -> Result<()> {
+        use TagAction::*;
+        // first do any local mutations
+        let action: TagAction = match recvd {
+            Receivable::Event(ev) => {
+                let _ = self.state.apply(ev.clone());
+                return self.handle_event(ev).await;
+            }
+            Receivable::TagCmd(cmd) => match cmd {
+                TagCmd::AddTagToWin(t) => {
+                    let win = self.get_focused_window().await?;
+                    let wid = win.id;
+                    self.windows.insert(wid, t);
+                    tracing::debug!("adding tag {} to {}", t, wid);
+                    ChangeWindow(wid)
+                }
+                TagCmd::RemoveTagFromWin(_) => {
+                    let win = self.get_focused_window().await?;
+                    let wid = win.id;
+                    self.windows.insert(wid, 0);
+                    tracing::debug!("resetting tag on {}", wid);
+                    ChangeWindow(wid)
+                }
+                TagCmd::ToggleTagOnWin(t) => {
+                    let win = self.get_focused_window().await?;
+                    let wid = win.id;
+                    tracing::debug!("{} has tag {:?}", wid, self.windows.get(&wid));
+                    let this_tag = *self.windows.entry(wid).or_insert(0);
+                    let toggle = if this_tag == t { 0 } else { t };
+                    tracing::debug!("toggling {} to tag {}", wid, toggle);
+                    self.windows.insert(wid, toggle);
+                    ChangeWindow(wid)
+                }
+
+                TagCmd::EnableTag(t) => {
+                    self.tags.insert(t, true);
+                    ChangeTag(t)
+                }
+                TagCmd::DisableTag(t) => {
+                    self.tags.insert(t, false);
+                    ChangeTag(t)
+                }
+                TagCmd::ToggleTag(t) => {
+                    let visible = *self.tags.entry(t).or_insert(false);
+                    tracing::debug!("toggling tag {} to {}", t, !visible);
+                    self.tags.insert(t, !visible);
+                    ChangeTag(t)
+                }
+            },
+        };
+        // then arrange corresponding state in the compositor
+        self.do_action(action).await
+    }
+
+    async fn handle_event(&mut self, ev: Event) -> Result<()> {
+        use Event::*;
+        match ev {
+            WindowOpenedOrChanged { window } => {
+                self.windows.entry(window.id).or_insert(0);
+                Ok(())
+            }
+            WindowClosed { id } => {
+                self.windows.remove(&id);
+                Ok(())
+            }
+            // WorkspaceActivated { .. } => (),
+            // WorkspacesChanged { .. } => (),
+            // WorkspaceUrgencyChanged { .. } => (),
+            // WindowsChanged { .. } => (),
+            // WindowUrgencyChanged { .. } => (),
+            _ => Ok(()),
         }
     }
 
     #[allow(unreachable_code)]
     pub async fn manage_tags(
-        self,
+        mut self,
         ev_rx: channel::Receiver<Event>,
         tag_rx: channel::Receiver<TagCmd>,
     ) -> Result<()> {
-        let mut state = EventStreamState::default();
-        let mut tags = NiriTag::new();
-        let mut socket = create_niri_socket().await?;
         // base tag is always visible
-        tags.tags.insert(0, true);
-
-        async fn on_focused_window<O>(socket: &mut BufReader<UnixStream>, mut ok: O) -> Result<()>
-        where
-            O: FnMut(Option<Window>) -> Result<()>,
-        {
-            let q = query(socket, Request::FocusedWindow).await?;
-            if let Reply::Ok(Response::FocusedWindow(win)) = q {
-                ok(win)
-            } else {
-                Err(anyhow!(
-                    "Invalid response from Niri when requesting FocusedWindow: {}",
-                    if q.is_err() {
-                        q.unwrap_err()
-                    } else {
-                        serde_json::to_string(&q.unwrap())?
-                    }
-                ))
-            }
-        }
+        self.tags.insert(0, true);
 
         loop {
             let recvd: Receivable =
@@ -65,130 +285,14 @@ impl NiriTag {
                 .await?;
             tracing::debug!("received {:?}", recvd);
 
-            let res = match recvd {
-                Receivable::Event(ev) => {
-                    let _ = state.apply(ev.clone());
-                    Ok(())
-                }
-                Receivable::TagCmd(cmd) => match cmd {
-                    TagCmd::Add(t) => {
-                        on_focused_window(&mut socket, |win| {
-                            if let Some(win) = win {
-                                let wid = win.id;
-                                tags.windows.insert(wid, t);
-                                Ok(())
-                            } else {
-                                Err(anyhow!("No focused window to tag"))
-                            }
-                        })
-                        .await
-                    }
-                    TagCmd::Remove(_) => {
-                        on_focused_window(&mut socket, |win| {
-                            if let Some(win) = win {
-                                let wid = win.id;
-                                tags.windows.insert(wid, 0);
-                                Ok(())
-                            } else {
-                                Err(anyhow!("No focused window to untag"))
-                            }
-                        })
-                        .await
-                    }
-                    TagCmd::Toggle(t) => {
-                        on_focused_window(&mut socket, |win| {
-                            if let Some(win) = win {
-                                let wid = win.id;
-                                let toggle = if *tags.windows.get(&wid).unwrap_or(&0) == t {
-                                    0
-                                } else {
-                                    t
-                                };
-                                tracing::debug!("toggling {} to tag {}", wid, toggle);
-                                tags.windows.insert(wid, toggle);
-                                Ok(())
-                            } else {
-                                Err(anyhow!("No focused window to untag"))
-                            }
-                        })
-                        .await
-                    }
-
-                    TagCmd::Enable(t) => {
-                        tags.tags.insert(t, true);
-                        Ok(())
-                    }
-                    TagCmd::Disable(t) => {
-                        tags.tags.insert(t, false);
-                        Ok(())
-                    }
-                    TagCmd::ToggleWs(t) => {
-                        tracing::debug!("toggling tag {}", t);
-                        tags.tags.insert(t, !tags.tags.get(&t).unwrap_or(&false));
-                        Ok(())
-                    }
-                },
-            };
+            let res = self.handle_recvd(recvd).await;
             match res {
                 Ok(()) => (),
                 Err(e) => tracing::error!("error occurred in manager loop: {}", e),
             }
-            // TODO: react selectively instead of brute forcing window state
-            // use Event::*;
-            // match ev {
-            //     WorkspaceActivated { .. } => (),
-            //     WorkspacesChanged { .. } => (),
-            //     WorkspaceUrgencyChanged { .. } => (),
-            //     WindowsChanged { .. } => (),
-            //     WindowOpenedOrChanged { .. } => (),
-            //     WindowUrgencyChanged { .. } => (),
-            //     WindowClosed { .. } => (),
-            //     _ => (),
-            // }
-            for (&wid, window) in state.windows.windows.iter() {
-                let (active, inactive): (Vec<_>, Vec<_>) = state
-                    .workspaces
-                    .workspaces
-                    .iter()
-                    .map(|(wsid, ws)| (wsid, ws.is_active))
-                    .partition(|(_, a)| *a);
-                if let Some(wsid) = window.workspace_id {
-                    if let Some(&window_tag) = tags.windows.get(&wid) {
-                        if let Some(&tag_enabled) = tags.tags.get(&window_tag) {
-                            if tag_enabled && inactive.contains(&(&wsid, false)) {
-                                tell(
-                                    &mut socket,
-                                    Request::Action(Action::MoveWindowToWorkspace {
-                                        window_id: Some(wid),
-                                        reference: WorkspaceReferenceArg::Index(0),
-                                        focus: false,
-                                    }),
-                                )
-                                .await?;
-                                tracing::debug!("making visible {}", wid);
-                            } else if !tag_enabled && active.contains(&(&wsid, true)) {
-                                let hidden = *inactive.first().unwrap().0;
-                                tell(
-                                    &mut socket,
-                                    Request::Action(Action::MoveWindowToWorkspace {
-                                        window_id: Some(wid),
-                                        reference: WorkspaceReferenceArg::Id(hidden),
-                                        focus: false,
-                                    }),
-                                )
-                                .await?;
-                                tracing::debug!("making hidden {}", wid);
-                            }
-                        } else {
-                            tags.windows.insert(wid, 0);
-                        }
-                    }
-                }
-            }
         }
         tracing::error!("Manager loop ended");
         unreachable!("Manager loop ended");
-        Ok(())
     }
 }
 

lib/main.rs

@@ -2,10 +2,10 @@ use serde::{Deserialize, Serialize};
 
 #[derive(Serialize, Deserialize, Debug)]
 pub enum TagCmd {
-    Add(u8),
-    Remove(u8),
-    Enable(u8),
-    Disable(u8),
-    Toggle(u8),
-    ToggleWs(u8),
+    AddTagToWin(u8),
+    RemoveTagFromWin(u8),
+    EnableTag(u8),
+    DisableTag(u8),
+    ToggleTagOnWin(u8),
+    ToggleTag(u8),
 }

module.nix

@@ -38,5 +38,6 @@ in
           PrivateTmp = true;
         };
       };
+      environment.systemPackages = [ cfg.package ];
     };
 }