Skip to main content

steel_core/block_entity/
storage.rs

1//! Block entity storage for chunks.
2
3use std::{fmt, ptr, sync::Arc};
4
5use rustc_hash::{FxHashMap, FxHashSet};
6use smallvec::SmallVec;
7use steel_utils::{BlockPos, BlockStateId, locks::SyncRwLock};
8
9#[cfg(test)]
10use super::BlockEntityLifecycleExt as _;
11use super::{BlockEntity, SharedBlockEntity};
12
13/// Storage for block entities in a chunk.
14///
15/// Ticker iteration is world-owned. This storage only owns chunk-persistent
16/// entities and lazy-promotion markers.
17pub(crate) struct BlockEntityStorage {
18    /// Related entity and marker state shares one lock. This makes replacement,
19    /// promotion, removal, and persistence snapshots linearizable.
20    entries: SyncRwLock<BlockEntityEntries>,
21}
22
23#[derive(Default)]
24struct BlockEntityEntries {
25    entities: FxHashMap<BlockPos, SharedBlockEntity>,
26    pending: FxHashSet<BlockPos>,
27}
28
29/// Atomic classification of one block-entity storage position.
30pub(crate) enum BlockEntityLookup {
31    /// A concrete entity currently owns the position.
32    Concrete(SharedBlockEntity),
33    /// A packed marker awaits lazy promotion.
34    Pending,
35    /// Neither a concrete entity nor a marker exists.
36    Absent,
37}
38
39/// Result of conditionally inserting a concrete entity into an empty slot.
40pub(crate) enum BlockEntityInsert {
41    /// Another concrete entity already owns the position.
42    Existing(SharedBlockEntity),
43    /// The new entity was inserted and may have staged callbacks.
44    Inserted(LifecycleDispatchers),
45}
46
47/// Entity ownership detached atomically with a replacing block-state write.
48pub(crate) struct DetachedBlockEntity {
49    /// The concrete entity that owned the old state, if any.
50    pub entity: Option<SharedBlockEntity>,
51    /// Whether this detachment owns dispatch of the queued removal callback.
52    pub dispatch_removed: bool,
53}
54
55#[derive(Default)]
56pub(crate) struct ClearedBlockEntities {
57    pub(crate) lifecycle_dispatchers: Vec<SharedBlockEntity>,
58    pub(crate) positions: Vec<BlockPos>,
59}
60
61pub(crate) type LifecycleDispatchers = SmallVec<[SharedBlockEntity; 2]>;
62
63impl BlockEntityStorage {
64    /// Creates a new empty block entity storage.
65    #[must_use]
66    pub(crate) fn new() -> Self {
67        Self {
68            entries: SyncRwLock::new(BlockEntityEntries::default()),
69        }
70    }
71
72    /// Gets a block entity at the given position.
73    #[must_use]
74    pub(crate) fn get(&self, pos: BlockPos) -> Option<SharedBlockEntity> {
75        self.entries.read().entities.get(&pos).cloned()
76    }
77
78    /// Returns whether `expected` is the exact current concrete owner.
79    #[must_use]
80    pub(crate) fn contains_same(&self, pos: BlockPos, expected: &SharedBlockEntity) -> bool {
81        self.entries
82            .read()
83            .entities
84            .get(&pos)
85            .is_some_and(|current| Arc::ptr_eq(current, expected))
86    }
87
88    /// Classifies a position with one lock acquisition and one entity-map probe.
89    #[must_use]
90    pub(crate) fn lookup(&self, pos: BlockPos) -> BlockEntityLookup {
91        let entries = self.entries.read();
92        if let Some(block_entity) = entries.entities.get(&pos) {
93            BlockEntityLookup::Concrete(Arc::clone(block_entity))
94        } else if entries.pending.contains(&pos) {
95            BlockEntityLookup::Pending
96        } else {
97            BlockEntityLookup::Absent
98        }
99    }
100
101    /// Returns all block entities in this storage.
102    #[must_use]
103    pub(crate) fn get_all(&self) -> Vec<SharedBlockEntity> {
104        self.entries
105            .read()
106            .entities
107            .values()
108            .filter(|entity| !entity.base().is_removed())
109            .cloned()
110            .collect()
111    }
112
113    /// Returns every concrete entry without applying `LevelChunk` lifecycle filtering.
114    #[must_use]
115    pub(crate) fn get_all_without_lifecycle_filter(&self) -> Vec<SharedBlockEntity> {
116        self.entries.read().entities.values().cloned().collect()
117    }
118
119    /// Atomically snapshots concrete entities and packed markers for persistence.
120    #[must_use]
121    pub(crate) fn save_snapshot(&self) -> (Vec<SharedBlockEntity>, Vec<BlockPos>) {
122        let entries = self.entries.read();
123        (
124            entries
125                .entities
126                .values()
127                .filter(|entity| !entity.base().is_removed())
128                .cloned()
129                .collect(),
130            entries.pending.iter().copied().collect(),
131        )
132    }
133
134    /// Atomically snapshots `Chunk` entries without applying Full lifecycle filtering.
135    ///
136    /// Vanilla `ProtoChunk` storage is a raw map: removed flags are neither changed nor consulted
137    /// until transfer into a `LevelChunk`.
138    #[must_use]
139    pub(crate) fn save_snapshot_without_lifecycle_filter(
140        &self,
141    ) -> (Vec<SharedBlockEntity>, Vec<BlockPos>) {
142        let entries = self.entries.read();
143        (
144            entries.entities.values().cloned().collect(),
145            entries.pending.iter().copied().collect(),
146        )
147    }
148
149    /// Returns packed block-entity positions without changing them.
150    #[must_use]
151    pub(crate) fn pending_positions(&self) -> Vec<BlockPos> {
152        self.entries.read().pending.iter().copied().collect()
153    }
154
155    /// Removes one invalid packed marker without constructing an entity.
156    pub(crate) fn remove_pending(&self, pos: BlockPos) {
157        self.entries.write().pending.remove(&pos);
158    }
159
160    /// Adds a packed marker only when no concrete entity owns the position.
161    pub(crate) fn set_pending(&self, pos: BlockPos) -> bool {
162        let mut entries = self.entries.write();
163        if entries.entities.contains_key(&pos) {
164            return false;
165        }
166        entries.pending.insert(pos)
167    }
168
169    /// Returns the number of block entities in this storage.
170    #[must_use]
171    pub(crate) fn len(&self) -> usize {
172        self.entries.read().entities.len()
173    }
174
175    /// Sets a `Chunk` block entity without invoking `LevelChunk` lifecycle callbacks.
176    ///
177    /// Vanilla `ProtoChunk` map replacement neither clears nor sets the removed flag.
178    #[must_use]
179    pub(crate) fn set_without_lifecycle(&self, block_entity: &SharedBlockEntity) -> bool {
180        let pos = block_entity.get_block_pos();
181        let mut entries = self.entries.write();
182        entries.pending.remove(&pos);
183        if entries
184            .entities
185            .get(&pos)
186            .is_some_and(|existing| Arc::ptr_eq(existing, block_entity))
187        {
188            return false;
189        }
190        entries.entities.insert(pos, Arc::clone(block_entity));
191        true
192    }
193
194    /// Adopts an existing pre-Full entity in place and stages its Full lifecycle updates.
195    ///
196    /// Returns `None` if `expected` no longer owns the position.
197    #[must_use]
198    pub(crate) fn adopt_if_same_staged(
199        &self,
200        pos: BlockPos,
201        expected: &SharedBlockEntity,
202        block_state: BlockStateId,
203    ) -> Option<LifecycleDispatchers> {
204        let dispatch = {
205            let entries = self.entries.write();
206            if !entries
207                .entities
208                .get(&pos)
209                .is_some_and(|current| Arc::ptr_eq(current, expected))
210            {
211                return None;
212            }
213            let dispatch_state = expected.base().queue_block_state_change(block_state);
214            let dispatch_clear = expected.base().queue_clear_removed();
215            dispatch_state || dispatch_clear
216        };
217        let mut lifecycle_dispatchers = LifecycleDispatchers::new();
218        if dispatch {
219            lifecycle_dispatchers.push(Arc::clone(expected));
220        }
221        Some(lifecycle_dispatchers)
222    }
223
224    /// Discards an invalid pre-Full entity only while it still owns the position.
225    ///
226    /// Promotion is an ownership transfer, not an unload, so this deliberately
227    /// queues no removal lifecycle event.
228    pub(crate) fn discard_if_same_without_lifecycle(
229        &self,
230        pos: BlockPos,
231        expected: &SharedBlockEntity,
232    ) -> bool {
233        let mut entries = self.entries.write();
234        if !entries
235            .entities
236            .get(&pos)
237            .is_some_and(|current| Arc::ptr_eq(current, expected))
238        {
239            return false;
240        }
241        entries.entities.remove(&pos);
242        true
243    }
244
245    /// Removes a block entity at the given position.
246    ///
247    /// Marks the entity as removed.
248    #[cfg(test)]
249    pub(crate) fn remove(&self, pos: BlockPos) -> bool {
250        let (removed, lifecycle_dispatchers) = self.remove_staged(pos);
251        for entity in lifecycle_dispatchers {
252            entity.dispatch_lifecycle_events();
253        }
254        removed
255    }
256
257    /// Removes an entity or marker while staging lifecycle callbacks for an outer lock boundary.
258    #[must_use]
259    pub(crate) fn remove_staged(&self, pos: BlockPos) -> (bool, LifecycleDispatchers) {
260        let (removed, removed_pending, dispatch_removed) = {
261            let mut entries = self.entries.write();
262            let removed = entries.entities.remove(&pos);
263            let removed_pending = entries.pending.remove(&pos);
264            let dispatch_removed = removed
265                .as_ref()
266                .is_some_and(|entity| entity.base().queue_set_removed());
267            (removed, removed_pending, dispatch_removed)
268        };
269        let mut lifecycle_dispatchers = LifecycleDispatchers::new();
270        if dispatch_removed && let Some(entity) = &removed {
271            lifecycle_dispatchers.push(Arc::clone(entity));
272        }
273        (removed.is_some() || removed_pending, lifecycle_dispatchers)
274    }
275
276    /// Detaches the old owner without invoking callbacks.
277    ///
278    /// The removal flag/event is queued while storage still owns the entity. This lets a
279    /// reentrant same-Arc insertion order its clear transition after removal even though the
280    /// caller delays callback dispatch until pre-removal side effects have run.
281    #[must_use]
282    pub(crate) fn detach_and_queue_removal(&self, pos: BlockPos) -> DetachedBlockEntity {
283        let mut entries = self.entries.write();
284        let entity = entries.entities.remove(&pos);
285        entries.pending.remove(&pos);
286        let dispatch_removed = entity
287            .as_ref()
288            .is_some_and(|entity| entity.base().queue_set_removed());
289        DetachedBlockEntity {
290            entity,
291            dispatch_removed,
292        }
293    }
294
295    /// Removes `Chunk` entity data without invoking `LevelChunk` lifecycle callbacks.
296    pub(crate) fn remove_without_lifecycle(&self, pos: BlockPos) -> bool {
297        let mut entries = self.entries.write();
298        let removed = entries.entities.remove(&pos).is_some();
299        let removed_pending = entries.pending.remove(&pos);
300        removed || removed_pending
301    }
302
303    /// Removes the entity only if `expected` still owns `pos`.
304    ///
305    /// This prevents a stale reader from deleting a concurrent replacement.
306    pub(crate) fn remove_if_same_and_removed(
307        &self,
308        pos: BlockPos,
309        expected: &SharedBlockEntity,
310    ) -> bool {
311        let mut entries = self.entries.write();
312        if !entries
313            .entities
314            .get(&pos)
315            .is_some_and(|current| Arc::ptr_eq(current, expected) && current.base().is_removed())
316        {
317            return false;
318        }
319        entries.entities.remove(&pos);
320        entries.pending.remove(&pos);
321        true
322    }
323
324    #[cfg(test)]
325    pub(crate) fn add_and_register(&self, block_entity: SharedBlockEntity) {
326        let block_state = block_entity.get_block_state();
327        let (_, lifecycle_dispatchers) = self.add_staged(&block_entity, block_state);
328        for entity in lifecycle_dispatchers {
329            entity.dispatch_lifecycle_events();
330        }
331    }
332
333    /// Adds an entity while staging callbacks for dispatch after outer chunk locks are dropped.
334    #[must_use]
335    pub(crate) fn add_staged(
336        &self,
337        block_entity: &SharedBlockEntity,
338        block_state: BlockStateId,
339    ) -> (bool, LifecycleDispatchers) {
340        self.set_inner(block_entity, block_state)
341    }
342
343    fn set_inner(
344        &self,
345        block_entity: &SharedBlockEntity,
346        block_state: BlockStateId,
347    ) -> (bool, LifecycleDispatchers) {
348        let pos = block_entity.get_block_pos();
349        let (inserted, dispatch_new, removed) = {
350            let mut entries = self.entries.write();
351            entries.pending.remove(&pos);
352
353            if entries
354                .entities
355                .get(&pos)
356                .is_some_and(|existing| Arc::ptr_eq(existing, block_entity))
357            {
358                let dispatch_state = block_entity.base().queue_block_state_change(block_state);
359                let dispatch_clear = block_entity.base().queue_clear_removed();
360                (false, dispatch_state || dispatch_clear, None)
361            } else {
362                let dispatch_state = block_entity.base().queue_block_state_change(block_state);
363                let dispatch_clear = block_entity.base().queue_clear_removed();
364                let removed = entries
365                    .entities
366                    .insert(pos, Arc::clone(block_entity))
367                    .map(|old| {
368                        let dispatch = old.base().queue_set_removed();
369                        (old, dispatch)
370                    });
371                (true, dispatch_state || dispatch_clear, removed)
372            }
373        };
374        let mut lifecycle_dispatchers = LifecycleDispatchers::new();
375        if dispatch_new {
376            lifecycle_dispatchers.push(Arc::clone(block_entity));
377        }
378        if let Some((old, true)) = removed {
379            lifecycle_dispatchers.push(old);
380        }
381        (inserted, lifecycle_dispatchers)
382    }
383
384    /// Inserts without replacing a concurrent concrete owner.
385    #[must_use]
386    pub(crate) fn insert_if_absent_staged(
387        &self,
388        block_entity: &SharedBlockEntity,
389        block_state: BlockStateId,
390    ) -> BlockEntityInsert {
391        let pos = block_entity.get_block_pos();
392        let dispatch_new = {
393            let mut entries = self.entries.write();
394            if let Some(existing) = entries.entities.get(&pos) {
395                return BlockEntityInsert::Existing(Arc::clone(existing));
396            }
397            entries.pending.remove(&pos);
398            let dispatch_state = block_entity.base().queue_block_state_change(block_state);
399            let dispatch_clear = block_entity.base().queue_clear_removed();
400            entries.entities.insert(pos, Arc::clone(block_entity));
401            dispatch_state || dispatch_clear
402        };
403        let mut lifecycle_dispatchers = LifecycleDispatchers::new();
404        if dispatch_new {
405            lifecycle_dispatchers.push(Arc::clone(block_entity));
406        }
407        BlockEntityInsert::Inserted(lifecycle_dispatchers)
408    }
409
410    fn promote_entry(
411        &self,
412        expected_pos: BlockPos,
413        block_entity: SharedBlockEntity,
414        block_state: Option<BlockStateId>,
415        update_lifecycle: bool,
416    ) -> (Option<SharedBlockEntity>, LifecycleDispatchers) {
417        let pos = block_entity.get_block_pos();
418        if pos != expected_pos {
419            return (None, LifecycleDispatchers::new());
420        }
421        let dispatch_new = {
422            let mut entries = self.entries.write();
423            if let Some(existing) = entries.entities.get(&pos) {
424                return (Some(Arc::clone(existing)), LifecycleDispatchers::new());
425            }
426            if !entries.pending.remove(&pos) {
427                return (None, LifecycleDispatchers::new());
428            }
429            let dispatch_state = block_state
430                .is_some_and(|state| block_entity.base().queue_block_state_change(state));
431            let dispatch_clear = update_lifecycle && block_entity.base().queue_clear_removed();
432            entries.entities.insert(pos, Arc::clone(&block_entity));
433            dispatch_state || dispatch_clear
434        };
435        let mut lifecycle_dispatchers = LifecycleDispatchers::new();
436        if dispatch_new {
437            lifecycle_dispatchers.push(Arc::clone(&block_entity));
438        }
439        (Some(block_entity), lifecycle_dispatchers)
440    }
441
442    /// Atomically replaces a packed marker with a concrete proto entity without lifecycle work.
443    pub(crate) fn promote_without_lifecycle(
444        &self,
445        expected_pos: BlockPos,
446        block_entity: SharedBlockEntity,
447    ) -> Option<SharedBlockEntity> {
448        self.promote_entry(expected_pos, block_entity, None, false)
449            .0
450    }
451
452    /// Promotes a marker while staging callbacks for dispatch after outer locks are dropped.
453    #[must_use]
454    pub(crate) fn promote_staged(
455        &self,
456        expected_pos: BlockPos,
457        block_state: BlockStateId,
458        block_entity: SharedBlockEntity,
459    ) -> (Option<SharedBlockEntity>, LifecycleDispatchers) {
460        self.promote_entry(expected_pos, block_entity, Some(block_state), true)
461    }
462
463    /// Updates cached state only while `block_entity` still owns `pos`.
464    #[must_use]
465    pub(crate) fn update_if_same_staged(
466        &self,
467        pos: BlockPos,
468        block_entity: &SharedBlockEntity,
469        block_state: BlockStateId,
470    ) -> (bool, LifecycleDispatchers) {
471        let dispatch_state = {
472            let entries = self.entries.write();
473            if !entries
474                .entities
475                .get(&pos)
476                .is_some_and(|current| Arc::ptr_eq(current, block_entity))
477            {
478                return (false, LifecycleDispatchers::new());
479            }
480            block_entity.base().queue_block_state_change(block_state)
481        };
482        let mut lifecycle_dispatchers = LifecycleDispatchers::new();
483        if dispatch_state {
484            lifecycle_dispatchers.push(Arc::clone(block_entity));
485        }
486        (true, lifecycle_dispatchers)
487    }
488
489    /// Removes `expected` only while it still owns `pos`, staging its callback.
490    #[must_use]
491    pub(crate) fn remove_if_same_staged(
492        &self,
493        pos: BlockPos,
494        expected: &dyn BlockEntity,
495    ) -> (bool, LifecycleDispatchers) {
496        let (removed, dispatch_removed) = {
497            let mut entries = self.entries.write();
498            if !entries
499                .entities
500                .get(&pos)
501                .is_some_and(|current| ptr::addr_eq(current.as_ref(), expected))
502            {
503                return (false, LifecycleDispatchers::new());
504            }
505            let Some(removed) = entries.entities.remove(&pos) else {
506                return (false, LifecycleDispatchers::new());
507            };
508            entries.pending.remove(&pos);
509            let dispatch_removed = removed.base().queue_set_removed();
510            (removed, dispatch_removed)
511        };
512        let mut lifecycle_dispatchers = LifecycleDispatchers::new();
513        if dispatch_removed {
514            lifecycle_dispatchers.push(removed);
515        }
516        (true, lifecycle_dispatchers)
517    }
518
519    /// Clears storage and returns entities whose lifecycle callback dispatcher this call owns.
520    ///
521    /// Callers that hold outer chunk-map or holder guards can drop them before dispatching.
522    #[must_use]
523    pub(crate) fn clear_and_stage_lifecycle_callbacks(&self) -> ClearedBlockEntities {
524        let mut entries = self.entries.write();
525        let mut lifecycle_dispatchers = Vec::with_capacity(entries.entities.len());
526        let mut positions = Vec::with_capacity(entries.entities.len());
527        for (&pos, entity) in &entries.entities {
528            positions.push(pos);
529            if entity.base().queue_set_removed() {
530                lifecycle_dispatchers.push(Arc::clone(entity));
531            }
532        }
533        entries.entities.clear();
534        entries.pending.clear();
535        ClearedBlockEntities {
536            lifecycle_dispatchers,
537            positions,
538        }
539    }
540
541    /// Clears `Chunk` entity data without invoking `LevelChunk` lifecycle callbacks.
542    pub(crate) fn clear_without_lifecycle(&self) {
543        let mut entries = self.entries.write();
544        entries.entities.clear();
545        entries.pending.clear();
546    }
547}
548
549impl Default for BlockEntityStorage {
550    fn default() -> Self {
551        Self::new()
552    }
553}
554
555impl fmt::Debug for BlockEntityStorage {
556    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
557        f.debug_struct("BlockEntityStorage")
558            .field("len", &self.len())
559            .finish_non_exhaustive()
560    }
561}
562
563#[cfg(test)]
564mod tests {
565    use std::sync::{
566        Arc, Weak,
567        atomic::{AtomicBool, Ordering},
568    };
569
570    use simdnbt::{borrow::BaseNbtCompound as BorrowedNbtCompound, owned::NbtCompound};
571    use steel_registry::{init_vanilla_registry, vanilla_block_entity_types, vanilla_blocks};
572    use steel_utils::{DowncastType, DowncastTypeKey, locks::SyncMutex};
573
574    use super::*;
575    use crate::block_entity::{BlockEntity, BlockEntityBase, entities::SignBlockEntity};
576
577    struct ReentrantLifecycleBlockEntity {
578        base: BlockEntityBase,
579        reenter_on_remove: AtomicBool,
580        events: SyncMutex<Vec<&'static str>>,
581    }
582
583    // SAFETY: This test-only key uniquely identifies this concrete test implementation.
584    unsafe impl DowncastType for ReentrantLifecycleBlockEntity {
585        const TYPE_KEY: DowncastTypeKey =
586            DowncastTypeKey::new("steel:test/block_entity/reentrant_lifecycle");
587    }
588
589    impl BlockEntity for ReentrantLifecycleBlockEntity {
590        fn base(&self) -> &BlockEntityBase {
591            &self.base
592        }
593
594        fn on_set_removed(&self) {
595            self.events.lock().push("removed");
596            if self.reenter_on_remove.swap(false, Ordering::AcqRel) {
597                self.clear_removed();
598            }
599        }
600
601        fn on_clear_removed(&self) {
602            self.events.lock().push("cleared");
603        }
604
605        fn on_block_state_changed(&self, _state: BlockStateId) {
606            self.events.lock().push("state");
607        }
608
609        fn load_additional(&self, _nbt: &BorrowedNbtCompound<'_>) {}
610
611        fn save_additional(&self, _nbt: &mut NbtCompound) {}
612    }
613
614    #[test]
615    fn readding_the_same_entity_preserves_ownership_and_clears_the_marker() {
616        init_vanilla_registry();
617        let storage = BlockEntityStorage::new();
618        let entity: SharedBlockEntity = Arc::new(SignBlockEntity::new(
619            Weak::new(),
620            BlockPos::new(1, 2, 3),
621            vanilla_blocks::OAK_SIGN.default_state(),
622        ));
623
624        assert!(storage.set_pending(entity.get_block_pos()));
625        let (concrete, pending) = storage.save_snapshot();
626        assert!(concrete.is_empty());
627        assert_eq!(pending, [entity.get_block_pos()]);
628
629        storage.add_and_register(Arc::clone(&entity));
630        storage.add_and_register(Arc::clone(&entity));
631
632        assert_eq!(storage.len(), 1);
633        assert!(!entity.is_removed());
634        let (concrete, pending) = storage.save_snapshot();
635        assert_eq!(concrete.len(), 1);
636        assert!(pending.is_empty());
637    }
638
639    #[test]
640    fn stale_removed_cleanup_cannot_delete_a_same_arc_revival() {
641        init_vanilla_registry();
642        let storage = BlockEntityStorage::new();
643        let entity: SharedBlockEntity = Arc::new(SignBlockEntity::new(
644            Weak::new(),
645            BlockPos::new(1, 2, 3),
646            vanilla_blocks::OAK_SIGN.default_state(),
647        ));
648        storage.add_and_register(Arc::clone(&entity));
649        entity.set_removed();
650        storage.add_and_register(Arc::clone(&entity));
651
652        assert!(!storage.remove_if_same_and_removed(entity.get_block_pos(), &entity));
653        let Some(current) = storage.get(entity.get_block_pos()) else {
654            panic!("revived entity should remain stored");
655        };
656        assert!(Arc::ptr_eq(&entity, &current));
657        assert!(!entity.is_removed());
658    }
659
660    #[test]
661    fn insert_if_absent_preserves_the_concurrent_owner() {
662        init_vanilla_registry();
663        let storage = BlockEntityStorage::new();
664        let pos = BlockPos::new(1, 2, 3);
665        let state = vanilla_blocks::OAK_SIGN.default_state();
666        let owner: SharedBlockEntity = Arc::new(SignBlockEntity::new(Weak::new(), pos, state));
667        let challenger: SharedBlockEntity = Arc::new(SignBlockEntity::new(Weak::new(), pos, state));
668        storage.add_and_register(Arc::clone(&owner));
669
670        let result = storage.insert_if_absent_staged(&challenger, state);
671        let BlockEntityInsert::Existing(existing) = result else {
672            panic!("the existing owner should win insertion");
673        };
674        assert!(Arc::ptr_eq(&owner, &existing));
675        let Some(stored) = storage.get(pos) else {
676            panic!("the existing owner should remain stored");
677        };
678        assert!(Arc::ptr_eq(&owner, &stored));
679    }
680
681    #[test]
682    fn lifecycle_callbacks_are_reentrant_and_keep_transition_order() {
683        init_vanilla_registry();
684        let concrete = Arc::new(ReentrantLifecycleBlockEntity {
685            base: BlockEntityBase::new(
686                &vanilla_block_entity_types::BARREL,
687                Weak::new(),
688                BlockPos::new(1, 2, 3),
689                vanilla_blocks::BARREL.default_state(),
690            ),
691            reenter_on_remove: AtomicBool::new(true),
692            events: SyncMutex::new(Vec::new()),
693        });
694        let entity: SharedBlockEntity = concrete.clone();
695        let storage = BlockEntityStorage::new();
696        storage.add_and_register(entity);
697
698        assert!(storage.remove(concrete.get_block_pos()));
699        assert_eq!(*concrete.events.lock(), ["removed", "cleared"]);
700        assert!(!concrete.is_removed());
701    }
702
703    #[test]
704    fn repeated_set_removed_calls_remain_observable() {
705        init_vanilla_registry();
706        let concrete = Arc::new(ReentrantLifecycleBlockEntity {
707            base: BlockEntityBase::new(
708                &vanilla_block_entity_types::BARREL,
709                Weak::new(),
710                BlockPos::new(1, 2, 3),
711                vanilla_blocks::BARREL.default_state(),
712            ),
713            reenter_on_remove: AtomicBool::new(false),
714            events: SyncMutex::new(Vec::new()),
715        });
716
717        concrete.set_removed();
718        concrete.set_removed();
719
720        assert_eq!(*concrete.events.lock(), ["removed", "removed"]);
721        assert!(concrete.is_removed());
722    }
723
724    #[test]
725    fn detached_dispatcher_preserves_same_arc_revival_order() {
726        init_vanilla_registry();
727        let concrete = Arc::new(ReentrantLifecycleBlockEntity {
728            base: BlockEntityBase::new(
729                &vanilla_block_entity_types::BARREL,
730                Weak::new(),
731                BlockPos::new(1, 2, 3),
732                vanilla_blocks::BARREL.default_state(),
733            ),
734            reenter_on_remove: AtomicBool::new(false),
735            events: SyncMutex::new(Vec::new()),
736        });
737        let entity: SharedBlockEntity = concrete.clone();
738        let storage = BlockEntityStorage::new();
739        storage.add_and_register(Arc::clone(&entity));
740
741        let detached = storage.detach_and_queue_removal(entity.get_block_pos());
742        let Some(detached_entity) = detached.entity else {
743            panic!("the stored entity should be detached");
744        };
745        assert!(Arc::ptr_eq(&entity, &detached_entity));
746        assert!(entity.is_removed());
747
748        let (_, revival_dispatchers) =
749            storage.add_staged(&entity, vanilla_blocks::BARREL.default_state());
750        assert!(revival_dispatchers.is_empty());
751        assert!(!entity.is_removed());
752        assert!(concrete.events.lock().is_empty());
753
754        if detached.dispatch_removed {
755            detached_entity.dispatch_lifecycle_events();
756        }
757        assert_eq!(*concrete.events.lock(), ["removed", "cleared"]);
758        let Some(current) = storage.get(entity.get_block_pos()) else {
759            panic!("the revived entity should remain stored");
760        };
761        assert!(Arc::ptr_eq(&entity, &current));
762    }
763
764    #[test]
765    fn cached_state_callback_is_staged_after_storage_commit() {
766        init_vanilla_registry();
767        let copper = vanilla_blocks::COPPER_CHEST.default_state();
768        let exposed = vanilla_blocks::EXPOSED_COPPER_CHEST.default_state();
769        let concrete = Arc::new(ReentrantLifecycleBlockEntity {
770            base: BlockEntityBase::new(
771                &vanilla_block_entity_types::CHEST,
772                Weak::new(),
773                BlockPos::new(1, 2, 3),
774                copper,
775            ),
776            reenter_on_remove: AtomicBool::new(false),
777            events: SyncMutex::new(Vec::new()),
778        });
779        let entity: SharedBlockEntity = concrete.clone();
780        let storage = BlockEntityStorage::new();
781        storage.add_and_register(Arc::clone(&entity));
782
783        let (_, lifecycle_dispatchers) = storage.add_staged(&entity, exposed);
784        assert_eq!(entity.get_block_state(), exposed);
785        assert!(concrete.events.lock().is_empty());
786        for dispatcher in lifecycle_dispatchers {
787            dispatcher.dispatch_lifecycle_events();
788        }
789        assert_eq!(*concrete.events.lock(), ["state"]);
790    }
791}