WayVR: Clean-up destroyed toplevel windows

src/backend/wayvr/comp.rs

@@ -157,6 +157,13 @@ impl XdgShellHandler for Application {
         surface.send_configure();
     }
 
+    fn toplevel_destroyed(&mut self, surface: ToplevelSurface) {
+        if let Some(client) = surface.wl_surface().client() {
+            self.wayvr_tasks
+                .send(WayVRTask::DropToplevel(client.id(), surface.clone()));
+        }
+    }
+
     fn new_popup(&mut self, _surface: PopupSurface, _positioner: PositionerState) {
         // Handle popup creation here
     }

src/backend/wayvr/display.rs

@@ -183,6 +183,11 @@ impl Display {
         self.reposition_windows();
     }
 
+    pub fn remove_window(&mut self, window_handle: window::WindowHandle) {
+        self.displayed_windows
+            .retain(|disp| disp.window_handle != window_handle);
+    }
+
     pub fn reposition_windows(&mut self) {
         let window_count = self.displayed_windows.len();
 

src/backend/wayvr/mod.rs

@@ -73,6 +73,7 @@ pub struct ExternalProcessRequest {
 #[derive(Clone)]
 pub enum WayVRTask {
     NewToplevel(ClientId, ToplevelSurface),
+    DropToplevel(ClientId, ToplevelSurface),
     NewExternalProcess(ExternalProcessRequest),
     DropOverlay(super::overlay::OverlayID),
     ProcessTerminationRequest(process::ProcessHandle),
@@ -350,32 +351,59 @@ impl WayVR {
                             continue;
                         }
 
-                        if let Some(process_handle) =
+                        let Some(process_handle) =
                             process::find_by_pid(&self.state.processes, client.pid)
-                        {
+                        else {
+                            log::error!(
+                                    "WayVR window creation failed: Unexpected process ID {}. It wasn't registered before.",
+                                    client.pid
+                                );
+                            continue;
+                        };
+
                         let window_handle = self
                             .state
                             .wm
                             .borrow_mut()
                             .create_window(client.display_handle, &toplevel);
 
-                            if let Some(display) =
+                        let Some(display) = self.state.displays.get_mut(&client.display_handle)
-                                self.state.displays.get_mut(&client.display_handle)
+                        else {
-                            {
+                            // This shouldn't happen, scream if it does
+                            log::error!("Could not attach window handle into display");
+                            continue;
+                        };
+
                         display.add_window(window_handle, process_handle, &toplevel);
                         self.state.signals.send(WayVRSignal::BroadcastStateChanged(
                             packet_server::WvrStateChanged::WindowCreated,
                         ));
-                            } else {
-                                // This shouldn't happen, scream if it does
-                                log::error!("Could not attach window handle into display");
                     }
-                        } else {
-                            log::error!(
-                            "WayVR window creation failed: Unexpected process ID {}. It wasn't registered before.",
-                            client.pid
-                        );
                 }
+                WayVRTask::DropToplevel(client_id, toplevel) => {
+                    for client in &self.state.manager.clients {
+                        if client.client.id() != client_id {
+                            continue;
+                        }
+
+                        let mut wm = self.state.wm.borrow_mut();
+                        let Some(window_handle) = wm.find_window_handle(&toplevel) else {
+                            log::warn!("DropToplevel: Couldn't find matching window handle");
+                            continue;
+                        };
+
+                        let Some(display) = self.state.displays.get_mut(&client.display_handle)
+                        else {
+                            log::warn!("DropToplevel: Couldn't find matching display");
+                            continue;
+                        };
+
+                        display.remove_window(window_handle);
+                        wm.remove_window(window_handle);
+
+                        drop(wm);
+
+                        display.reposition_windows();
                     }
                 }
                 WayVRTask::ProcessTerminationRequest(process_handle) => {

src/backend/wayvr/window.rs

@@ -77,6 +77,10 @@ impl WindowManager {
     ) -> WindowHandle {
         self.windows.add(Window::new(display_handle, toplevel))
     }
+
+    pub fn remove_window(&mut self, window_handle: WindowHandle) {
+        self.windows.remove(&window_handle);
+    }
 }
 
 gen_id!(WindowVec, Window, WindowCell, WindowHandle);