# HG changeset patch
# User alfadur
# Date 1711497464 -10800
# Node ID e915ed28726e9e4a373f00de4c17b1d4bebb77d6
# Parent  8ba2b5007c298eec299c4320c8b61a25a3e1df94
thread server reference through dispatch to ensure rooms can be auto dropped

diff -r 8ba2b5007c29 -r e915ed28726e rust/hedgewars-server/src/core/server.rs
--- a/rust/hedgewars-server/src/core/server.rs	Wed Mar 27 02:19:44 2024 +0300
+++ b/rust/hedgewars-server/src/core/server.rs	Wed Mar 27 02:57:44 2024 +0300
@@ -270,7 +270,7 @@
     }
 
     #[inline]
-    pub fn get_room_control(&mut self, client_id: ClientId) -> Option<HwRoomControl> {
+    pub fn get_room_control(&mut self, client_id: ClientId) -> HwRoomOrServer {
         HwRoomControl::new(self, client_id)
     }
 
@@ -345,7 +345,15 @@
         client_id: ClientId,
         room_id: RoomId,
         room_password: Option<&str>,
-    ) -> Result<(&HwClient, Option<&HwClient>, &HwRoom, impl Iterator<Item = &HwClient> + Clone), JoinRoomError> {
+    ) -> Result<
+        (
+            &HwClient,
+            Option<&HwClient>,
+            &HwRoom,
+            impl Iterator<Item = &HwClient> + Clone,
+        ),
+        JoinRoomError,
+    > {
         use JoinRoomError::*;
         let room = &mut self.rooms[room_id];
         let client = &mut self.clients[client_id];
@@ -382,7 +390,15 @@
         client_id: ClientId,
         room_name: &str,
         room_password: Option<&str>,
-    ) -> Result<(&HwClient, Option<&HwClient>, &HwRoom, impl Iterator<Item = &HwClient> + Clone), JoinRoomError> {
+    ) -> Result<
+        (
+            &HwClient,
+            Option<&HwClient>,
+            &HwRoom,
+            impl Iterator<Item = &HwClient> + Clone,
+        ),
+        JoinRoomError,
+    > {
         use JoinRoomError::*;
         let room = self.rooms.iter().find(|(_, r)| r.name == room_name);
         if let Some((_, room)) = room {
@@ -548,6 +564,21 @@
     }
 }
 
+pub enum HwRoomOrServer<'a> {
+    Room(HwRoomControl<'a>),
+    Server(&'a mut HwServer),
+}
+
+impl<'a> HwRoomOrServer<'a> {
+    #[inline]
+    pub fn into_room(self) -> Option<HwRoomControl<'a>> {
+        match self {
+            HwRoomOrServer::Room(control) => Some(control),
+            HwRoomOrServer::Server(_) => None,
+        }
+    }
+}
+
 pub struct HwRoomControl<'a> {
     server: &'a mut HwServer,
     client_id: ClientId,
@@ -557,23 +588,16 @@
 
 impl<'a> HwRoomControl<'a> {
     #[inline]
-    pub fn new(server: &'a mut HwServer, client_id: ClientId) -> Option<Self> {
+    pub fn new(server: &'a mut HwServer, client_id: ClientId) -> HwRoomOrServer {
         if let Some(room_id) = server.clients[client_id].room_id {
-            Some(Self {
+            HwRoomOrServer::Room(Self {
                 server,
                 client_id,
                 room_id,
                 is_room_removed: false,
             })
         } else {
-            None
-        }
-    }
-
-    #[inline]
-    pub fn cleanup_room(self) {
-        if self.is_room_removed {
-            self.server.rooms.remove(self.room_id);
+            HwRoomOrServer::Server(server)
         }
     }
 
@@ -615,11 +639,6 @@
         )
     }
 
-    pub fn change_client<'b: 'a>(self, client_id: ClientId) -> Option<HwRoomControl<'a>> {
-        let room_id = self.room_id;
-        HwRoomControl::new(self.server, client_id).filter(|c| c.room_id == room_id)
-    }
-
     fn remove_from_room(&mut self, client_id: ClientId) -> LeaveRoomResult {
         let (client, room) = self
             .server
@@ -1132,6 +1151,15 @@
     }
 }
 
