niri-tag/daemon/manager.rs
2025-06-21 23:22:10 +10:00

386 lines
15 KiB
Rust

use crate::socket::{create_niri_socket, query, tell};
use anyhow::{Context, Result, anyhow};
use niri_ipc::{
Action, Event, Reply, Request, Response, Window, Workspace, WorkspaceReferenceArg,
state::{EventStreamState, EventStreamStatePart},
};
use niri_tag::{TagCmd, TagEvent, TagState};
use smol::{
channel::{self, Sender},
future,
io::BufReader,
net::unix::UnixStream,
};
use std::collections::HashMap;
pub struct NiriTag {
tags: HashMap<u8, bool>,
windows: HashMap<u64, u8>,
state: EventStreamState,
socket: BufReader<UnixStream>,
ev_tx: channel::Sender<TagEvent>,
}
enum TagAction {
ChangeWindow(u64),
ChangeTag(u8),
}
impl NiriTag {
pub async fn new(ev_tx: channel::Sender<TagEvent>) -> Result<Self> {
Ok(Self {
tags: HashMap::new(),
windows: HashMap::new(),
state: EventStreamState::default(),
socket: create_niri_socket().await?,
ev_tx,
})
}
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::*;
let send_event = async |tx: Sender<TagEvent>, ev| {
smol::spawn(async move {
tx.send(ev)
.await
.inspect_err(|e| tracing::error!("Failed to send event: {}", e))
})
.detach();
};
let add_tag = async |tx: Sender<TagEvent>, windows: &HashMap<u64, u8>, t| {
if windows
.iter()
.filter(|(_, tag)| **tag == t)
.collect::<HashMap<_, _>>()
.is_empty()
{
send_event(tx, TagEvent::TagOccupied(t)).await;
}
};
let rm_tag = async |tx: Sender<TagEvent>, windows: &HashMap<u64, u8>, wid, old_tag| {
if old_tag != 0
&& windows
.iter()
.filter(|(w, tag)| **tag == old_tag && **w != wid)
.collect::<Vec<(_, _)>>()
.is_empty()
{
send_event(tx, TagEvent::TagEmpty(old_tag)).await;
}
};
// 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::FullState(tx) => {
tracing::debug!("received request for full state");
let fullstate: HashMap<u8, TagState> = self
.tags
.iter()
.map(|(&t, &enabled)| {
(
t,
TagState {
enabled,
occupied: self.windows.values().filter(|w_t| **w_t == t).count()
> 0,
urgent: false, // urgency is TODO
},
)
})
.collect();
return tx.send(fullstate).await.map_err(|e| anyhow!(e));
}
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);
let tx = self.ev_tx.clone();
add_tag(tx, &self.windows, t).await;
ChangeWindow(wid)
}
TagCmd::RemoveTagFromWin(_) => {
let win = self.get_focused_window().await?;
let wid = win.id;
let old_tag = self.windows.insert(wid, 0).unwrap_or(0);
tracing::debug!("resetting tag on {}", wid);
let tx = self.ev_tx.clone();
rm_tag(tx, &self.windows, wid, old_tag).await;
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 };
let tx = self.ev_tx.clone();
if toggle == 0 {
rm_tag(tx, &self.windows, wid, this_tag).await;
} else {
add_tag(tx, &self.windows, t).await;
}
tracing::debug!("toggling {} to tag {}", wid, toggle);
self.windows.insert(wid, toggle);
ChangeWindow(wid)
}
TagCmd::EnableTag(t) => {
self.tags.insert(t, true);
send_event(self.ev_tx.clone(), TagEvent::TagEnabled(t)).await;
ChangeTag(t)
}
TagCmd::DisableTag(t) => {
self.tags.insert(t, false);
send_event(self.ev_tx.clone(), TagEvent::TagDisabled(t)).await;
ChangeTag(t)
}
TagCmd::ToggleTag(t) => {
let visible = *self.tags.entry(t).or_insert(false);
if visible {
send_event(self.ev_tx.clone(), TagEvent::TagEnabled(t)).await;
} else {
send_event(self.ev_tx.clone(), TagEvent::TagDisabled(t)).await;
}
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 { windows } => {
for w in windows {
self.windows.entry(w.id).or_insert(0);
let action = self.do_action(TagAction::ChangeWindow(w.id)).await;
if let Err(e) = action {
tracing::warn!("Failed to ChangeWindow on {}: {}", w.id, e);
}
}
Ok(())
}
// WindowUrgencyChanged { .. } => (),
_ => Ok(()),
}
}
#[allow(unreachable_code)]
pub async fn manage_tags(
mut self,
ev_rx: channel::Receiver<Event>,
tag_rx: channel::Receiver<TagCmd>,
fullstate_rx: channel::Receiver<channel::Sender<HashMap<u8, TagState>>>,
) -> Result<()> {
// base tag is always visible
self.tags.insert(0, true);
loop {
let recvd: Receivable = future::or(
async { ev_rx.recv().await.map(Receivable::Event) },
future::or(
async { tag_rx.recv().await.map(Receivable::TagCmd) },
async { fullstate_rx.recv().await.map(Receivable::FullState) },
),
)
.await?;
tracing::debug!("manager received {:?}", recvd);
let res = self.handle_recvd(recvd).await;
match res {
Ok(()) => (),
Err(e) => tracing::error!("error occurred in manager loop: {}", e),
}
}
tracing::error!("Manager loop ended");
unreachable!("Manager loop ended");
}
}
#[derive(Debug)]
enum Receivable {
Event(Event),
TagCmd(TagCmd),
FullState(channel::Sender<HashMap<u8, TagState>>),
}