Add support for niri workspaces (#39)

src/shared_state.rs

@@ -23,4 +23,8 @@ impl SharedState {
     pub fn get_hyprland(&mut self) -> Option<&mut wm_info_provider::HyprlandInfoProvider> {
         self.wm_info_provider.as_any().downcast_mut()
     }
+
+    pub fn get_niri(&mut self) -> Option<&mut wm_info_provider::NiriInfoProvider> {
+        self.wm_info_provider.as_any().downcast_mut()
+    }
 }

src/wm_info_provider.rs

@@ -19,6 +19,9 @@ pub use river::*;
 mod hyprland;
 pub use hyprland::*;
 
+mod niri;
+pub use niri::*;
+
 pub trait WmInfoProvider {
     fn register(&self, _: &mut EventLoop) {}
 
@@ -62,6 +65,10 @@ pub fn bind(
         return Box::new(hyprland);
     }
 
+    if let Some(niri) = NiriInfoProvider::new() {
+        return Box::new(niri);
+    }
+
     Box::new(DummyInfoProvider)
 }
 

src/wm_info_provider/niri.rs

@@ -0,0 +1,232 @@
+#![allow(clippy::collapsible_else_if)]
+
+use std::io::{self, Write};
+use std::os::fd::AsRawFd;
+use std::os::unix::net::UnixStream;
+use std::path::PathBuf;
+
+use serde::de::IgnoredAny;
+
+use super::*;
+use crate::event_loop;
+use crate::utils::read_to_vec;
+
+pub struct NiriInfoProvider {
+    ipc: Ipc,
+    workspaces: Vec<IpcWorkspace>,
+}
+
+impl NiriInfoProvider {
+    pub fn new() -> Option<Self> {
+        let ns = std::env::var("NIRI_SOCKET").ok()?;
+        let ipc = Ipc::new(&ns)?;
+        Some(Self {
+            workspaces: Vec::new(),
+            ipc,
+        })
+    }
+
+    fn set_workspace(&self, idx: u32) {
+        let _ = self.ipc.exec(&format!(
+            r#"{{"Action":{{"FocusWorkspace":{{"reference":{{"Index":{idx}}}}}}}}}"#
+        ));
+    }
+}
+
+impl WmInfoProvider for NiriInfoProvider {
+    fn register(&self, event_loop: &mut EventLoop) {
+        event_loop.register_with_fd(self.ipc.sock.as_raw_fd(), |ctx| {
+            match niri_cb(ctx.conn, ctx.state) {
+                Ok(()) => Ok(event_loop::Action::Keep),
+                Err(e) => {
+                    ctx.state.set_error(ctx.conn, "niri", e);
+                    Ok(event_loop::Action::Unregister)
+                }
+            }
+        });
+    }
+
+    fn get_tags(&self, output: &Output) -> Vec<Tag> {
+        // Niri always generates an empty workspace rather than having an explicit workspace
+        // creation command, so we make the last workspace active only if the user is looking at
+        // it. This makes the behavior of `hide_inactive_tags` useful for Niri. Because we're
+        // looking for the last element, we have to create an intermediate vector to get the
+        // length.
+        let output_workspaces: Vec<_> = self
+            .workspaces
+            .iter()
+            .filter(|ws| ws.output == output.name)
+            .collect();
+        output_workspaces
+            .iter()
+            .enumerate()
+            .map(|(i, ws)| Tag {
+                id: ws.idx,
+                name: ws.name.clone().map_or_else(
+                    || ws.idx.to_string(),
+                    |name| format!("{0} / {1}", ws.idx, name),
+                ),
+                is_focused: ws.is_active,
+                is_active: i < output_workspaces.len() - 1 || ws.is_focused,
+                is_urgent: false,
+            })
+            .collect()
+    }
+
+    fn click_on_tag(
+        &mut self,
+        _: &mut Connection<State>,
+        output: &Output,
+        _: WlSeat,
+        tag_id: Option<u32>,
+        btn: PointerBtn,
+    ) {
+        match btn {
+            PointerBtn::Left => {
+                if let Some(tag_id) = tag_id {
+                    self.set_workspace(tag_id);
+                }
+            }
+            PointerBtn::WheelUp | PointerBtn::WheelDown => {
+                if let Some(active_i) = self
+                    .workspaces
+                    .iter()
+                    .position(|ws| ws.output == output.name && ws.is_focused)
+                {
+                    if btn == PointerBtn::WheelUp {
+                        if let Some(prev) = self.workspaces[..active_i]
+                            .iter()
+                            .rfind(|ws| ws.output == output.name)
+                        {
+                            self.set_workspace(prev.idx);
+                        }
+                    } else {
+                        if let Some(next) = self.workspaces[active_i..]
+                            .iter()
+                            .skip(1)
+                            .find(|ws| ws.output == output.name)
+                        {
+                            self.set_workspace(next.idx);
+                        }
+                    }
+                }
+            }
+            _ => (),
+        }
+    }
+
+    fn as_any(&mut self) -> &mut dyn Any {
+        self
+    }
+}
+
+fn niri_cb(conn: &mut Connection<State>, state: &mut State) -> io::Result<()> {
+    let niri = state.shared_state.get_niri().unwrap();
+    let mut updated = false;
+    loop {
+        match niri.ipc.next_event() {
+            Ok(IpcEvent::WorkspacesChanged { workspaces }) => {
+                niri.workspaces = workspaces;
+                niri.workspaces.sort_by_key(|w| w.idx);
+                updated = true;
+            }
+            Ok(IpcEvent::WorkspaceActivated { id, focused }) => {
+                if let Some(new_active) = niri.workspaces.iter().position(|ws| ws.id == id) {
+                    // Clear the previous active workspace and apply it to the new one.
+                    if let Some(previous_active) = niri.workspaces.iter().position(|ws| {
+                        ws.is_active && ws.output == niri.workspaces[new_active].output
+                    }) {
+                        niri.workspaces[previous_active].is_active = false;
+                        niri.workspaces[new_active].is_active = true;
+                        updated = true;
+                    }
+                    if focused {
+                        if let Some(previous_focused) =
+                            niri.workspaces.iter().position(|ws| ws.is_focused)
+                        {
+                            niri.workspaces[previous_focused].is_focused = false;
+                            niri.workspaces[new_active].is_focused = true;
+                            updated = true;
+                        }
+                    }
+                }
+            }
+            Ok(IpcEvent::Ok(_)) => continue,
+            Ok(IpcEvent::Ignored(_)) => continue,
+            Err(e) if e.kind() == io::ErrorKind::WouldBlock => break,
+            Err(e) => return Err(e),
+        }
+    }
+    if updated {
+        state.tags_updated(conn, None);
+    }
+    Ok(())
+}
+
+#[derive(Debug)]
+struct Ipc {
+    sock_path: PathBuf,
+    sock: UnixStream,
+    sock_buf: Vec<u8>,
+}
+
+impl Ipc {
+    fn new(ns: &str) -> Option<Self> {
+        let sock_path = PathBuf::from(ns);
+        let mut sock = UnixStream::connect(sock_path.clone()).ok()?;
+        sock.set_nonblocking(true).ok()?;
+        sock.write_all("\"EventStream\"\n".as_bytes()).ok()?;
+        Some(Self {
+            sock_path,
+            sock,
+            sock_buf: Vec::new(),
+        })
+    }
+
+    fn exec(&self, cmd: &str) -> io::Result<()> {
+        let mut sock = UnixStream::connect(&self.sock_path)?;
+        sock.write_all(cmd.as_bytes())?;
+        sock.flush()?;
+        Ok(())
+    }
+
+    fn next_event(&mut self) -> io::Result<IpcEvent> {
+        loop {
+            if let Some(i) = memchr::memchr(b'\n', &self.sock_buf) {
+                let event = String::from_utf8_lossy(&self.sock_buf[..i]).into_owned();
+                self.sock_buf.drain(..=i);
+                return Ok(serde_json::from_str(&event)?);
+            }
+            if read_to_vec(&self.sock, &mut self.sock_buf)? == 0 {
+                return Err(io::Error::new(
+                    io::ErrorKind::BrokenPipe,
+                    "niri socked disconnected",
+                ));
+            }
+        }
+    }
+}
+
+#[derive(Debug, serde::Deserialize)]
+struct IpcWorkspace {
+    id: u32,  // Niri's internal id is monotonic, only used for comparison.
+    idx: u32, // idx is the user-facing workspace number.
+    name: Option<String>,
+    output: String,
+    is_focused: bool,
+    is_active: bool, // Niri's is_active means the workspace is visible on a display.
+}
+
+#[derive(Debug, serde::Deserialize)]
+enum IpcEvent {
+    Ok(IgnoredAny),
+    WorkspacesChanged {
+        workspaces: Vec<IpcWorkspace>,
+    },
+    WorkspaceActivated {
+        id: u32,
+        focused: bool,
+    },
+    #[serde(untagged)]
+    Ignored(IgnoredAny),
+}