Properly abstract EntityPools

Shatur · Shatur · commit 34bb99555ce8 · 2025-11-05T11:44:01.000+02:00
Since it's used in many places, it's better to abstract it to avoid
forgetting to call `clear`.
diff --git a/src/server.rs b/src/server.rs
@@ -1,3 +1,4 @@
+pub mod client_pools;
 pub mod message;
 pub mod related_entities;
 pub(super) mod removal_buffer;
@@ -29,23 +30,20 @@ use crate::{
     postcard_utils,
     prelude::*,
     server::{
-        replication_messages::mutations::{EntityMutations, MutationsSplit},
-        visibility::registry::FilterRegistry,
+        replication_messages::mutations::MutationsSplit, visibility::registry::FilterRegistry,
     },
     shared::{
         backend::channels::ClientChannel,
         message::server_message::message_buffer::MessageBuffer,
         replication::{
             client_ticks::{ClientTicks, EntityTicks},
-            registry::{
-                FnsId, ReplicationRegistry, component_mask::ComponentMask, ctx::SerializeCtx,
-                serde_fns::SerdeFns,
-            },
+            registry::{FnsId, ReplicationRegistry, ctx::SerializeCtx, serde_fns::SerdeFns},
             rules::{ReplicationRules, component::ComponentRule},
             track_mutate_messages::TrackMutateMessages,
         },
     },
 };
+use client_pools::ClientPools;
 use related_entities::RelatedEntities;
 use removal_buffer::{RemovalBuffer, RemovalReader};
 use replication_messages::{
@@ -242,7 +240,7 @@ fn cleanup_acks(
         let min_timestamp = time.elapsed().saturating_sub(mutations_timeout);
         for mut ticks in &mut clients {
             ticks.cleanup_older_mutations(min_timestamp, |mutate_info| {
-                pools.push_entities(mem::take(&mut mutate_info.entities));
+                pools.recycle_entities(mem::take(&mut mutate_info.entities));
             });
         }
     }
@@ -263,7 +261,7 @@ fn receive_acks(
                         )
                     });
                     if let Some(entities) = ticks.ack_mutate_message(client, mutate_index) {
-                        pools.push_entities(entities);
+                        pools.recycle_entities(entities);
                     }
                 }
                 Err(e) => {
@@ -532,7 +530,7 @@ fn collect_despawns(
                 // spawn and despawn could happen during the same tick.
                 trace!("writing despawn for `{entity}` for client `{client_entity}`");
                 message.add_despawn(entity_range.clone());
-                pools.push_components(entity_ticks.components);
+                pools.recycle_components(entity_ticks.components);
             }
             visibility.remove_despawned(entity);
             priority.remove(&entity);
@@ -545,7 +543,7 @@ fn collect_despawns(
                 trace!("writing visibility lost for `{entity}` for client `{client_entity}`");
                 let entity_range = serialized.write_entity(entity)?;
                 message.add_despawn(entity_range);
-                pools.push_components(entity_ticks.components);
+                pools.recycle_components(entity_ticks.components);
             }
             priority.remove(&entity);
         }
