steel_core/block_entity/
registry.rs

1//! Block entity registry for creating block entity instances.
2
3use std::io::Cursor;
4use std::ops::Deref;
5use std::sync::{Arc, OnceLock, Weak};
6
7use simdnbt::borrow::NbtCompound as BorrowedRootNbtCompound;
8use simdnbt::borrow::{
9    BaseNbtCompound as BorrowedNbtCompound, read_compound as read_borrowed_compound,
10};
11use simdnbt::owned::NbtCompound;
12use steel_registry::block_entity_type::BlockEntityTypeRef;
13use steel_registry::vanilla_block_entity_types;
14use steel_registry::{REGISTRY, RegistryEntry, RegistryExt};
15use steel_utils::{BlockPos, BlockStateId};
16
17use super::SharedBlockEntity;
18use super::entities::{
19    BarrelBlockEntity, BeehiveBlockEntity, ComparatorBlockEntity, DaylightDetectorBlockEntity,
20    EndGatewayBlockEntity, EndPortalBlockEntity, PistonMovingBlockEntity, PotentSulfurBlockEntity,
21    RawBlockEntity, SignBlockEntity,
22};
23use crate::world::World;
24
25/// Factory function type for creating block entities.
26///
27/// Takes the world, position and block state, returns a new block entity instance.
28pub type BlockEntityFactory = fn(Weak<World>, BlockPos, BlockStateId) -> SharedBlockEntity;
29
30/// Registry entry for a block entity type.
31struct BlockEntityEntry {
32    /// Factory function to create instances.
33    factory: Option<BlockEntityFactory>,
34}
35
36/// Registry for block entity factories.
37///
38/// Maps `BlockEntityType` to factory functions that can create block entity instances.
39/// This is used when loading block entities from disk or when blocks with entities
40/// are placed.
41pub struct BlockEntityRegistry {
42    entries: Vec<BlockEntityEntry>,
43}
44
45impl BlockEntityRegistry {
46    /// Creates a new empty registry with entries for all block entity types.
47    #[must_use]
48    pub fn new() -> Self {
49        let count = REGISTRY.block_entity_types.len();
50        let entries = (0..count)
51            .map(|_| BlockEntityEntry { factory: None })
52            .collect();
53
54        Self { entries }
55    }
56
57    /// Registers a factory function for a block entity type.
58    pub fn register(&mut self, block_entity_type: BlockEntityTypeRef, factory: BlockEntityFactory) {
59        let id = block_entity_type.id();
60        self.entries[id].factory = Some(factory);
61    }
62
63    /// Creates a new block entity instance.
64    ///
65    /// Returns `None` if no factory is registered for the given type.
66    #[must_use]
67    pub fn create(
68        &self,
69        block_entity_type: BlockEntityTypeRef,
70        level: Weak<World>,
71        pos: BlockPos,
72        state: BlockStateId,
73    ) -> Option<SharedBlockEntity> {
74        let id = block_entity_type.id();
75        self.entries.get(id)?.factory.map(|f| f(level, pos, state))
76    }
77
78    /// Creates a block entity, falling back to an NBT-preserving raw entity.
79    ///
80    /// Use this for disk/worldgen paths where an unimplemented block entity type must still
81    /// survive save/load. Gameplay paths that require concrete behavior should call
82    /// [`Self::create`] and handle `None`.
83    #[must_use]
84    pub fn create_or_raw(
85        &self,
86        block_entity_type: BlockEntityTypeRef,
87        level: Weak<World>,
88        pos: BlockPos,
89        state: BlockStateId,
90    ) -> SharedBlockEntity {
91        let id = block_entity_type.id();
92        if let Some(factory) = self.entries.get(id).and_then(|entry| entry.factory) {
93            factory(level, pos, state)
94        } else {
95            Arc::new(RawBlockEntity::new(block_entity_type, level, pos, state))
96        }
97    }
98
99    /// Creates a block entity and loads borrowed NBT, falling back to raw preservation.
100    #[must_use]
101    pub fn create_and_load_or_raw(
102        &self,
103        block_entity_type: BlockEntityTypeRef,
104        level: Weak<World>,
105        pos: BlockPos,
106        state: BlockStateId,
107        nbt: &BorrowedNbtCompound<'_>,
108    ) -> SharedBlockEntity {
109        let id = block_entity_type.id();
110        if let Some(factory) = self.entries.get(id).and_then(|entry| entry.factory) {
111            let entity = factory(level, pos, state);
112            entity.load_additional(nbt);
113            entity
114        } else {
115            let nbt_view: BorrowedRootNbtCompound<'_, '_> = nbt.into();
116            Arc::new(RawBlockEntity::with_data(
117                block_entity_type,
118                level,
119                pos,
120                state,
121                nbt_view.to_owned(),
122            ))
123        }
124    }
125
126    /// Creates a block entity and loads owned NBT, falling back to raw preservation.
127    #[must_use]
128    pub fn create_and_load_owned_or_raw(
129        &self,
130        block_entity_type: BlockEntityTypeRef,
131        level: Weak<World>,
132        pos: BlockPos,
133        state: BlockStateId,
134        nbt: NbtCompound,
135    ) -> SharedBlockEntity {
136        let id = block_entity_type.id();
137        if let Some(factory) = self.entries.get(id).and_then(|entry| entry.factory) {
138            let entity = factory(level, pos, state);
139            let mut nbt_bytes = Vec::new();
140            nbt.write(&mut nbt_bytes);
141            if let Ok(borrowed) = read_borrowed_compound(&mut Cursor::new(&nbt_bytes)) {
142                entity.load_additional(&borrowed);
143            } else {
144                log::warn!(
145                    "failed to reborrow owned NBT for block entity {}",
146                    block_entity_type.key()
147                );
148            }
149            entity
150        } else {
151            Arc::new(RawBlockEntity::with_data(
152                block_entity_type,
153                level,
154                pos,
155                state,
156                nbt,
157            ))
158        }
159    }
160
161    /// Returns whether a factory is registered for the given type.
162    #[must_use]
163    pub fn has_factory(&self, block_entity_type: BlockEntityTypeRef) -> bool {
164        let id = block_entity_type.id();
165        self.entries.get(id).is_some_and(|e| e.factory.is_some())
166    }
167}
168
169impl Default for BlockEntityRegistry {
170    fn default() -> Self {
171        Self::new()
172    }
173}
174
175/// Wrapper for the global block entity registry that implements `Deref`.
176pub struct BlockEntityRegistryLock(OnceLock<BlockEntityRegistry>);
177
178impl Deref for BlockEntityRegistryLock {
179    type Target = BlockEntityRegistry;
180
181    fn deref(&self) -> &Self::Target {
182        self.0.get().expect("Block entity registry not initialized")
183    }
184}
185
186impl BlockEntityRegistryLock {
187    /// Sets the registry. Returns `Err` if already initialized.
188    pub fn set(&self, registry: BlockEntityRegistry) -> Result<(), BlockEntityRegistry> {
189        self.0.set(registry)
190    }
191
192    fn get_or_init(&self, init: impl FnOnce() -> BlockEntityRegistry) -> &BlockEntityRegistry {
193        self.0.get_or_init(init)
194    }
195}
196
197/// Global block entity registry.
198///
199/// Access via deref: `BLOCK_ENTITIES.create(type, pos, state)`
200pub static BLOCK_ENTITIES: BlockEntityRegistryLock = BlockEntityRegistryLock(OnceLock::new());
201
202/// Initializes the global block entity registry.
203///
204/// This should be called after the main registry is frozen. Repeated calls are a no-op.
205pub fn init_block_entities() {
206    BLOCK_ENTITIES.get_or_init(|| {
207        let mut registry = BlockEntityRegistry::new();
208
209        // Register sign block entity factory
210        registry.register(&vanilla_block_entity_types::SIGN, |level, pos, state| {
211            Arc::new(SignBlockEntity::new(level, pos, state))
212        });
213
214        // Register hanging sign block entity factory
215        registry.register(
216            &vanilla_block_entity_types::HANGING_SIGN,
217            |level, pos, state| Arc::new(SignBlockEntity::new_hanging(level, pos, state)),
218        );
219
220        // Register barrel block entity factory
221        registry.register(&vanilla_block_entity_types::BARREL, |level, pos, state| {
222            Arc::new(BarrelBlockEntity::new(level, pos, state))
223        });
224
225        // Register beehive block entity factory
226        registry.register(&vanilla_block_entity_types::BEEHIVE, |level, pos, state| {
227            Arc::new(BeehiveBlockEntity::new(level, pos, state))
228        });
229
230        // Register comparator block entity factory
231        registry.register(
232            &vanilla_block_entity_types::COMPARATOR,
233            |level, pos, state| Arc::new(ComparatorBlockEntity::new(level, pos, state)),
234        );
235
236        registry.register(
237            &vanilla_block_entity_types::DAYLIGHT_DETECTOR,
238            |level, pos, state| Arc::new(DaylightDetectorBlockEntity::new(level, pos, state)),
239        );
240
241        registry.register(&vanilla_block_entity_types::PISTON, |level, pos, state| {
242            Arc::new(PistonMovingBlockEntity::new(level, pos, state))
243        });
244
245        // Register End gateway block entity factory
246        registry.register(
247            &vanilla_block_entity_types::END_GATEWAY,
248            |level, pos, state| Arc::new(EndGatewayBlockEntity::new(level, pos, state)),
249        );
250
251        // Register End portal block entity factory
252        registry.register(
253            &vanilla_block_entity_types::END_PORTAL,
254            |level, pos, state| Arc::new(EndPortalBlockEntity::new(level, pos, state)),
255        );
256
257        // Register potent sulfur block entity factory
258        registry.register(
259            &vanilla_block_entity_types::POTENT_SULFUR,
260            |level, pos, state| Arc::new(PotentSulfurBlockEntity::new(level, pos, state)),
261        );
262
263        registry
264    });
265}
steel_core/block_entity/registry.rs

steel_core/block_entity/
registry.rs
