From 05b3baa5d99057076429b17b86dc41fc872da926 Mon Sep 17 00:00:00 2001
From: Aleksander <aleksander@oo8.dev>
Date: Sun, 9 Mar 2025 22:04:46 +0100
Subject: [PATCH] WayVR: Clean-up destroyed toplevel windows

---
 src/backend/wayvr/comp.rs    |  7 ++++
 src/backend/wayvr/display.rs |  5 +++
 src/backend/wayvr/mod.rs     | 74 +++++++++++++++++++++++++-----------
 src/backend/wayvr/window.rs  |  4 ++
 4 files changed, 67 insertions(+), 23 deletions(-)

diff --git a/src/backend/wayvr/comp.rs b/src/backend/wayvr/comp.rs
index 2af7ddc..9c4dc88 100644
--- a/src/backend/wayvr/comp.rs
+++ b/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
     }
diff --git a/src/backend/wayvr/display.rs b/src/backend/wayvr/display.rs
index cd51d06..01f1445 100644
--- a/src/backend/wayvr/display.rs
+++ b/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();
 
diff --git a/src/backend/wayvr/mod.rs b/src/backend/wayvr/mod.rs
index 55c29c1..cd0ab47 100644
--- a/src/backend/wayvr/mod.rs
+++ b/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)
-                        {
-                            let window_handle = self
-                                .state
-                                .wm
-                                .borrow_mut()
-                                .create_window(client.display_handle, &toplevel);
-
-                            if let Some(display) =
-                                self.state.displays.get_mut(&client.display_handle)
-                            {
-                                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 {
+                        else {
                             log::error!(
-                            "WayVR window creation failed: Unexpected process ID {}. It wasn't registered before.",
-                            client.pid
-                        );
+                                    "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);
+
+                        let Some(display) = 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,
+                        ));
+                    }
+                }
+                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) => {
diff --git a/src/backend/wayvr/window.rs b/src/backend/wayvr/window.rs
index 6466303..893b7d4 100644
--- a/src/backend/wayvr/window.rs
+++ b/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);