use std::{marker::PhantomData, ptr::NonNull, cell::{Ref, RefMut}}; use crate::{World, ComponentColumn, DynTypeId, Tick, Component}; use super::{Fetch, Query, AsQuery}; /// Fetcher for borrowing components from archetypes. pub struct FetchBorrow<'a, T> { col: &'a ComponentColumn, _phantom: PhantomData<&'a T> } impl<'a, T> Fetch<'a> for FetchBorrow<'a, T> where T: 'a, { type Item = Ref<'a, T>; fn dangling() -> Self { unreachable!() } unsafe fn get_item(&mut self, entity: crate::world::ArchetypeEntityId) -> Self::Item { self.col.get(entity.0 as _) } } /// A Query for borrowing components from archetypes. /// /// Since [`AsQuery`] is implemented for `&T`, you can use this query like this: /// ```nobuild /// for ts in world.view::<&T>() { /// println!("Got a &T!"); /// } /// ``` pub struct QueryBorrow { type_id: DynTypeId, _phantom: PhantomData } impl Default for QueryBorrow { fn default() -> Self { Self { type_id: DynTypeId::of::(), _phantom: PhantomData, } } } impl Copy for QueryBorrow {} impl Clone for QueryBorrow { fn clone(&self) -> Self { *self } } impl QueryBorrow { pub fn new() -> Self { Self::default() } } impl Query for QueryBorrow where T: 'static { type Item<'a> = Ref<'a, T>; type Fetch<'a> = FetchBorrow<'a, T>; fn new() -> Self { QueryBorrow::::new() } fn can_visit_archetype(&self, archetype: &crate::archetype::Archetype) -> bool { archetype.has_column(self.type_id) } unsafe fn fetch<'a>(&self, _world: &'a World, archetype: &'a crate::archetype::Archetype, tick: crate::Tick) -> Self::Fetch<'a> { let _ = tick; let col = archetype.get_column(self.type_id) .expect("You ignored 'can_visit_archetype'!"); FetchBorrow { col, _phantom: PhantomData, } } } impl AsQuery for QueryBorrow { type Query = Self; } impl AsQuery for &T { type Query = QueryBorrow; } /// A fetcher for mutably borrowing components from archetypes. pub struct FetchBorrowMut<'a, T> { col: NonNull, tick: Tick, _phantom: PhantomData<&'a T> } impl<'a, T> Fetch<'a> for FetchBorrowMut<'a, T> where T: 'a, { type Item = RefMut<'a, T>; fn dangling() -> Self { unreachable!() } unsafe fn get_item(&mut self, entity: crate::world::ArchetypeEntityId) -> Self::Item { let col = unsafe { self.col.as_mut() }; col.get_mut(entity.0 as _, &self.tick) } } /// A Query for mutably borrowing components from archetypes. /// /// Since [`AsQuery`] is implemented for `&mut T`, you can use this query like this: /// ```nobuild /// for ts in world.view::<&mut T>() { /// println!("Got an &mut T!"); /// } /// ``` pub struct QueryBorrowMut { type_id: DynTypeId, _phantom: PhantomData } impl Default for QueryBorrowMut { fn default() -> Self { Self { type_id: DynTypeId::of::(), _phantom: PhantomData, } } } impl Copy for QueryBorrowMut {} impl Clone for QueryBorrowMut { fn clone(&self) -> Self { *self } } impl QueryBorrowMut { pub fn new() -> Self { Self::default() } } impl Query for QueryBorrowMut where T: 'static { type Item<'a> = RefMut<'a, T>; type Fetch<'a> = FetchBorrowMut<'a, T>; const MUTATES: bool = true; fn new() -> Self { QueryBorrowMut::::new() } fn can_visit_archetype(&self, archetype: &crate::archetype::Archetype) -> bool { archetype.has_column(self.type_id) } unsafe fn fetch<'a>(&self, _world: &'a World, archetype: &'a crate::archetype::Archetype, tick: crate::Tick) -> Self::Fetch<'a> { let col = archetype.get_column(self.type_id) .expect("You ignored 'can_visit_archetype'!"); // TODO: find a way to get the component column mutable with a borrowed archetype so its tick can be updated. // the fetcher needs to tick the entities tick in the archetype let col = NonNull::from(col); FetchBorrowMut { col, tick, _phantom: PhantomData, } } } impl AsQuery for QueryBorrowMut { type Query = Self; } impl AsQuery for &mut T { type Query = QueryBorrowMut; } #[cfg(test)] mod tests { use std::{marker::PhantomData, ptr::NonNull}; use crate::{archetype::{Archetype, ArchetypeId}, bundle::Bundle, query::{Fetch, ViewState}, tests::Vec2, World, DynTypeId, Entity, EntityId, Tick}; use super::{QueryBorrow, QueryBorrowMut, FetchBorrowMut}; /// Creates a world with two entities, one at Vec(10, 50) and /// the other at Vec2(25, 30). fn prepare_world() -> World { let mut world = World::new(); world.spawn((Vec2::new(10.0, 50.0),)); world.spawn((Vec2::new(25.0, 30.0),)); world } #[test] fn borrow_query() { let world = prepare_world(); let borrow = QueryBorrow:: { type_id: DynTypeId::of::(), _phantom: std::marker::PhantomData, }; let archetypes: Vec<&Archetype> = world.archetypes.values().collect(); let v = ViewState::, ()>::new(&world, borrow, (), archetypes); for e in v.into_iter() { println!("Found entity at {:?}", e); } } #[test] fn borrow_mut_query() { let world = prepare_world(); let borrow = QueryBorrowMut:: { type_id: DynTypeId::of::(), _phantom: std::marker::PhantomData, }; let archetypes: Vec<&Archetype> = world.archetypes.values().collect(); let v = ViewState::, ()>::new(&world, borrow, (), archetypes); let mut orig = vec![]; for mut v in v.into_iter() { orig.push(v.clone()); v.x += 10.0; v.y += 10.0; } // Now make sure the changes were actually made let borrow = QueryBorrow:: { type_id: DynTypeId::of::(), _phantom: std::marker::PhantomData, }; let archetypes: Vec<&Archetype> = world.archetypes.values().collect(); let v = ViewState::, ()>::new(&world, borrow, (), archetypes); for (new, orig) in v.into_iter().zip(orig.iter()) { assert!(new.x - orig.x == 10.0); assert!(new.y - orig.y == 10.0); } println!("They were modified!"); } #[test] fn only_one_mutable_borrow() { let info = (Vec2::new(0.0, 0.0),).info(); let mut a = Archetype::from_bundle_info(ArchetypeId(0), info); for i in 0..10 { a.add_entity(Entity { id: EntityId(i), generation: 0, }, (Vec2::rand(),), &Tick::default()); } let col = a.get_column(DynTypeId::of::()).unwrap(); let mut bmut = FetchBorrowMut:: { col: NonNull::from(col), tick: Tick::default(), _phantom: PhantomData, }; let item = unsafe { bmut.get_item(crate::ArchetypeEntityId(0)) }; // ensure only one mutable borrow can be done at a time assert!(col.try_borrow_mut_ptr().is_err()); drop(item); assert!(col.try_borrow_mut_ptr().is_ok()); } }