diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index 6d0fd231fc81f..2013ec583b272 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -58,11 +58,12 @@ pub struct World { pub(crate) storages: Storages, pub(crate) bundles: Bundles, pub(crate) removed_components: SparseSet>, - /// Access cache used by [WorldCell]. - pub(crate) archetype_component_access: ArchetypeComponentAccess, main_thread_validator: MainThreadValidator, pub(crate) change_tick: AtomicU32, pub(crate) last_change_tick: u32, + /// Access cache used by [WorldCell]. + pub(crate) archetype_component_access: ArchetypeComponentAccess, + pub(crate) entity_component_access: EntityComponentAccess, } impl Default for World { @@ -75,12 +76,13 @@ impl Default for World { storages: Default::default(), bundles: Default::default(), removed_components: Default::default(), - archetype_component_access: Default::default(), main_thread_validator: Default::default(), // Default value is `1`, and `last_change_tick`s default to `0`, such that changes // are detected on first system runs and for direct world queries. change_tick: AtomicU32::new(1), last_change_tick: 0, + archetype_component_access: Default::default(), + entity_component_access: Default::default(), } } } diff --git a/crates/bevy_ecs/src/world/world_cell.rs b/crates/bevy_ecs/src/world/world_cell.rs index 9db47dfd8c0a2..4bdc70b7f7b87 100644 --- a/crates/bevy_ecs/src/world/world_cell.rs +++ b/crates/bevy_ecs/src/world/world_cell.rs @@ -1,5 +1,10 @@ +use bevy_utils::HashMap; + use crate::{ archetype::ArchetypeComponentId, + component::ComponentId, + entity::Entity, + prelude::Component, storage::SparseSet, system::Resource, world::{Mut, World}, @@ -15,30 +20,18 @@ use std::{ /// World in a way that violates Rust's mutability rules will panic thanks to runtime checks. pub struct WorldCell<'w> { pub(crate) world: &'w mut World, - pub(crate) access: Rc>, + pub(crate) resource_access: Rc>, + pub(crate) component_access: Rc>, } +#[derive(Default)] pub(crate) struct ArchetypeComponentAccess { access: SparseSet, } -impl Default for ArchetypeComponentAccess { - fn default() -> Self { - Self { - access: SparseSet::new(), - } - } -} - const UNIQUE_ACCESS: usize = 0; const BASE_ACCESS: usize = 1; impl ArchetypeComponentAccess { - const fn new() -> Self { - Self { - access: SparseSet::new(), - } - } - fn read(&mut self, id: ArchetypeComponentId) -> bool { let id_access = self.access.get_or_insert_with(id, || BASE_ACCESS); if *id_access == UNIQUE_ACCESS { @@ -70,35 +63,105 @@ impl ArchetypeComponentAccess { } } +#[derive(Default)] +pub(crate) struct EntityComponentAccess { + access: HashMap<(Entity, ComponentId), usize>, +} + +impl EntityComponentAccess { + fn read(&mut self, id: (Entity, ComponentId)) -> bool { + let id_access = self.access.entry(id).or_insert(BASE_ACCESS); + if *id_access == UNIQUE_ACCESS { + false + } else { + *id_access += 1; + true + } + } + + fn drop_read(&mut self, id: (Entity, ComponentId)) { + let id_access = self.access.entry(id).or_insert(BASE_ACCESS); + *id_access -= 1; + } + + fn write(&mut self, id: (Entity, ComponentId)) -> bool { + let id_access = self.access.entry(id).or_insert(BASE_ACCESS); + if *id_access == BASE_ACCESS { + *id_access = UNIQUE_ACCESS; + true + } else { + false + } + } + + fn drop_write(&mut self, id: (Entity, ComponentId)) { + let id_access = self.access.entry(id).or_insert(BASE_ACCESS); + *id_access = BASE_ACCESS; + } +} + impl<'w> Drop for WorldCell<'w> { fn drop(&mut self) { - let mut access = self.access.borrow_mut(); - // give world ArchetypeComponentAccess back to reuse allocations - std::mem::swap(&mut self.world.archetype_component_access, &mut *access); + // give world {Archetype,Entity}ComponentAccess back to reuse allocations + let mut resource_access = self.resource_access.borrow_mut(); + std::mem::swap( + &mut self.world.archetype_component_access, + &mut *resource_access, + ); + let mut component_access = self.component_access.borrow_mut(); + std::mem::swap( + &mut self.world.entity_component_access, + &mut *component_access, + ); + } +} + +enum WorldCellId { + Resource(ArchetypeComponentId), + Component((Entity, ComponentId)), +} +impl From for WorldCellId { + fn from(id: ArchetypeComponentId) -> Self { + WorldCellId::Resource(id) } } pub struct WorldBorrow<'w, T> { value: &'w T, - archetype_component_id: ArchetypeComponentId, - access: Rc>, + id: WorldCellId, + resource_access: Rc>, + component_access: Rc>, } impl<'w, T> WorldBorrow<'w, T> { fn new( value: &'w T, - archetype_component_id: ArchetypeComponentId, - access: Rc>, + id: WorldCellId, + resource_access: Rc>, + component_access: Rc>, ) -> Self { - assert!( - access.borrow_mut().read(archetype_component_id), - "Attempted to immutably access {}, but it is already mutably borrowed", - std::any::type_name::(), - ); + match id { + WorldCellId::Resource(archetype_component_id) => { + assert!( + resource_access.borrow_mut().read(archetype_component_id), + "Attempted to immutably access {}, but it is already mutably borrowed", + std::any::type_name::(), + ); + } + WorldCellId::Component((entity, component_id)) => { + assert!( + component_access.borrow_mut().read((entity, component_id)), + "Attempted to immutably access component {} on entity {:?}, but it is already mutably borrowed", + std::any::type_name::(), + entity, + ); + } + } Self { value, - archetype_component_id, - access, + id, + resource_access, + component_access, } } } @@ -114,32 +177,49 @@ impl<'w, T> Deref for WorldBorrow<'w, T> { impl<'w, T> Drop for WorldBorrow<'w, T> { fn drop(&mut self) { - let mut access = self.access.borrow_mut(); - access.drop_read(self.archetype_component_id); + match self.id { + WorldCellId::Resource(id) => self.resource_access.borrow_mut().drop_read(id), + WorldCellId::Component(id) => self.component_access.borrow_mut().drop_read(id), + } } } pub struct WorldBorrowMut<'w, T> { value: Mut<'w, T>, - archetype_component_id: ArchetypeComponentId, - access: Rc>, + id: WorldCellId, + resource_access: Rc>, + component_access: Rc>, } impl<'w, T> WorldBorrowMut<'w, T> { fn new( value: Mut<'w, T>, - archetype_component_id: ArchetypeComponentId, - access: Rc>, + id: WorldCellId, + resource_access: Rc>, + component_access: Rc>, ) -> Self { - assert!( - access.borrow_mut().write(archetype_component_id), - "Attempted to mutably access {}, but it is already mutably borrowed", - std::any::type_name::(), - ); + match id { + WorldCellId::Resource(archetype_component_id) => { + assert!( + resource_access.borrow_mut().write(archetype_component_id), + "Attempted to mutably access {}, but it is already mutably borrowed", + std::any::type_name::(), + ); + } + WorldCellId::Component((entity, component)) => { + assert!( + component_access.borrow_mut().write((entity, component)), + "Attempted to mutably access component {} at entity {:?}, but it is already mutably borrowed", + std::any::type_name::(), + entity + ); + } + } Self { value, - archetype_component_id, - access, + id, + resource_access, + component_access, } } } @@ -161,22 +241,27 @@ impl<'w, T> DerefMut for WorldBorrowMut<'w, T> { impl<'w, T> Drop for WorldBorrowMut<'w, T> { fn drop(&mut self) { - let mut access = self.access.borrow_mut(); - access.drop_write(self.archetype_component_id); + match self.id { + WorldCellId::Resource(id) => { + self.resource_access.borrow_mut().drop_write(id); + } + WorldCellId::Component(id) => { + self.component_access.borrow_mut().drop_write(id); + } + } } } impl<'w> WorldCell<'w> { pub(crate) fn new(world: &'w mut World) -> Self { // this is cheap because ArchetypeComponentAccess::new() is const / allocation free - let access = std::mem::replace( - &mut world.archetype_component_access, - ArchetypeComponentAccess::new(), - ); + let resource_access = std::mem::take(&mut world.archetype_component_access); + let component_access = std::mem::take(&mut world.entity_component_access); // world's ArchetypeComponentAccess is recycled to cut down on allocations Self { world, - access: Rc::new(RefCell::new(access)), + resource_access: Rc::new(RefCell::new(resource_access)), + component_access: Rc::new(RefCell::new(component_access)), } } @@ -188,8 +273,9 @@ impl<'w> WorldCell<'w> { Some(WorldBorrow::new( // SAFETY: ComponentId matches TypeId unsafe { self.world.get_resource_with_id(component_id)? }, - archetype_component_id, - self.access.clone(), + archetype_component_id.into(), + self.resource_access.clone(), + self.component_access.clone(), )) } @@ -223,8 +309,9 @@ impl<'w> WorldCell<'w> { self.world .get_resource_unchecked_mut_with_id(component_id)? }, - archetype_component_id, - self.access.clone(), + archetype_component_id.into(), + self.resource_access.clone(), + self.component_access.clone(), )) } @@ -255,8 +342,9 @@ impl<'w> WorldCell<'w> { Some(WorldBorrow::new( // SAFETY: ComponentId matches TypeId unsafe { self.world.get_non_send_with_id(component_id)? }, - archetype_component_id, - self.access.clone(), + archetype_component_id.into(), + self.resource_access.clone(), + self.component_access.clone(), )) } @@ -290,8 +378,9 @@ impl<'w> WorldCell<'w> { self.world .get_non_send_unchecked_mut_with_id(component_id)? }, - archetype_component_id, - self.access.clone(), + archetype_component_id.into(), + self.resource_access.clone(), + self.component_access.clone(), )) } @@ -313,19 +402,49 @@ impl<'w> WorldCell<'w> { ), } } + + /// Gets a reference to the component of the given type + pub fn get_component(&self, entity: Entity) -> Option> { + let component_id = self.world.components.get_id(TypeId::of::())?; + Some(WorldBorrow::new( + // SAFETY: ComponentId matches TypeId + self.world.get(entity)?, + WorldCellId::Component((entity, component_id)), + self.resource_access.clone(), + self.component_access.clone(), + )) + } + /// Gets a mutable reference to the component of the given type + pub fn get_component_mut(&self, entity: Entity) -> Option> { + let component_id = self.world.components.get_id(TypeId::of::())?; + let last_change_tick = self.world.last_change_tick(); + let change_tick = self.world.read_change_tick(); + Some(WorldBorrowMut::new( + // SAFETY: ComponentId matches TypeId + unsafe { + self.world + .get_entity(entity)? + .get_unchecked_mut(last_change_tick, change_tick)? + }, + WorldCellId::Component((entity, component_id)), + self.resource_access.clone(), + self.component_access.clone(), + )) + } } #[cfg(test)] mod tests { use super::BASE_ACCESS; use crate as bevy_ecs; + use crate::prelude::Component; use crate::{archetype::ArchetypeId, system::Resource, world::World}; use std::any::TypeId; - #[derive(Resource)] + #[derive(Component, Resource)] struct A(u32); - #[derive(Resource)] + #[derive(Component, Resource)] struct B(u64); #[test] @@ -431,4 +550,33 @@ mod tests { let _value_a = cell.resource::(); let _value_b = cell.resource::(); } + + #[test] + #[should_panic] + fn world_cell_component_mut_and_ref() { + let mut world = World::default(); + let entity = world.spawn().insert(A(1)).id(); + let cell = world.cell(); + let _value_a = cell.get_component_mut::(entity).unwrap(); + let _value_b = cell.get_component::(entity).unwrap(); + } + + #[test] + fn world_cell_component_ref_and_ref() { + let mut world = World::default(); + let entity = world.spawn().insert(A(1)).id(); + let cell = world.cell(); + let _value_a = cell.get_component::(entity).unwrap(); + let _value_b = cell.get_component::(entity).unwrap(); + } + + #[test] + fn world_cell_component_mut_and_ref_different_entities() { + let mut world = World::default(); + let entity_1 = world.spawn().insert(A(1)).id(); + let entity_2 = world.spawn().insert(A(1)).id(); + let cell = world.cell(); + let _value_1 = cell.get_component_mut::(entity_1).unwrap(); + let _value_2 = cell.get_component::(entity_2).unwrap(); + } }