From bcb60ea8394f7de78e71fb6b964ba329d07479e1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gerg=C5=91=20S=C3=A1lyi?= <salyigergo94@gmail.com>
Date: Wed, 9 Apr 2025 19:23:29 +0200
Subject: [PATCH] Change niri event subscription loop to cache the state of
 current workspaces

---
 src/compositors/niri.rs | 106 +++++++++++++++++++++-------------------
 1 file changed, 57 insertions(+), 49 deletions(-)

diff --git a/src/compositors/niri.rs b/src/compositors/niri.rs
index 0be3939..8435c58 100644
--- a/src/compositors/niri.rs
+++ b/src/compositors/niri.rs
@@ -1,6 +1,6 @@
-
+use std::io;
 use super::{CompositorInterface, WorkspaceVisible, EventSender};
-use niri_ipc::{socket::Socket, Event, Request, Response};
+use niri_ipc::{socket::Socket, Event, Request, Response, Workspace};
 
 pub struct NiriConnectionTask {}
 
@@ -8,60 +8,68 @@ impl NiriConnectionTask {
     pub fn new() -> Self {
         NiriConnectionTask {}
     }
-
-    fn query_workspace(&self, id: u64) -> (String, String) {
-        if let Ok((Ok(Response::Workspaces(workspaces)), _)) = Socket::connect()
-            .expect("failed to connect to niri socket")
-            .send(Request::Workspaces)
-        {
-            if let Some(workspace) = workspaces.into_iter().find(|w| w.id == id) {
-                return (
-                    workspace.name.unwrap_or_else(|| format!("{}", workspace.idx)),
-                    workspace.output.unwrap_or_else(String::new),
-                );
-            }
-            panic!("unknown workspace id");
-        } else {
-            panic!("niri workspace query failed");
-        }
-    }
 }
+
 impl CompositorInterface for NiriConnectionTask {
     fn request_visible_workspaces(&mut self) -> Vec<WorkspaceVisible> {
-        if let Ok((Ok(Response::Workspaces(workspaces)), _)) = Socket::connect()
-            .expect("failed to connect to niri socket")
-            .send(Request::Workspaces)
-        {
-            workspaces
-                .into_iter()
-                .filter(|w| w.is_active)
-                .map(|workspace| WorkspaceVisible {
-                    output: workspace.output.unwrap_or_default(),
-                    workspace_name: workspace.name
-                        .unwrap_or_else(|| format!("{}", workspace.idx)),
-                })
-                .collect()
-        } else {
-            panic!("unable to retrieve niri workspaces")
-        }
+        request_workspaces().into_iter()
+            .filter(|w| w.is_active)
+            .map(|workspace| WorkspaceVisible {
+                output: workspace.output.unwrap_or_default(),
+                workspace_name: workspace.name
+                    .unwrap_or_else(|| format!("{}", workspace.idx)),
+            })
+            .collect()
     }
 
     fn subscribe_event_loop(self, event_sender: EventSender) {
-        if let Ok((Ok(Response::Handled), mut callback)) = Socket::connect()
-            .expect("failed to connect to niri socket")
-            .send(Request::EventStream)
-        {
-            while let Ok(event) = callback() {
-                if let Event::WorkspaceActivated { id, focused: _ } = event {
-                    let (workspace_name, output) = self.query_workspace(id);
-                    event_sender.send(WorkspaceVisible {
-                        output,
-                        workspace_name,
-                    });
-                }
+        let mut workspaces_state = request_workspaces();
+        let mut callback = request_event_stream();
+        while let Ok(event) = callback() {
+            match event {
+                Event::WorkspaceActivated { id, focused: _ } => {
+                    let visible_workspace =
+                        find_workspace(&workspaces_state, id);
+                    event_sender.send(visible_workspace);
+                },
+                Event::WorkspacesChanged { workspaces } => {
+                    workspaces_state = workspaces
+                },
+                _ => {},
             }
-        } else {
-            panic!("failed to subscribe to event stream");
         }
     }
 }
+
+fn find_workspace(workspaces: &[Workspace], id: u64) -> WorkspaceVisible {
+    let workspace = workspaces.iter()
+        .find(|workspace| workspace.id == id)
+        .expect("Unknown niri workspace id {id}");
+    let workspace_name = workspace.name.clone()
+        .unwrap_or_else(|| format!("{}", workspace.idx));
+    let output = workspace.output.clone().unwrap_or_default();
+    WorkspaceVisible { output, workspace_name }
+}
+
+fn request_event_stream() -> impl FnMut() -> Result<Event, io::Error> {
+    let Ok((Ok(Response::Handled), callback)) = Socket::connect()
+        .expect("failed to connect to niri socket")
+        .send(Request::EventStream)
+    else {
+        panic!("failed to subscribe to event stream");
+    };
+    callback
+}
+
+fn request_workspaces() -> Vec<Workspace> {
+    let response = Socket::connect()
+        .expect("failed to connect to niri socket")
+        .send(Request::Workspaces)
+        .expect("failed to send niri ipc request")
+        .0
+        .expect("niri workspace query failed");
+    let Response::Workspaces(workspaces) = response else {
+        panic!("unexpected response from niri");
+    };
+    workspaces
+}