steel_core/entity/
registry.rs

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

steel_core/entity/
registry.rs