+impl<'a> Drop for HwRoomControl<'a> {
+    #[inline]
+    fn drop(&mut self) {
+        if self.is_room_removed {
+            self.server.rooms.remove(self.room_id);
+        }
+    }
+}
+
 fn allocate_room(rooms: &mut Slab<HwRoom>) -> &mut HwRoom {
     let entry = rooms.vacant_entry();
     let room = HwRoom::new(entry.key());
diff -r 8ba2b5007c29 -r e915ed28726e rust/hedgewars-server/src/handlers.rs
--- a/rust/hedgewars-server/src/handlers.rs	Wed Mar 27 02:19:44 2024 +0300
+++ b/rust/hedgewars-server/src/handlers.rs	Wed Mar 27 02:57:44 2024 +0300
@@ -15,7 +15,7 @@
     core::{
         anteroom::HwAnteroom,
         room::RoomSave,
-        server::HwServer,
+        server::{HwRoomOrServer, HwServer},
         types::{ClientId, Replay, RoomId},
     },
     utils,
@@ -372,8 +372,10 @@
                         }
                     }
                     _ => match state.server.get_room_control(client_id) {
-                        None => inlobby::handle(&mut state.server, client_id, response, message),
-                        Some(control) => inroom::handle(control, response, message),
+                        HwRoomOrServer::Room(control) => inroom::handle(control, response, message),
+                        HwRoomOrServer::Server(server) => {
+                            inlobby::handle(server, client_id, response, message)
+                        }
                     },
                 }
             }
diff -r 8ba2b5007c29 -r e915ed28726e rust/hedgewars-server/src/handlers/common.rs
--- a/rust/hedgewars-server/src/handlers/common.rs	Wed Mar 27 02:19:44 2024 +0300
+++ b/rust/hedgewars-server/src/handlers/common.rs	Wed Mar 27 02:57:44 2024 +0300
@@ -350,9 +350,11 @@
     let client = server.client(client_id);
     let nick = client.nick.clone();
 
-    if let Some(mut room_control) = server.get_room_control(client_id) {
-        let room_id = room_control.room().id;
-        let result = room_control.leave_room();
+    if let Some((room_id, result)) = server
+        .get_room_control(client_id)
+        .into_room()
+        .map(|mut control| (control.room().id, control.leave_room()))
+    {
         get_room_leave_result(server, server.room(room_id), &msg, result, response);
     }
 
diff -r 8ba2b5007c29 -r e915ed28726e rust/hedgewars-server/src/handlers/inlobby.rs
--- a/rust/hedgewars-server/src/handlers/inlobby.rs	Wed Mar 27 02:19:44 2024 +0300
+++ b/rust/hedgewars-server/src/handlers/inlobby.rs	Wed Mar 27 02:57:44 2024 +0300
@@ -77,7 +77,13 @@
                     match server.join_room(client_id, room_id, None) {
                         Err(error) => super::common::get_room_join_error(error, response),
                         Ok((client, master, room, room_clients)) => {
-                            super::common::get_room_join_data(client, master, room, room_clients, response)
+                            super::common::get_room_join_data(
+                                client,
+                                master,
+                                room,
+                                room_clients,
+                                response,
+                            )
                         }
                     }
                 } else {
diff -r 8ba2b5007c29 -r e915ed28726e rust/hedgewars-server/src/handlers/inroom.rs
--- a/rust/hedgewars-server/src/handlers/inroom.rs	Wed Mar 27 02:19:44 2024 +0300
+++ b/rust/hedgewars-server/src/handlers/inroom.rs	Wed Mar 27 02:57:44 2024 +0300
@@ -124,7 +124,6 @@
                 result,
                 response,
             );
-            room_control.cleanup_room();
         }
         Chat(msg) => {
             response.add(