Skip to main content

naia_shared/world/entity/
global_entity_map.rs

1use std::{collections::HashMap, hash::Hash};
2
3use crate::{
4    BigMap, EntityAndGlobalEntityConverter, EntityDoesNotExistError, GlobalEntity, RemoteEntity,
5};
6
7/// Bidirectional map between world-local entities and their stable [`GlobalEntity`] identifiers.
8pub struct GlobalEntityMap<E: Copy + Eq + Hash + Send + Sync> {
9    entity_to_global_map: HashMap<E, GlobalEntity>,
10    global_to_entity_map: BigMap<GlobalEntity, Option<E>>,
11    reserved_global_entities: HashMap<RemoteEntity, GlobalEntity>,
12}
13
14impl<E: Copy + Eq + Hash + Send + Sync> Default for GlobalEntityMap<E> {
15    fn default() -> Self {
16        Self::new()
17    }
18}
19
20impl<E: Copy + Eq + Hash + Send + Sync> GlobalEntityMap<E> {
21    /// Creates an empty map with no entity registrations.
22    pub fn new() -> Self {
23        Self {
24            entity_to_global_map: HashMap::new(),
25            global_to_entity_map: BigMap::new(),
26            reserved_global_entities: HashMap::new(),
27        }
28    }
29
30    /// Returns the number of world entities currently registered in the map.
31    pub fn entity_count(&self) -> usize {
32        self.entity_to_global_map.len()
33    }
34
35    #[cfg(feature = "test_utils")]
36    #[doc(hidden)]
37    pub fn set_global_entity_counter_for_test(&mut self, value: u64) {
38        self.global_to_entity_map.set_current_index_for_test(value);
39    }
40}
41
42impl<E: Copy + Eq + Hash + Send + Sync> EntityAndGlobalEntityConverter<E> for GlobalEntityMap<E> {
43    fn global_entity_to_entity(
44        &self,
45        global_entity: &GlobalEntity,
46    ) -> Result<E, EntityDoesNotExistError> {
47        match self.global_to_entity_map.get(global_entity) {
48            Some(world_entity_opt) => {
49                if let Some(world_entity) = world_entity_opt {
50                    return Ok(*world_entity);
51                }
52                // warn!(
53                //     "Global entity {:?} exists but does not map to a world entity yet, it is reserved.",
54                //     global_entity
55                // );
56                Err(EntityDoesNotExistError)
57            }
58            None => Err(EntityDoesNotExistError),
59        }
60    }
61
62    fn entity_to_global_entity(
63        &self,
64        world_entity: &E,
65    ) -> Result<GlobalEntity, EntityDoesNotExistError> {
66        match self.entity_to_global_map.get(world_entity) {
67            Some(global_entity) => Ok(*global_entity),
68            None => Err(EntityDoesNotExistError),
69        }
70    }
71}
72
73/// Extends [`EntityAndGlobalEntityConverter`] with the ability to create, reserve, and destroy entity mappings.
74pub trait GlobalEntitySpawner<E: Copy + Eq + Hash + Send + Sync>:
75    EntityAndGlobalEntityConverter<E>
76{
77    /// Registers `world_entity` in the map, reusing a reserved [`GlobalEntity`] for `remote_entity_opt` if one exists.
78    fn spawn(&mut self, world_entity: E, remote_entity_opt: Option<RemoteEntity>) -> GlobalEntity;
79    /// Pre-allocates a [`GlobalEntity`] slot for `remote_entity` before the local world entity is spawned.
80    fn reserve_global_entity(&mut self, remote_entity: RemoteEntity) -> GlobalEntity;
81    /// Removes the mapping keyed by `global_entity`, panicking if it does not exist.
82    fn despawn_by_global(&mut self, global_entity: &GlobalEntity);
83    /// Removes the mapping keyed by `world_entity`, panicking if it does not exist.
84    fn despawn_by_world(&mut self, world_entity: &E);
85    /// Returns `self` as an [`EntityAndGlobalEntityConverter`] reference for read-only lookups.
86    fn to_converter(&self) -> &dyn EntityAndGlobalEntityConverter<E>;
87}
88
89impl<E: Copy + Eq + Hash + Send + Sync> GlobalEntitySpawner<E> for GlobalEntityMap<E> {
90    fn spawn(&mut self, world_entity: E, remote_entity_opt: Option<RemoteEntity>) -> GlobalEntity {
91        let global_entity_opt;
92
93        if let Some(remote_entity) = remote_entity_opt {
94            if let Some(global_entity) = self.reserved_global_entities.remove(&remote_entity) {
95                // global entity was reserved, update the mapping
96                let Some(entry) = self.global_to_entity_map.get_mut(&global_entity) else {
97                    panic!(
98                        "Global entity {:?} does not exist in the global to entity map",
99                        global_entity
100                    );
101                };
102                *entry = Some(world_entity);
103                global_entity_opt = Some(global_entity);
104            } else {
105                // global entity was not reserved, create a new one
106                global_entity_opt = None;
107            }
108        } else {
109            // local spawn, no remote entity
110            global_entity_opt = None;
111        };
112
113        let global_entity = if let Some(global_entity) = global_entity_opt {
114            global_entity
115        } else {
116            self.global_to_entity_map.insert(Some(world_entity))
117        };
118
119        self.entity_to_global_map
120            .insert(world_entity, global_entity);
121
122        global_entity
123    }
124
125    fn reserve_global_entity(&mut self, remote_entity: RemoteEntity) -> GlobalEntity {
126        if self.reserved_global_entities.contains_key(&remote_entity) {
127            panic!(
128                "Remote entity {:?} already has a reserved global entity",
129                remote_entity
130            );
131        }
132
133        let global_entity = self.global_to_entity_map.insert(None);
134        self.reserved_global_entities
135            .insert(remote_entity, global_entity);
136
137        // warn!(
138        //     "Reserving global entity {:?}, for remote entity {:?}",
139        //     global_entity, remote_entity
140        // );
141
142        global_entity
143    }
144
145    fn despawn_by_global(&mut self, global_entity: &GlobalEntity) {
146        let Some(Some(world_entity)) = self.global_to_entity_map.remove(global_entity) else {
147            panic!(
148                "Global entity {:?} does not exist in the global to entity map",
149                global_entity
150            );
151        };
152        self.entity_to_global_map.remove(&world_entity);
153    }
154
155    fn despawn_by_world(&mut self, world_entity: &E) {
156        let global_entity = self.entity_to_global_map.remove(world_entity).unwrap();
157        self.global_to_entity_map.remove(&global_entity);
158    }
159
160    fn to_converter(&self) -> &dyn EntityAndGlobalEntityConverter<E> {
161        self
162    }
163}