feat(workspaces): niri support

Co-authored-by: Jake Stanger <mail@jstanger.dev>

src/clients/compositor/hyprland.rs

@@ -327,11 +327,8 @@ impl Client {
 }
 
 impl WorkspaceClient for Client {
-    fn focus(&self, id: String) {
-        let identifier = id.parse::<i32>().map_or_else(
-            |_| WorkspaceIdentifierWithSpecial::Name(&id),
-            WorkspaceIdentifierWithSpecial::Id,
-        );
+    fn focus(&self, id: i64) {
+        let identifier = WorkspaceIdentifierWithSpecial::Id(id as i32);
 
         if let Err(e) = Dispatch::call(DispatchType::Workspace(identifier)) {
             error!("Couldn't focus workspace '{id}': {e:#}");

src/clients/compositor/mod.rs

@@ -8,6 +8,8 @@ use tracing::debug;
 
 #[cfg(feature = "workspaces+hyprland")]
 pub mod hyprland;
+#[cfg(feature = "workspaces+niri")]
+pub mod niri;
 #[cfg(feature = "workspaces+sway")]
 pub mod sway;
 
@@ -16,6 +18,8 @@ pub enum Compositor {
     Sway,
     #[cfg(feature = "workspaces+hyprland")]
     Hyprland,
+    #[cfg(feature = "workspaces+niri")]
+    Niri,
     Unsupported,
 }
 
@@ -29,6 +33,8 @@ impl Display for Compositor {
                 Self::Sway => "Sway",
                 #[cfg(feature = "workspaces+hyprland")]
                 Self::Hyprland => "Hyprland",
+                #[cfg(feature = "workspaces+niri")]
+                Self::Niri => "Niri",
                 Self::Unsupported => "Unsupported",
             }
         )
@@ -49,6 +55,11 @@ impl Compositor {
                 if #[cfg(feature = "workspaces+hyprland")] { Self::Hyprland }
                 else { tracing::error!("Not compiled with Hyprland support"); Self::Unsupported }
             }
+        } else if std::env::var("NIRI_SOCKET").is_ok() {
+            cfg_if! {
+                if #[cfg(feature = "workspaces+niri")] { Self::Niri }
+                else {tracing::error!("Not compiled with Niri support"); Self::Unsupported }
+            }
         } else {
             Self::Unsupported
         }
@@ -68,7 +79,7 @@ impl Compositor {
             Self::Hyprland => clients
                 .hyprland()
                 .map(|client| client as Arc<dyn KeyboardLayoutClient + Send + Sync>),
-            Self::Unsupported => Err(Report::msg("Unsupported compositor").note(
+            Self::Niri | Self::Unsupported => Err(Report::msg("Unsupported compositor").note(
                 "Currently keyboard layout functionality are only supported by Sway and Hyprland",
             )),
         }
@@ -90,8 +101,10 @@ impl Compositor {
             Self::Hyprland => clients
                 .hyprland()
                 .map(|client| client as Arc<dyn WorkspaceClient + Send + Sync>),
+            #[cfg(feature = "workspaces+niri")]
+            Self::Niri => Ok(Arc::new(niri::Client::new())),
             Self::Unsupported => Err(Report::msg("Unsupported compositor")
-                .note("Currently workspaces are only supported by Sway and Hyprland")),
+                .note("Currently workspaces are only supported by Sway, Niri and Hyprland")),
         }
     }
 }
@@ -125,6 +138,10 @@ impl Visibility {
         Self::Visible { focused: true }
     }
 
+    pub fn is_visible(self) -> bool {
+        matches!(self, Self::Visible { .. })
+    }
+
     pub fn is_focused(self) -> bool {
         if let Self::Visible { focused } = self {
             focused
@@ -170,8 +187,8 @@ pub enum WorkspaceUpdate {
 }
 
 pub trait WorkspaceClient: Debug + Send + Sync {
-    /// Requests the workspace with this name is focused.
-    fn focus(&self, name: String);
+    /// Requests the workspace with this id is focused.
+    fn focus(&self, id: i64);
 
     /// Creates a new to workspace event receiver.
     fn subscribe(&self) -> broadcast::Receiver<WorkspaceUpdate>;

src/clients/compositor/niri/connection.rs

@@ -0,0 +1,117 @@
+/// Taken from the `niri_ipc` crate.
+/// Only a relevant snippet has been extracted
+/// to reduce compile times.
+use crate::clients::compositor::Workspace as IronWorkspace;
+use crate::{await_sync, clients::compositor::Visibility};
+use color_eyre::eyre::{eyre, Result};
+use core::str;
+use serde::{Deserialize, Serialize};
+use std::{env, path::Path};
+use tokio::{
+    io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
+    net::UnixStream,
+};
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub enum Request {
+    Action(Action),
+    EventStream,
+}
+
+pub type Reply = Result<Response, String>;
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub enum Response {
+    Handled,
+    Workspaces(Vec<Workspace>),
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub enum Action {
+    FocusWorkspace { reference: WorkspaceReferenceArg },
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
+pub enum WorkspaceReferenceArg {
+    Name(String),
+    Id(u64),
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct Workspace {
+    pub id: u64,
+    pub idx: u8,
+    pub name: Option<String>,
+    pub output: Option<String>,
+    pub is_active: bool,
+    pub is_focused: bool,
+}
+
+impl From<&Workspace> for IronWorkspace {
+    fn from(workspace: &Workspace) -> IronWorkspace {
+        // Workspaces in niri don't neccessarily have names.
+        // If the niri workspace has a name then it is assigned as is,
+        // but if it does not have a name, the monitor index is used.
+        Self {
+            id: workspace.id as i64,
+            name: workspace.name.clone().unwrap_or(workspace.idx.to_string()),
+            monitor: workspace.output.clone().unwrap_or_default(),
+            visibility: if workspace.is_active {
+                Visibility::Visible {
+                    focused: workspace.is_focused,
+                }
+            } else {
+                Visibility::Hidden
+            },
+        }
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub enum Event {
+    WorkspacesChanged { workspaces: Vec<Workspace> },
+    WorkspaceActivated { id: u64, focused: bool },
+    Other,
+}
+
+#[derive(Debug)]
+pub struct Connection(UnixStream);
+impl Connection {
+    pub async fn connect() -> Result<Self> {
+        let socket_path =
+            env::var_os("NIRI_SOCKET").ok_or_else(|| eyre!("NIRI_SOCKET not found!"))?;
+        Self::connect_to(socket_path).await
+    }
+
+    pub async fn connect_to(path: impl AsRef<Path>) -> Result<Self> {
+        let raw_stream = UnixStream::connect(path.as_ref()).await?;
+        let stream = raw_stream;
+        Ok(Self(stream))
+    }
+
+    pub async fn send(
+        &mut self,
+        request: Request,
+    ) -> Result<(Reply, impl FnMut() -> Result<Event> + '_)> {
+        let Self(stream) = self;
+        let mut buf = serde_json::to_string(&request)?;
+
+        stream.write_all(buf.as_bytes()).await?;
+        stream.shutdown().await?;
+
+        buf.clear();
+        let mut reader = BufReader::new(stream);
+        reader.read_line(&mut buf).await?;
+        let reply = serde_json::from_str(&buf)?;
+
+        let events = move || {
+            buf.clear();
+            await_sync(async {
+                reader.read_line(&mut buf).await.unwrap_or(0);
+            });
+            let event: Event = serde_json::from_str(&buf).unwrap_or(Event::Other);
+            Ok(event)
+        };
+        Ok((reply, events))
+    }
+}

src/clients/compositor/niri/mod.rs

@@ -0,0 +1,183 @@
+use crate::{clients::compositor::Visibility, send, spawn};
+use color_eyre::Report;
+use tracing::{error, warn};
+
+use tokio::sync::broadcast;
+
+use super::{Workspace as IronWorkspace, WorkspaceClient, WorkspaceUpdate};
+mod connection;
+
+use connection::{Action, Connection, Event, Request, WorkspaceReferenceArg};
+
+#[derive(Debug)]
+pub struct Client {
+    tx: broadcast::Sender<WorkspaceUpdate>,
+    _rx: broadcast::Receiver<WorkspaceUpdate>,
+}
+
+impl Client {
+    pub fn new() -> Self {
+        let (tx, rx) = broadcast::channel(32);
+        let tx2 = tx.clone();
+
+        spawn(async move {
+            let mut conn = Connection::connect().await?;
+            let (_, mut event_listener) = conn.send(Request::EventStream).await?;
+
+            let mut workspace_state: Vec<IronWorkspace> = Vec::new();
+            let mut first_event = true;
+
+            loop {
+                let events = match event_listener() {
+                    Ok(Event::WorkspacesChanged { workspaces }) => {
+                        // Niri only has a WorkspacesChanged Event and Ironbar has 4 events which have to be handled: Add, Remove, Rename and Move.
+                        // This is handled by keeping a previous state of workspaces and comparing with the new state for changes.
+                        let new_workspaces: Vec<IronWorkspace> = workspaces
+                            .into_iter()
+                            .map(|w| IronWorkspace::from(&w))
+                            .collect();
+
+                        let mut updates: Vec<WorkspaceUpdate> = vec![];
+
+                        if first_event {
+                            updates.push(WorkspaceUpdate::Init(new_workspaces.clone()));
+                            first_event = false;
+                        } else {
+                            // first pass - add/update
+                            for workspace in &new_workspaces {
+                                let old_workspace =
+                                    workspace_state.iter().find(|w| w.id == workspace.id);
+
+                                match old_workspace {
+                                    None => updates.push(WorkspaceUpdate::Add(workspace.clone())),
+                                    Some(old_workspace) => {
+                                        if workspace.name != old_workspace.name {
+                                            updates.push(WorkspaceUpdate::Rename {
+                                                id: workspace.id,
+                                                name: workspace.name.clone(),
+                                            });
+                                        }
+
+                                        if workspace.monitor != old_workspace.monitor {
+                                            updates.push(WorkspaceUpdate::Move(workspace.clone()));
+                                        }
+                                    }
+                                }
+                            }
+
+                            // second pass - delete
+                            for workspace in &workspace_state {
+                                let exists = new_workspaces.iter().any(|w| w.id == workspace.id);
+
+                                if !exists {
+                                    updates.push(WorkspaceUpdate::Remove(workspace.id));
+                                }
+                            }
+                        }
+
+                        workspace_state = new_workspaces;
+                        updates
+                    }
+
+                    Ok(Event::WorkspaceActivated { id, focused }) => {
+                        // workspace with id is activated, if focus is true then it is also focused
+                        // if focused is true then focus has changed => find old focused workspace. set it to inactive and set current
+                        //
+                        // we use indexes here as both new/old need to be mutable
+
+                        if let Some(new_index) =
+                            workspace_state.iter().position(|w| w.id == id as i64)
+                        {
+                            if focused {
+                                if let Some(old_index) = workspace_state
+                                    .iter()
+                                    .position(|w| w.visibility.is_focused())
+                                {
+                                    workspace_state[new_index].visibility = Visibility::focused();
+
+                                    if workspace_state[old_index].monitor
+                                        == workspace_state[new_index].monitor
+                                    {
+                                        workspace_state[old_index].visibility = Visibility::Hidden;
+                                    } else {
+                                        workspace_state[old_index].visibility =
+                                            Visibility::visible();
+                                    }
+
+                                    vec![WorkspaceUpdate::Focus {
+                                        old: Some(workspace_state[old_index].clone()),
+                                        new: workspace_state[new_index].clone(),
+                                    }]
+                                } else {
+                                    workspace_state[new_index].visibility = Visibility::focused();
+
+                                    vec![WorkspaceUpdate::Focus {
+                                        old: None,
+                                        new: workspace_state[new_index].clone(),
+                                    }]
+                                }
+                            } else {
+                                // if focused is false means active workspace on a particular monitor has changed =>
+                                // change all workspaces on monitor to inactive and change current workspace as active
+                                workspace_state[new_index].visibility = Visibility::visible();
+
+                                if let Some(old_index) = workspace_state.iter().position(|w| {
+                                    (w.visibility.is_focused() || w.visibility.is_visible())
+                                        && w.monitor == workspace_state[new_index].monitor
+                                }) {
+                                    workspace_state[old_index].visibility = Visibility::Hidden;
+
+                                    vec![]
+                                } else {
+                                    vec![]
+                                }
+                            }
+                        } else {
+                            warn!("No workspace with id for new focus/visible workspace found");
+                            vec![]
+                        }
+                    }
+                    Ok(Event::Other) => {
+                        vec![]
+                    }
+                    Err(err) => {
+                        error!("{err:?}");
+                        break;
+                    }
+                };
+
+                for event in events {
+                    send!(tx, event);
+                }
+            }
+
+            Ok::<(), Report>(())
+        });
+
+        Self { tx: tx2, _rx: rx }
+    }
+}
+
+impl WorkspaceClient for Client {
+    fn focus(&self, id: i64) {
+        // this does annoyingly require spawning a separate connection for every focus call
+        // the alternative is sticking the conn behind a mutex which could perform worse
+        spawn(async move {
+            let mut conn = Connection::connect().await?;
+
+            let command = Request::Action(Action::FocusWorkspace {
+                reference: WorkspaceReferenceArg::Id(id as u64),
+            });
+
+            if let Err(err) = conn.send(command).await {
+                error!("failed to send command: {err:?}");
+            }
+
+            Ok::<(), Report>(())
+        });
+    }
+
+    fn subscribe(&self) -> broadcast::Receiver<WorkspaceUpdate> {
+        self.tx.subscribe()
+    }
+}

src/clients/compositor/sway.rs

@@ -2,20 +2,36 @@ use super::{
     KeyboardLayoutClient, KeyboardLayoutUpdate, Visibility, Workspace, WorkspaceClient,
     WorkspaceUpdate,
 };
+use crate::clients::sway::Client;
 use crate::{await_sync, error, send, spawn};
+use color_eyre::Report;
 use swayipc_async::{InputChange, InputEvent, Node, WorkspaceChange, WorkspaceEvent};
 use tokio::sync::broadcast::{channel, Receiver};
 
-use crate::clients::sway::Client;
-
 impl WorkspaceClient for Client {
-    fn focus(&self, id: String) {
+    fn focus(&self, id: i64) {
         let client = self.connection().clone();
         spawn(async move {
             let mut client = client.lock().await;
-            if let Err(e) = client.run_command(format!("workspace {id}")).await {
-                error!("Couldn't focus workspace '{id}': {e:#}");
+
+            let name = client
+                .get_workspaces()
+                .await?
+                .into_iter()
+                .find(|w| w.id == id)
+                .map(|w| w.name);
+
+            let Some(name) = name else {
+                return Err(Report::msg(format!("couldn't find workspace with id {id}")));
+            };
+
+            if let Err(e) = client.run_command(format!("workspace {name}")).await {
+                return Err(Report::msg(format!(
+                    "Couldn't focus workspace '{id}': {e:#}"
+                )));
             }
+
+            Ok(())
         });
     }
 
@@ -24,6 +40,7 @@ impl WorkspaceClient for Client {
 
         let client = self.connection().clone();
 
+        // TODO: this needs refactoring
         await_sync(async {
             let mut client = client.lock().await;
             let workspaces = client.get_workspaces().await.expect("to get workspaces");
@@ -35,7 +52,7 @@ impl WorkspaceClient for Client {
 
             drop(client);
 
-            self.add_listener::<swayipc_async::WorkspaceEvent>(move |event| {
+            self.add_listener::<WorkspaceEvent>(move |event| {
                 let update = WorkspaceUpdate::from(event.clone());
                 send!(tx, update);
             })