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

steel_core/entity/
registry.rs
