From 0f952885e58862c04a9d61a46d045bdcb97a9780 Mon Sep 17 00:00:00 2001
From: jeff425 <goodhue.jeff@gmail.com>
Date: Tue, 9 May 2023 21:30:33 -0700
Subject: [PATCH] Added broadcast_message_scoped to allow sending messages only
 to users with a specific entity in their scope

---
 adapters/bevy/server/src/server.rs   |  4 ++++
 server/src/server.rs                 | 27 ++++++++++++++++++++-------
 server/src/world/entity_scope_map.rs |  9 +++++++--
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/adapters/bevy/server/src/server.rs b/adapters/bevy/server/src/server.rs
index 7f7dee2d6..8bd8cd405 100644
--- a/adapters/bevy/server/src/server.rs
+++ b/adapters/bevy/server/src/server.rs
@@ -57,6 +57,10 @@ impl<'w> Server<'w> {
         self.server.broadcast_message::<C, M>(message);
     }
 
+    pub fn broadcast_message_scoped<C: Channel, M: Message>(&mut self, entity: &Entity, message: &M) {
+        self.server.broadcast_message_scoped::<C, M>(entity, message);
+    }
+
     pub fn receive_tick_buffer_messages(&mut self, tick: &Tick) -> TickBufferMessages {
         self.server.receive_tick_buffer_messages(tick)
     }
diff --git a/server/src/server.rs b/server/src/server.rs
index 6060f9a51..9e6631c19 100644
--- a/server/src/server.rs
+++ b/server/src/server.rs
@@ -286,6 +286,25 @@ impl<E: Copy + Eq + Hash + Send + Sync> Server<E> {
         })
     }
 
+    /// Sends a message to all connected users that contain the given entity in their scope using a given channel
+    pub fn broadcast_message_scoped<C: Channel, M: Message>(&mut self, entity: &E, message: &M) {
+        let cloned_message = M::clone_box(message);
+        self.broadcast_message_scoped_inner(entity, &ChannelKind::of::<C>(), cloned_message);
+    }
+
+    fn broadcast_message_scoped_inner(
+        &mut self,
+        entity: &E,
+        channel_kind: &ChannelKind,
+        message_box: Box<dyn Message>,
+    ) {
+        self.user_keys().iter().for_each(|user_key| {
+            if self.entity_scope_map.get(user_key, entity) {
+                self.send_message_inner(user_key, channel_kind, message_box.clone())
+            }
+        })
+    }
+
     pub fn receive_tick_buffer_messages(&mut self, tick: &Tick) -> TickBufferMessages {
         let mut tick_buffer_messages = TickBufferMessages::new();
         for (_user_address, connection) in self.user_connections.iter_mut() {
@@ -1281,13 +1300,7 @@ impl<E: Copy + Eq + Hash + Send + Sync> Server<E> {
                                 let currently_in_scope =
                                     connection.base.host_world_manager.host_has_entity(entity);
 
-                                let should_be_in_scope = if let Some(in_scope) =
-                                    self.entity_scope_map.get(user_key, entity)
-                                {
-                                    *in_scope
-                                } else {
-                                    false
-                                };
+                                let should_be_in_scope = self.entity_scope_map.get(user_key, entity);
 
                                 if should_be_in_scope {
                                     if !currently_in_scope {
diff --git a/server/src/world/entity_scope_map.rs b/server/src/world/entity_scope_map.rs
index e737cf267..f6664e2ed 100644
--- a/server/src/world/entity_scope_map.rs
+++ b/server/src/world/entity_scope_map.rs
@@ -20,13 +20,18 @@ impl<E: Copy + Eq + Hash> EntityScopeMap<E> {
         }
     }
 
-    pub fn get(&self, user_key: &UserKey, entity: &E) -> Option<&bool> {
+    pub fn get(&self, user_key: &UserKey, entity: &E) -> bool {
         let key = (*user_key, *entity);
 
-        self.main_map.get(&key)
+        self.main_map.get(&key).map(|k| *k).unwrap_or(false)
     }
 
     pub fn insert(&mut self, user_key: UserKey, entity: E, in_scope: bool) {
+        if self.main_map.contains_key(&(user_key, entity)) {
+            // Only need to update scope status
+            self.main_map.insert((user_key, entity), in_scope);
+            return;
+        }
         self.entities_of_user
             .entry(user_key)
             .or_insert_with(|| HashSet::new());