1use parking_lot::ArcMutexGuard;
8use parking_lot::RawMutex;
9use rustc_hash::FxHashMap;
10use std::borrow::Borrow;
11use std::mem;
12use std::ops::{Deref, DerefMut};
13use std::sync::Arc;
14use steel_utils::locks::Shared;
15use steel_utils::{Downcast as _, DowncastType, locks::SyncMutex};
16
17use crate::{
18 block_entity::{BlockEntityBase, SharedBlockEntity},
19 inventory::container::Container,
20 player::{Player, player_inventory::PlayerInventory},
21};
22use steel_registry::item_stack::ItemStack;
23
24type ContainerChangedCallback = Arc<dyn Fn() + Send + Sync>;
25
26#[derive(Clone)]
27struct ContainerOwner {
28 block_entity: Arc<BlockEntityBase>,
29 after_changed: Option<ContainerChangedCallback>,
30}
31
32pub type SharedContainer = Shared<dyn Container>;
34
35struct LockedContainer(ArcMutexGuard<RawMutex, dyn Container>);
36
37impl Deref for LockedContainer {
38 type Target = dyn Container;
39
40 fn deref(&self) -> &Self::Target {
41 &*self.0
42 }
43}
44
45impl DerefMut for LockedContainer {
46 fn deref_mut(&mut self) -> &mut Self::Target {
47 &mut *self.0
48 }
49}
50
51#[derive(Clone)]
56pub struct ContainerRef {
57 id: ContainerId,
58 source: SharedContainer,
59 owner: Option<ContainerOwner>,
60}
61
62impl<T> From<Shared<T>> for ContainerRef
63where
64 T: Container + 'static,
65{
66 fn from(container: Shared<T>) -> Self {
67 let id = ContainerId::from_arc(&container);
68 let container: SharedContainer = container;
69 Self {
70 id,
71 source: container,
72 owner: None,
73 }
74 }
75}
76
77impl<T> From<&Shared<T>> for ContainerRef
78where
79 T: Container + 'static,
80{
81 fn from(container: &Shared<T>) -> Self {
82 container.clone().into()
83 }
84}
85
86impl From<&ContainerRef> for ContainerRef {
87 fn from(r: &ContainerRef) -> Self {
88 r.clone()
89 }
90}
91
92impl From<&SharedContainer> for ContainerRef {
93 fn from(container: &SharedContainer) -> Self {
94 container.clone().into()
95 }
96}
97
98impl From<SharedContainer> for ContainerRef {
99 fn from(container: SharedContainer) -> Self {
100 Self {
101 id: ContainerId::from_arc(&container),
102 source: container,
103 owner: None,
104 }
105 }
106}
107
108impl ContainerRef {
109 #[must_use]
113 pub fn from_block_entity(block_entity: SharedBlockEntity) -> Option<Self> {
114 block_entity.container_ref()
115 }
116
117 #[must_use]
124 pub fn owned_by_block_entity(container: SharedContainer, owner: Arc<BlockEntityBase>) -> Self {
125 Self {
126 id: ContainerId::from_arc(&container),
127 source: container,
128 owner: Some(ContainerOwner {
129 block_entity: owner,
130 after_changed: None,
131 }),
132 }
133 }
134
135 #[must_use]
142 pub(crate) fn owned_by_block_entity_with_callback(
143 container: SharedContainer,
144 owner: Arc<BlockEntityBase>,
145 after_changed: ContainerChangedCallback,
146 ) -> Self {
147 Self {
148 id: ContainerId::from_arc(&container),
149 source: container,
150 owner: Some(ContainerOwner {
151 block_entity: owner,
152 after_changed: Some(after_changed),
153 }),
154 }
155 }
156
157 #[must_use]
159 pub const fn container_id(&self) -> ContainerId {
160 self.id
161 }
162
163 #[must_use]
165 pub fn still_valid(&self, player: &Player) -> bool {
166 self.owner
167 .as_ref()
168 .is_none_or(|owner| owner.block_entity.is_valid_container_for(player))
169 }
170
171 fn lock(&self) -> LockedContainer {
173 LockedContainer(SyncMutex::lock_arc(&self.source))
174 }
175}
176
177pub struct ContainerLockGuard {
197 sources: Vec<(ContainerId, ContainerRef)>,
199 guards: Vec<(ContainerId, LockedContainer)>,
201 id_to_index: FxHashMap<ContainerId, usize>,
203}
204
205impl ContainerLockGuard {
206 #[must_use]
212 pub fn lock_all<C>(containers: &[C]) -> Self
213 where
214 C: Borrow<ContainerRef>,
215 {
216 let mut sources: Vec<_> = containers
217 .iter()
218 .map(|c| (c.borrow().container_id(), c.borrow().clone()))
219 .collect();
220
221 sources.sort_by_key(|(id, _)| *id);
223
224 sources.dedup_by_key(|(id, _)| *id);
226
227 let mut guards = Vec::with_capacity(sources.len());
229 for (id, container) in &sources {
230 let guard = container.lock();
231 guards.push((*id, guard));
232 }
233
234 let id_to_index = guards
236 .iter()
237 .enumerate()
238 .map(|(idx, (id, _))| (*id, idx))
239 .collect();
240
241 Self {
242 sources,
243 guards,
244 id_to_index,
245 }
246 }
247
248 pub fn get_disjoint_mut<const N: usize>(
252 &mut self,
253 ids: [ContainerId; N],
254 ) -> Option<[&mut dyn Container; N]> {
255 let mut indices = [0usize; N];
256 for (i, id) in ids.iter().enumerate() {
257 indices[i] = *self.id_to_index.get(id)?;
258 }
259 let entries = self.guards.get_disjoint_mut(indices).ok()?;
260 Some(entries.map(|(_, locked)| &mut **locked as &mut dyn Container))
261 }
262
263 #[must_use]
268 pub fn relock(self, containers: &[&ContainerRef]) -> Self {
269 drop(self);
271 Self::lock_all(containers)
273 }
274
275 #[must_use]
277 pub fn get(&self, id: impl Into<ContainerId>) -> Option<&dyn Container> {
278 self.id_to_index
279 .get(&id.into())
280 .and_then(|&idx| self.guards.get(idx))
281 .map(|(_, guard)| &**guard as &dyn Container)
282 }
283
284 pub fn get_mut(&mut self, id: impl Into<ContainerId>) -> Option<&mut dyn Container> {
289 self.id_to_index
290 .get(&id.into())
291 .copied()
292 .and_then(|idx| self.guards.get_mut(idx))
293 .map(|(_, guard)| &mut **guard as &mut dyn Container)
294 }
295
296 pub fn set_item(&mut self, id: impl Into<ContainerId>, slot: usize, stack: ItemStack) -> bool {
301 let id = id.into();
302 let Some(&index) = self.id_to_index.get(&id) else {
303 return false;
304 };
305 self.guards[index].1.set_item(slot, stack);
306 let owner = self.sources[index].1.owner.clone();
307 self.notify_owner(owner);
308 true
309 }
310
311 pub fn remove_item(
313 &mut self,
314 id: impl Into<ContainerId>,
315 slot: usize,
316 amount: i32,
317 ) -> Option<ItemStack> {
318 let id = id.into();
319 let index = *self.id_to_index.get(&id)?;
320 let removed = self.guards.get_mut(index)?.1.remove_item(slot, amount);
321 if !removed.is_empty() {
322 let owner = self.sources.get(index)?.1.owner.clone();
323 self.notify_owner(owner);
324 }
325 Some(removed)
326 }
327
328 pub fn set_changed(&mut self, id: impl Into<ContainerId>) -> bool {
331 let id = id.into();
332 let Some(&index) = self.id_to_index.get(&id) else {
333 return false;
334 };
335 self.guards[index].1.set_changed();
336 let owner = self.sources[index].1.owner.clone();
337 self.notify_owner(owner);
338 true
339 }
340
341 pub(crate) fn run_unlocked<R>(&mut self, callback: impl FnOnce() -> R) -> R {
344 drop(mem::take(&mut self.guards));
345 let result = callback();
346 self.guards = self
347 .sources
348 .iter()
349 .map(|(id, container)| (*id, container.lock()))
350 .collect();
351 result
352 }
353
354 fn notify_owner(&mut self, owner: Option<ContainerOwner>) {
355 let Some(owner) = owner else {
356 return;
357 };
358
359 self.run_unlocked(|| {
360 owner.block_entity.set_changed();
361 if let Some(after_changed) = owner.after_changed {
362 after_changed();
363 }
364 });
365 }
366
367 #[must_use]
369 pub fn get_typed<T>(&self, id: impl Into<ContainerId>) -> Option<&T>
370 where
371 T: Container + DowncastType,
372 {
373 self.get(id)?.downcast_ref::<T>()
374 }
375
376 pub fn get_typed_mut<T>(&mut self, id: impl Into<ContainerId>) -> Option<&mut T>
381 where
382 T: Container + DowncastType,
383 {
384 self.get_mut(id)?.downcast_mut::<T>()
385 }
386
387 pub fn get_two_typed_mut<A, B>(
392 &mut self,
393 first: impl Into<ContainerId>,
394 second: impl Into<ContainerId>,
395 ) -> Option<(&mut A, &mut B)>
396 where
397 A: Container + DowncastType,
398 B: Container + DowncastType,
399 {
400 let first_index = *self.id_to_index.get(&first.into())?;
401 let second_index = *self.id_to_index.get(&second.into())?;
402 if first_index == second_index {
403 return None;
404 }
405
406 let (first, second): (&mut dyn Container, &mut dyn Container) =
407 if first_index < second_index {
408 let (before_second, from_second) = self.guards.split_at_mut(second_index);
409 (&mut *before_second[first_index].1, &mut *from_second[0].1)
410 } else {
411 let (before_first, from_first) = self.guards.split_at_mut(first_index);
412 (&mut *from_first[0].1, &mut *before_first[second_index].1)
413 };
414 Some((first.downcast_mut::<A>()?, second.downcast_mut::<B>()?))
415 }
416
417 #[must_use]
419 pub fn contains(&self, id: ContainerId) -> bool {
420 self.id_to_index.contains_key(&id)
421 }
422}
423
424#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
429pub struct ContainerId(usize);
430
431impl ContainerId {
432 pub fn from_arc<T: ?Sized>(arc: &Arc<T>) -> Self {
434 Self(Arc::as_ptr(arc).cast::<()>() as usize)
435 }
436}
437
438impl From<&Shared<PlayerInventory>> for ContainerId {
439 fn from(value: &Shared<PlayerInventory>) -> Self {
440 Self::from_arc(value)
441 }
442}
443
444#[cfg(test)]
445mod tests {
446 use std::sync::{Arc, Weak};
447
448 use steel_registry::blocks::block_state_ext::BlockStateExt as _;
449 use steel_registry::blocks::properties::{BlockStateProperties, Direction};
450 use steel_registry::{
451 init_vanilla_registry, item_stack::ItemStack, vanilla_block_entity_types, vanilla_blocks,
452 vanilla_items,
453 };
454 use steel_utils::types::UpdateFlags;
455 use steel_utils::{BlockPos, ChunkPos, locks::SyncMutex};
456
457 use super::{ContainerId, ContainerLockGuard, ContainerRef};
458 use crate::behavior::{BLOCK_BEHAVIORS, init_behaviors};
459 use crate::block_entity::{
460 SharedBlockEntity,
461 entities::{BarrelBlockEntity, UnimplementedBlockEntity},
462 init_block_entities,
463 };
464 use crate::inventory::container::{Container, CraftingContainer, ResultContainer};
465 use crate::inventory::slots::{NormalSlot, Slot as _};
466 use crate::test_support::{fresh_test_world, insert_ready_full_chunk};
467
468 #[test]
469 fn erased_container_ref_preserves_id_and_typed_access() {
470 let crafting = Arc::new(SyncMutex::new(CraftingContainer::new(2, 2)));
471 let id = ContainerId::from_arc(&crafting);
472 let container_ref = ContainerRef::from(Arc::clone(&crafting));
473
474 assert_eq!(container_ref.container_id(), id);
475
476 let mut guard = ContainerLockGuard::lock_all(&[&container_ref]);
477 let Some(typed) = guard.get_typed::<CraftingContainer>(id) else {
478 panic!("erased crafting container should retain its concrete type");
479 };
480 assert_eq!((typed.width(), typed.height()), (2, 2));
481 assert!(guard.get_typed::<ResultContainer>(id).is_none());
482 assert!(guard.get_typed_mut::<CraftingContainer>(id).is_some());
483 }
484
485 #[test]
486 fn block_entity_container_capability_is_independently_lockable() {
487 init_vanilla_registry();
488 let barrel = Arc::new(BarrelBlockEntity::new(
489 Weak::new(),
490 BlockPos::new(1, 2, 3),
491 vanilla_blocks::BARREL.default_state(),
492 ));
493 let block_entity: SharedBlockEntity = barrel.clone();
494 let Some(container_ref) = ContainerRef::from_block_entity(block_entity) else {
495 panic!("barrel block entity should expose Container");
496 };
497 let id = container_ref.container_id();
498
499 let guard = ContainerLockGuard::lock_all(&[&container_ref]);
500 assert_eq!(guard.get(id).map(Container::get_container_size), Some(27));
501 }
502
503 #[test]
504 fn non_container_block_entity_ref_is_rejected() {
505 init_vanilla_registry();
506 let block_entity: SharedBlockEntity = Arc::new(UnimplementedBlockEntity::new(
507 &vanilla_block_entity_types::END_PORTAL,
508 Weak::new(),
509 BlockPos::new(1, 2, 3),
510 vanilla_blocks::END_PORTAL.default_state(),
511 ));
512
513 assert!(ContainerRef::from_block_entity(block_entity).is_none());
514 }
515
516 #[test]
517 fn unlocked_callback_releases_and_reacquires_every_container() {
518 let crafting = Arc::new(SyncMutex::new(CraftingContainer::new(2, 2)));
519 let result = Arc::new(SyncMutex::new(ResultContainer::new()));
520 let crafting_ref = ContainerRef::from(Arc::clone(&crafting));
521 let result_ref = ContainerRef::from(Arc::clone(&result));
522 let mut guard = ContainerLockGuard::lock_all(&[&crafting_ref, &result_ref]);
523
524 let both_unlocked = guard.run_unlocked(|| {
525 let crafting_guard = crafting.try_lock();
526 let result_guard = result.try_lock();
527 crafting_guard.is_some() && result_guard.is_some()
528 });
529
530 assert!(both_unlocked);
531 assert!(guard.get(crafting_ref.container_id()).is_some());
532 assert!(guard.get(result_ref.container_id()).is_some());
533 }
534
535 #[test]
536 fn barrel_change_reenters_analog_read_without_holding_container_lock() {
537 init_vanilla_registry();
538 init_behaviors();
539 init_block_entities();
540 let world = fresh_test_world("barrel_comparator_reentry");
541 let barrel_pos = BlockPos::new(8, 64, 8);
542 let comparator_pos = barrel_pos.west();
543 insert_ready_full_chunk(&world, ChunkPos::from_block_pos(barrel_pos));
544 assert!(world.set_block(
545 comparator_pos.below(),
546 vanilla_blocks::STONE.default_state(),
547 UpdateFlags::UPDATE_NONE,
548 ));
549 assert!(world.set_block(
550 barrel_pos,
551 vanilla_blocks::BARREL.default_state(),
552 UpdateFlags::UPDATE_NONE,
553 ));
554 let comparator_state = vanilla_blocks::COMPARATOR
555 .default_state()
556 .set_value(&BlockStateProperties::HORIZONTAL_FACING, Direction::East);
557 assert!(world.set_block(comparator_pos, comparator_state, UpdateFlags::UPDATE_NONE,));
558 assert!(!world.has_scheduled_block_tick(comparator_pos, &vanilla_blocks::COMPARATOR));
559
560 let container_ref = ContainerRef::from_block_entity(
561 world
562 .get_block_entity(barrel_pos)
563 .expect("barrel should create its block entity"),
564 )
565 .expect("barrel should expose a container capability");
566 let slot = NormalSlot::new(container_ref.clone(), 0);
567 let mut guard = ContainerLockGuard::lock_all(&[&container_ref]);
568 slot.set_item(&mut guard, ItemStack::new(&vanilla_items::STONE));
569 drop(guard);
570
571 let analog = BLOCK_BEHAVIORS
572 .get_behavior(&vanilla_blocks::BARREL)
573 .get_analog_output_signal(
574 world.get_block_state(barrel_pos),
575 world.as_ref(),
576 barrel_pos,
577 Direction::West,
578 );
579 assert_eq!(analog, 1);
580 assert!(
581 world.has_scheduled_block_tick(comparator_pos, &vanilla_blocks::COMPARATOR),
582 "barrel callback did not schedule the comparator despite analog output {analog}"
583 );
584 }
585}