fix: take care of removed outputs

Author: typester

src/gql.rs

@@ -18,6 +18,7 @@ pub enum RiverEventType {
     OutputUrgentTags,
     OutputLayoutName,
     OutputLayoutNameClear,
+    OutputRemoved,
     SeatFocusedOutput,
     SeatUnfocusedOutput,
     SeatFocusedView,
@@ -33,6 +34,7 @@ impl From<&river::Event> for RiverEventType {
             OutputUrgentTags { .. } => RiverEventType::OutputUrgentTags,
             OutputLayoutName { .. } => RiverEventType::OutputLayoutName,
             OutputLayoutNameClear { .. } => RiverEventType::OutputLayoutNameClear,
+            OutputRemoved { .. } => RiverEventType::OutputRemoved,
             SeatFocusedOutput { .. } => RiverEventType::SeatFocusedOutput,
             SeatUnfocusedOutput { .. } => RiverEventType::SeatUnfocusedOutput,
             SeatFocusedView { .. } => RiverEventType::SeatFocusedView,
@@ -222,6 +224,25 @@ impl RiverSnapshot {
                     state.layout_name = None;
                 });
             }
+            OutputRemoved { id, name } => {
+                let gql_id = id_to_graphql(id);
+                let key = gql_id.to_string();
+                if let Some(state) = self.outputs.remove(&key) {
+                    if let Some(name_value) = state.name {
+                        self.output_names.remove(&name_value);
+                    }
+                } else if let Some(name_value) = name.as_ref() {
+                    self.output_names.remove(name_value);
+                }
+                let clear_focus = self
+                    .seat_focused_output
+                    .as_ref()
+                    .map(|focused| focused.output_id == gql_id)
+                    .unwrap_or(false);
+                if clear_focus {
+                    self.seat_focused_output = None;
+                }
+            }
             SeatFocusedOutput { id, name } => {
                 self.seat_focused_output = Some(NamedOutputId {
                     output_id: id_to_graphql(id),
@@ -444,6 +465,7 @@ fn event_types_for_name(name: &str) -> Vec<RiverEventType> {
             RiverEventType::OutputLayoutName,
             RiverEventType::OutputLayoutNameClear,
         ],
+        "OutputRemoved" => vec![RiverEventType::OutputRemoved],
         "SeatFocusedOutput" => vec![RiverEventType::SeatFocusedOutput],
         "SeatUnfocusedOutput" => vec![RiverEventType::SeatUnfocusedOutput],
         "SeatFocusedView" => vec![RiverEventType::SeatFocusedView],
@@ -473,6 +495,7 @@ fn event_output_name<'a>(event: &'a river::Event) -> Option<&'a str> {
         | OutputUrgentTags { name, .. }
         | OutputLayoutName { name, .. }
         | OutputLayoutNameClear { name, .. }
+        | OutputRemoved { name, .. }
         | SeatFocusedOutput { name, .. }
         | SeatUnfocusedOutput { name, .. } => name.as_deref(),
 
@@ -527,6 +550,7 @@ pub enum RiverEvent {
     OutputViewTags(GOutputViewTags),
     OutputUrgentTags(GOutputUrgentTags),
     OutputLayoutName(GOutputLayoutName),
+    OutputRemoved(GOutputRemoved),
     SeatFocusedOutput(GSeatFocusedOutput),
     SeatUnfocusedOutput(GSeatUnfocusedOutput),
     SeatFocusedView(GSeatFocusedView),
@@ -632,6 +656,22 @@ impl GOutputLayoutName {
     }
 }
 
+#[derive(Clone)]
+pub struct GOutputRemoved {
+    pub output_id: ID,
+    pub name: Option<String>,
+}
+#[Object(name = "OutputRemoved")]
+impl GOutputRemoved {
+    async fn output_id(&self) -> &ID {
+        &self.output_id
+    }
+
+    async fn name(&self) -> Option<&str> {
+        self.name.as_deref()
+    }
+}
+
 // no-op clear event omitted in minimal schema
 
 #[derive(Clone)]
@@ -746,6 +786,13 @@ fn make_river_event(value: river::Event, include_lists: bool) -> RiverEvent {
             output_name: name,
             layout: String::new(),
         }),
+        OutputRemoved {
+            id: output_id,
+            name,
+        } => RiverEvent::OutputRemoved(GOutputRemoved {
+            output_id: id_to_graphql(&output_id),
+            name,
+        }),
         SeatFocusedOutput {
             id: output_id,
             name,

src/river.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
 
@@ -52,6 +52,10 @@ pub enum Event {
         id: ObjectId,
         name: Option<String>,
     },
+    OutputRemoved {
+        id: ObjectId,
+        name: Option<String>,
+    },
 
     SeatFocusedOutput {
         id: ObjectId,
@@ -70,8 +74,8 @@ pub enum Event {
 }
 
 struct State {
-    outputs: Vec<WlOutput>,
-    seats: Vec<WlSeat>,
+    outputs: HashMap<u32, WlOutput>,
+    seats: HashMap<u32, WlSeat>,
     manager: Option<ZriverStatusManagerV1>,
     output_statuses: Vec<ZriverOutputStatusV1>,
     seat_statuses: Vec<ZriverSeatStatusV1>,
@@ -83,8 +87,8 @@ struct State {
 impl State {
     fn new(tx: UnboundedSender<Event>) -> Self {
         Self {
-            outputs: Vec::new(),
-            seats: Vec::new(),
+            outputs: HashMap::new(),
+            seats: HashMap::new(),
             manager: None,
             output_statuses: Vec::new(),
             seat_statuses: Vec::new(),
@@ -115,13 +119,13 @@ impl State {
 
     fn create_status_for_all(&mut self, qh: &QueueHandle<Self>) {
         if self.manager.is_some() {
-            let outs = self.outputs.clone();
-            for o in &outs {
-                self.maybe_create_status_for_output(qh, o);
+            let outputs: Vec<_> = self.outputs.values().cloned().collect();
+            for output in &outputs {
+                self.maybe_create_status_for_output(qh, output);
             }
-            let seats = self.seats.clone();
-            for s in &seats {
-                self.maybe_create_status_for_seat(qh, s);
+            let seats: Vec<_> = self.seats.values().cloned().collect();
+            for seat in &seats {
+                self.maybe_create_status_for_seat(qh, seat);
             }
         }
     }
@@ -188,15 +192,13 @@ impl Dispatch<WlRegistry, ()> for State {
             } => match interface.as_str() {
                 "wl_output" => {
                     let output = registry.bind::<WlOutput, _, _>(name, version.min(4), qh, ());
-                    state.outputs.push(output);
-                    let last = state.outputs.last().unwrap().clone();
-                    state.maybe_create_status_for_output(qh, &last);
+                    state.maybe_create_status_for_output(qh, &output);
+                    state.outputs.insert(name, output);
                 }
                 "wl_seat" => {
                     let seat = registry.bind::<WlSeat, _, _>(name, version.min(5), qh, ());
-                    state.seats.push(seat);
-                    let last = state.seats.last().unwrap().clone();
-                    state.maybe_create_status_for_seat(qh, &last);
+                    state.maybe_create_status_for_seat(qh, &seat);
+                    state.seats.insert(name, seat);
                 }
                 "zriver_status_manager_v1" => {
                     let mgr =
@@ -206,6 +208,11 @@ impl Dispatch<WlRegistry, ()> for State {
                 }
                 _ => {}
             },
+            wl_registry::Event::GlobalRemove { name } => {
+                if !state.remove_output(name) {
+                    state.seats.remove(&name);
+                }
+            }
             _ => {}
         }
     }
@@ -344,6 +351,31 @@ fn parse_u32_array(bytes: &[u8]) -> Vec<u32> {
     v
 }
 
+impl State {
+    fn remove_output(&mut self, global: u32) -> bool {
+        let Some(output) = self.outputs.remove(&global) else {
+            return false;
+        };
+        let id = output.id();
+        let label = self.output_label(&id);
+        let protocol_id = id.protocol_id();
+
+        let mut removed_status_ids = HashSet::new();
+        for (status_id, owner) in &self.output_status_owner {
+            if owner.protocol_id() == protocol_id {
+                removed_status_ids.insert(*status_id);
+            }
+        }
+        self.output_status_owner
+            .retain(|status_id, _| !removed_status_ids.contains(status_id));
+        self.output_statuses
+            .retain(|status| !removed_status_ids.contains(&status.id().protocol_id()));
+        self.output_info.remove(&protocol_id);
+        let _ = self.tx.send(Event::OutputRemoved { id, name: label });
+        true
+    }
+}
+
 pub struct RiverStatus;
 
 impl RiverStatus {