@@ -753,7 +751,7 @@ fn collect_changes(
                             entity_ticks.system_tick = change_tick.this_run();
                             entity_ticks.server_tick = server_tick;
                             entity_ticks.components |= &components;
-                            pools.push_components(components);
+                            pools.recycle_components(components);
                         }
                         Entry::Vacant(entry) => {
                             entry.insert(EntityTicks {
@@ -960,41 +958,3 @@ pub struct AuthorizedClient;
 /// See its documentation for more details.
 #[derive(Component, Deref, DerefMut, Debug, Default, Clone)]
 pub struct PriorityMap(EntityHashMap<f32>);
-
-/// Pools for various client components to reuse allocated capacity.
-///
-/// All data is cleared before the insertion.
-#[derive(Resource, Default)]
-struct ClientPools {
-    /// Entities with bitvecs for components from
-    /// [`MutateInfo`](crate::shared::replication::client_ticks::MutateInfo).
-    entities: Vec<Vec<(Entity, ComponentMask)>>,
-    /// Bitvecs for components from [`Updates`], [`Mutations`] and
-    /// [`MutateInfo`](crate::shared::replication::client_ticks::MutateInfo).
-    ///
-    /// Only heap-allocated instances are stored.
-    components: Vec<ComponentMask>,
-    /// Ranges from [`Updates`] and [`Mutations`].
-    ranges: Vec<Vec<Range<usize>>>,
-    /// Entities from [`Mutations`].
-    mutations: Vec<Vec<EntityMutations>>,
-}
-
-impl ClientPools {
-    fn push_entities(&mut self, mut entities: Vec<(Entity, ComponentMask)>) {
-        for (_, mut components) in entities.drain(..) {
-            if components.is_heap() {
-                components.clear();
-                self.components.push(components);
-            }
-        }
-        self.entities.push(entities);
-    }
-
-    fn push_components(&mut self, mut components: ComponentMask) {
-        if components.is_heap() {
-            components.clear();
-            self.components.push(components);
-        }
-    }
-}
diff --git a/src/server/client_pools.rs b/src/server/client_pools.rs
@@ -0,0 +1,77 @@
+use core::ops::Range;
+
+use bevy::prelude::*;
+
+use super::replication_messages::mutations::EntityMutations;
+use crate::shared::replication::registry::component_mask::ComponentMask;
+
+/// Pools for various client components to reuse allocated capacity.
+///
+/// All data is cleared before the insertion.
+#[derive(Resource, Default)]
+pub(super) struct ClientPools {
+    /// Entities with bitvecs for components from
+    /// [`MutateInfo`](crate::shared::replication::client_ticks::MutateInfo).
+    entities: Vec<Vec<(Entity, ComponentMask)>>,
+    /// Bitvecs for components from [`Updates`], [`Mutations`] and
+    /// [`MutateInfo`](crate::shared::replication::client_ticks::MutateInfo).
+    ///
+    /// Only heap-allocated instances are stored.
+    components: Vec<ComponentMask>,
+    /// Ranges from [`Updates`] and [`Mutations`].
+    ranges: Vec<Vec<Range<usize>>>,
+    /// Entities from [`Mutations`].
+    mutations: Vec<Vec<EntityMutations>>,
+}
+
+impl ClientPools {
+    pub(super) fn recycle_entities(&mut self, mut entities: Vec<(Entity, ComponentMask)>) {
+        for (_, mut components) in entities.drain(..) {
+            if components.is_heap() {
+                components.clear();
+                self.components.push(components);
+            }
+        }
+        self.entities.push(entities);
+    }
+
+    pub(super) fn recycle_components(&mut self, mut components: ComponentMask) {
+        if components.is_heap() {
+            components.clear();
+            self.components.push(components);
+        }
+    }
+
+    pub(super) fn recycle_ranges(&mut self, ranges: impl Iterator<Item = Vec<Range<usize>>>) {
+        self.ranges.extend(ranges.map(|mut ranges| {
+            ranges.clear();
+            ranges
+        }));
+    }
+
+    pub(super) fn recycle_mutations(
+        &mut self,
+        mutations: impl Iterator<Item = Vec<EntityMutations>>,
+    ) {
+        self.mutations.extend(mutations.map(|mut mutations| {
+            mutations.clear();
+            mutations
+        }));
+    }
+
+    pub(super) fn take_entities(&mut self) -> Vec<(Entity, ComponentMask)> {
+        self.entities.pop().unwrap_or_default()
+    }
+
+    pub(super) fn take_components(&mut self) -> ComponentMask {
+        self.components.pop().unwrap_or_default()
+    }
+
+    pub(super) fn take_ranges(&mut self) -> Vec<Range<usize>> {
+        self.ranges.pop().unwrap_or_default()
+    }
+
+    pub(super) fn take_mutations(&mut self) -> Vec<EntityMutations> {
+        self.mutations.pop().unwrap_or_default()
+    }
+}
diff --git a/src/server/replication_messages/mutations.rs b/src/server/replication_messages/mutations.rs
@@ -65,15 +65,14 @@ impl Mutations {
         graph_index: Option<usize>,
         entity_range: Range<usize>,
     ) {
-        let components = pools.ranges.pop().unwrap_or_default();
         let mutations = EntityMutations {
             entity,
             ranges: ChangeRanges {
                 entity: entity_range,
                 components_len: 0,
-                components,
+                components: pools.take_ranges(),
             },
-            components: pools.components.pop().unwrap_or_default(),
+            components: pools.take_components(),
         };
 
         match graph_index {
@@ -155,7 +154,7 @@ impl Mutations {
             server_tick,
             system_tick,
             timestamp,
-            entities: pools.entities.pop().unwrap_or_default(),
+            entities: pools.take_entities(),
         };
         let mut mutate_index = ticks.next_mutate_index();
         let mut chunks = EntityChunks::new(&mut self.related, &mut self.standalone);
@@ -185,7 +184,7 @@ impl Mutations {
                     server_tick,
                     system_tick,
                     timestamp,
-                    entities: pools.entities.pop().unwrap_or_default(),
+                    entities: pools.take_entities(),
                 };
                 chunks_range.start = chunks_range.end;
                 header_size = metadata_size + serialized_size(&mutate_index)?; // Recalculate since the mutate index changed.
@@ -263,11 +262,7 @@ impl Mutations {
             .iter_mut()
             .chain(iter::once(&mut self.standalone))
         {
-            let ranges = entities.drain(..).map(|mut mutations| {
-                mutations.ranges.components.clear();
-                mutations.ranges.components
-            });
-            pools.ranges.extend(ranges);
+            pools.recycle_ranges(entities.drain(..).map(|m| m.ranges.components));
             // We don't take component masks because they are moved to `MutateInfo` during sending.
         }
     }
@@ -279,14 +274,8 @@ impl Mutations {
         match self.related.len().cmp(&graphs_count) {
             Ordering::Less => self
                 .related
-                .resize_with(graphs_count, || pools.mutations.pop().unwrap_or_default()),
-            Ordering::Greater => {
-                let mutations = self.related.drain(graphs_count..).map(|mut mutations| {
-                    mutations.clear();
-                    mutations
-                });
-                pools.mutations.extend(mutations);
-            }
+                .resize_with(graphs_count, || pools.take_mutations()),
+            Ordering::Greater => pools.recycle_mutations(self.related.drain(graphs_count..)),
             Ordering::Equal => (),
         }
     }
diff --git a/src/server/replication_messages/updates.rs b/src/server/replication_messages/updates.rs
@@ -1,4 +1,4 @@
-use core::{mem, ops::Range};
+use core::{iter, mem, ops::Range};
 
 use bevy::prelude::*;
 use postcard::experimental::serialized_size;
@@ -120,7 +120,7 @@ impl Updates {
         self.removals.push(RemovalRanges {
             entity,
             ids_len: 0,
-            ids: pools.ranges.pop().unwrap_or_default(),
+            ids: pools.take_ranges(),
         });
         self.removals_entity_added = true;
     }
@@ -159,7 +159,7 @@ impl Updates {
         self.changes.push(ChangeRanges {
             entity,
             components_len: 0,
-            components: pools.ranges.pop().unwrap_or_default(),
+            components: pools.take_ranges(),
         });
         self.changed_entity_added = true;
     }
@@ -183,19 +183,18 @@ impl Updates {
     /// Takes last mutated entity with its component chunks from the mutate message.
     pub(crate) fn take_added_entity(&mut self, pools: &mut ClientPools, mutations: &mut Mutations) {
         debug_assert!(mutations.entity_added());
-        let mut entity_mutations = mutations.pop().expect("entity should be written");
+        let entity_mutations = mutations.pop().expect("entity should be written");
 
         if !self.changed_entity_added {
             self.changes.push(entity_mutations.ranges);
         } else {
             let changes = self.changes.last_mut().expect("entity should be written");
             debug_assert_eq!(entity_mutations.ranges.entity, changes.entity);
             changes.extend(&entity_mutations.ranges);
-            entity_mutations.ranges.components.clear();
-            pools.ranges.push(entity_mutations.ranges.components);
+            pools.recycle_ranges(iter::once(entity_mutations.ranges.components));
 
             self.changed_components |= &entity_mutations.components;
-            pools.push_components(entity_mutations.components);
+            pools.recycle_components(entity_mutations.components);
         }
     }
 
@@ -358,15 +357,8 @@ impl Updates {
         self.despawns.clear();
         self.despawns_len = 0;
 
-        let change_ranges = self.changes.drain(..).map(|mut changes| {
-            changes.components.clear();
-            changes.components
-        });
-        let removal_ranges = self.removals.drain(..).map(|mut removals| {
-            removals.ids.clear();
-            removals.ids
-        });
-        pools.ranges.extend(change_ranges.chain(removal_ranges));
+        pools.recycle_ranges(self.changes.drain(..).map(|c| c.components));
+        pools.recycle_ranges(self.removals.drain(..).map(|c| c.ids));
     }
 }
 
