steel_core/entity/
registry.rs

1//! Entity registry for creating entity instances.
2
3use std::ops::Deref;
4use std::sync::{OnceLock, Weak};
5
6use glam::DVec3;
7use simdnbt::borrow::{
8    BaseNbtCompound as BorrowedNbtCompound, NbtCompound as BorrowedNbtCompoundView,
9};
10use steel_registry::RegistryExt;
11use steel_registry::entity_type::EntityTypeRef;
12use steel_registry::{REGISTRY, RegistryEntry};
13use uuid::Uuid;
14
15use super::generated_entities::register_entity_factories;
16use super::{
17    EntityBaseLoad, EntityBaseSaveData, EntityFireFreezeState, SharedEntity, next_entity_id,
18};
19use crate::world::World;
20
21/// Factory function type for creating entities.
22///
23/// Takes the entity type, entity ID, spawn position, and world reference.
24/// Returns a new entity instance. The entity ID should be obtained from
25/// `next_entity_id()`.
26pub type EntityFactory = fn(EntityTypeRef, i32, DVec3, Weak<World>) -> SharedEntity;
27
28/// Factory function type for loading entities from disk.
29///
30/// Takes the entity type and all base entity fields needed for reconstruction.
31pub type EntityLoadFactory = fn(EntityTypeRef, EntityBaseLoad) -> SharedEntity;
32
33/// Entity load request before the registry assigns a runtime ID.
34pub struct EntityLoadRequest {
35    /// Entity type to instantiate.
36    pub entity_type: EntityTypeRef,
37    /// Restored entity position.
38    pub position: DVec3,
39    /// Persisted entity UUID.
40    pub uuid: Uuid,
41    /// Restored velocity.
42    pub velocity: DVec3,
43    /// Restored yaw and pitch.
44    pub rotation: (f32, f32),
45    /// Restored accumulated fall distance.
46    pub fall_distance: f64,
47    /// Restored vanilla fire/freeze state.
48    pub fire_freeze: EntityFireFreezeState,
49    /// Restored ground-contact flag.
50    pub on_ground: bool,
51    /// Restored shared vanilla save data.
52    pub save_data: EntityBaseSaveData,
53    /// World reference for the loaded entity.
54    pub world: Weak<World>,
55}
56
57impl EntityLoadRequest {
58    fn into_base_load(self) -> (EntityTypeRef, EntityBaseLoad) {
59        (
60            self.entity_type,
61            EntityBaseLoad {
62                id: next_entity_id(),
63                position: self.position,
64                uuid: self.uuid,
65                velocity: self.velocity,
66                rotation: self.rotation,
67                fall_distance: self.fall_distance,
68                fire_freeze: self.fire_freeze,
69                on_ground: self.on_ground,
70                save_data: self.save_data,
71                world: self.world,
72            },
73        )
74    }
75}
76
77/// Registry entry for an entity type.
78struct EntityEntry {
79    /// Factory function to create new instances.
80    factory: Option<EntityFactory>,
81    /// Factory function to load instances from disk.
82    load_factory: Option<EntityLoadFactory>,
83}
84
85/// Registry for entity factories.
86///
87/// Maps `EntityType` to factory functions that can create entity instances.
88/// This is used when loading entities from disk or when entities are spawned.
89pub struct EntityRegistry {
90    entries: Vec<EntityEntry>,
91}
92
93impl EntityRegistry {
94    /// Completes the registered-entity portion of vanilla `Entity.load` after
95    /// the load factory has reconstructed the entity's base state.
96    fn finish_registered_load(entity: &SharedEntity, nbt: &BorrowedNbtCompound<'_>) {
97        let yaw = entity.rotation().0;
98        if let Some(living) = entity.as_living_entity() {
99            living.set_y_head_rot(yaw);
100            living.set_y_body_rot(yaw);
101        }
102
103        let nbt: BorrowedNbtCompoundView<'_, '_> = nbt.into();
104        entity.load_additional(nbt);
105        entity.set_old_position_to_current();
106        entity.base().set_old_rotation_to_current();
107        entity.sync_base_entity_data();
108    }
109
110    /// Creates a new empty registry with entries for all entity types.
111    #[must_use]
112    pub fn new() -> Self {
113        let count = REGISTRY.entity_types.len();
114        let entries = (0..count)
115            .map(|_| EntityEntry {
116                factory: None,
117                load_factory: None,
118            })
119            .collect();
120
121        Self { entries }
122    }
123
124    /// Registers a factory function for an entity type.
125    ///
126    /// # Panics
127    ///
128    /// Panics if a factory is already registered for the entity type.
129    pub fn register(&mut self, entity_type: EntityTypeRef, factory: EntityFactory) {
130        let id = entity_type.id();
131        assert!(
132            self.entries[id].factory.is_none(),
133            "entity factory for {} is already registered",
134            entity_type.key
135        );
136        self.entries[id].factory = Some(factory);
137    }
138
139    /// Registers a load factory function for an entity type.
140    ///
141    /// The load factory is used when loading entities from disk.
142    ///
143    /// # Panics
144    ///
145    /// Panics if a load factory is already registered for the entity type.
146    pub fn register_load(&mut self, entity_type: EntityTypeRef, factory: EntityLoadFactory) {
147        let id = entity_type.id();
148        assert!(
149            self.entries[id].load_factory.is_none(),
150            "entity load factory for {} is already registered",
151            entity_type.key
152        );
153        self.entries[id].load_factory = Some(factory);
154    }
155
156    /// Creates a new entity instance.
157    ///
158    /// Returns `None` if no factory is registered for the given type.
159    #[must_use]
160    pub fn create(
161        &self,
162        entity_type: EntityTypeRef,
163        entity_id: i32,
164        pos: DVec3,
165        world: Weak<World>,
166    ) -> Option<SharedEntity> {
167        let id = entity_type.id();
168        self.entries
169            .get(id)?
170            .factory
171            .map(|f| f(entity_type, entity_id, pos, world))
172    }
173
174    /// Creates an entity from persisted data.
175    ///
176    /// Returns `None` and logs a warning when Steel has no load factory for the type.
177    /// Skipped entities and their passenger trees are not retained when the chunk is saved.
178    #[must_use]
179    pub fn create_and_load(
180        &self,
181        request: EntityLoadRequest,
182        nbt: &BorrowedNbtCompound<'_>,
183    ) -> Option<SharedEntity> {
184        let entity_type = request.entity_type;
185        let id = entity_type.id();
186        let Some(load_factory) = self.entries.get(id).and_then(|entry| entry.load_factory) else {
187            log::warn!(
188                "Skipping unsupported entity {} while loading",
189                entity_type.key
190            );
191            return None;
192        };
193
194        let (_, load) = request.into_base_load();
195        let entity = load_factory(entity_type, load);
196        Self::finish_registered_load(&entity, nbt);
197        Some(entity)
198    }
199
200    /// Returns whether a factory is registered for the given type.
201    #[must_use]
202    pub fn has_factory(&self, entity_type: EntityTypeRef) -> bool {
203        let id = entity_type.id();
204        self.entries.get(id).is_some_and(|e| e.factory.is_some())
205    }
206}
207
208impl Default for EntityRegistry {
209    fn default() -> Self {
210        Self::new()
211    }
212}
213
214/// Wrapper for the global entity registry that implements `Deref`.
215pub struct EntityRegistryLock(OnceLock<EntityRegistry>);
216
217impl Deref for EntityRegistryLock {
218    type Target = EntityRegistry;
219
220    fn deref(&self) -> &Self::Target {
221        self.0.get().expect("Entity registry not initialized")
222    }
223}
224
225impl EntityRegistryLock {
226    /// Sets the registry. Returns `Err` if already initialized.
227    pub fn set(&self, registry: EntityRegistry) -> Result<(), EntityRegistry> {
228        self.0.set(registry)
229    }
230
231    fn get_or_init(&self, init: impl FnOnce() -> EntityRegistry) -> &EntityRegistry {
232        self.0.get_or_init(init)
233    }
234
235    /// Returns the initialized registry, if entity factories have been installed.
236    #[must_use]
237    pub fn get(&self) -> Option<&EntityRegistry> {
238        self.0.get()
239    }
240}
241
242/// Global entity registry.
243///
244/// Access via deref: `ENTITIES.create(type, entity_id, pos)`
245pub static ENTITIES: EntityRegistryLock = EntityRegistryLock(OnceLock::new());
246
247/// Initializes the global entity registry, once the main registry is frozen.
248pub fn init_entities() {
249    ENTITIES.get_or_init(|| {
250        let mut registry = EntityRegistry::new();
251        register_entity_factories(&mut registry);
252        registry
253    });
254}
255
256#[cfg(test)]
257mod tests {
258    use std::io::Cursor;
259
260    use simdnbt::borrow::read_compound as read_borrowed_compound;
261    use simdnbt::owned::NbtCompound;
262    use steel_registry::init_vanilla_registry;
263    use steel_registry::vanilla_entities;
264
265    use super::*;
266    use crate::test_support::TestEntity;
267
268    #[test]
269    fn registered_living_load_restores_current_head_and_body_yaw() {
270        init_vanilla_registry();
271        let mut registry = EntityRegistry::new();
272        register_entity_factories(&mut registry);
273        let mut bytes = Vec::new();
274        NbtCompound::new().write(&mut bytes);
275        let borrowed = read_borrowed_compound(&mut Cursor::new(&bytes))
276            .unwrap_or_else(|error| panic!("test nbt should reborrow: {error}"));
277
278        let Some(entity) = registry.create_and_load(
279            EntityLoadRequest {
280                entity_type: &vanilla_entities::PIG,
281                position: DVec3::ZERO,
282                uuid: Uuid::from_u128(1),
283                velocity: DVec3::ZERO,
284                rotation: (135.0, -20.0),
285                fall_distance: 0.0,
286                fire_freeze: EntityFireFreezeState::new(),
287                on_ground: false,
288                save_data: EntityBaseSaveData::new(),
289                world: Weak::new(),
290            },
291            &borrowed,
292        ) else {
293            panic!("registered pig should load");
294        };
295
296        assert_eq!(entity.rotation(), (135.0, -20.0));
297        assert_eq!(entity.base().old_rotation(), (135.0, -20.0));
298        let Some(living) = entity.as_living_entity() else {
299            panic!("registered pig should expose living behavior");
300        };
301        let living_rotation = living.living_rotation_state();
302        assert_eq!(living_rotation.y_head_rot().to_bits(), 135.0_f32.to_bits());
303        assert_eq!(living_rotation.y_body_rot().to_bits(), 135.0_f32.to_bits());
304        assert_eq!(living_rotation.y_head_rot_o().to_bits(), 0.0_f32.to_bits());
305        assert_eq!(living_rotation.y_body_rot_o().to_bits(), 0.0_f32.to_bits());
306    }
307
308    #[test]
309    fn create_and_load_skips_unregistered_entity_type() {
310        init_vanilla_registry();
311        let registry = EntityRegistry::new();
312        let mut bytes = Vec::new();
313        NbtCompound::new().write(&mut bytes);
314        let borrowed =
315            read_borrowed_compound(&mut Cursor::new(&bytes)).expect("test nbt should reborrow");
316
317        let entity = registry.create_and_load(
318            EntityLoadRequest {
319                entity_type: &vanilla_entities::VILLAGER,
320                position: DVec3::new(1.0, 2.0, 3.0),
321                uuid: Uuid::from_u128(1),
322                velocity: DVec3::new(0.1, 0.0, 0.2),
323                rotation: (45.0, 10.0),
324                fall_distance: 2.25,
325                fire_freeze: EntityFireFreezeState::new(),
326                on_ground: true,
327                save_data: EntityBaseSaveData {
328                    no_gravity: true,
329                    invulnerable: true,
330                    ..EntityBaseSaveData::new()
331                },
332                world: Weak::new(),
333            },
334            &borrowed,
335        );
336
337        assert!(entity.is_none());
338    }
339
340    #[test]
341    fn create_forwards_entity_type_to_factory() {
342        init_vanilla_registry();
343        let mut registry = EntityRegistry::new();
344        registry.register(
345            &vanilla_entities::OAK_BOAT,
346            |entity_type, id, pos, world| TestEntity::shared(id, pos, world, entity_type),
347        );
348
349        let Some(entity) =
350            registry.create(&vanilla_entities::OAK_BOAT, 5, DVec3::ZERO, Weak::new())
351        else {
352            panic!("registered entity factory should create an entity");
353        };
354
355        assert_eq!(entity.entity_type(), &vanilla_entities::OAK_BOAT);
356    }
357}
steel_core/entity/registry.rs

steel_core/entity/
registry.rs
