lyra-engine/lyra-ecs/src/query/borrow.rs

use std::{marker::PhantomData, any::TypeId, ptr::NonNull, cell::{Ref, RefCell, RefMut}};

use crate::{world::World, ComponentColumn};

use super::{Fetch, Query, AsQuery, DefaultQuery};

/// Fetcher for borrowing components from archetypes.
pub struct FetchBorrow<'a, T> {
    col: &'a ComponentColumn,
    size: usize,
    _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 {
        let ptr = self.col.borrow_ptr();
        Ref::map(ptr, |ptr| {
            let ptr = NonNull::new_unchecked(ptr.as_ptr()
                .add(entity.0 as usize * self.size))
                .cast();
            &*ptr.as_ptr()
        })
    }
}

/// 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<T> {
    type_id: TypeId,
    _phantom: PhantomData<T>
}

impl<T: 'static> QueryBorrow<T> {
    pub fn new() -> Self {
        Self {
            type_id: TypeId::of::<T>(),
            _phantom: PhantomData,
        }
    }
}

impl<T> Query for QueryBorrow<T>
where
    T: 'static
{
    type Item<'a> = Ref<'a, T>;

    type Fetch<'a> = FetchBorrow<'a, T>;

    fn can_visit_archetype(&self, archetype: &crate::archetype::Archetype) -> bool {
        archetype.columns.iter().any(|c| c.info.type_id == self.type_id)
    }

    unsafe fn fetch<'a>(&self, _world: &'a World, _arch_id: crate::archetype::ArchetypeId, archetype: &'a crate::archetype::Archetype) -> Self::Fetch<'a> {
        let col = archetype.columns.iter().find(|c| c.info.type_id == self.type_id)
            .expect("You ignored 'can_visit_archetype'!");

        FetchBorrow {
            col,
            size: col.info.layout.size(),
            _phantom: PhantomData,
        }
    }
}

impl<T: 'static> AsQuery for QueryBorrow<T> {
    type Query = Self;
}

impl<T: 'static> DefaultQuery for QueryBorrow<T> {
    fn default_query() -> Self {
        QueryBorrow::<T>::new()
    }
}

impl<T: 'static> AsQuery for &T {
    type Query = QueryBorrow<T>;
}

impl<T: 'static> DefaultQuery for &T {
    fn default_query() -> QueryBorrow<T> {
        QueryBorrow::<T>::new()
    }
}

/// A fetcher for mutably borrowing components from archetypes.
pub struct FetchBorrowMut<'a, T> {
    col: &'a ComponentColumn,
    size: usize,
    _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 ptr = self.col.borrow_mut_ptr();
        RefMut::map(ptr, |ptr| {
            let ptr = NonNull::new_unchecked(ptr.as_ptr()
                .add(entity.0 as usize * self.size))
                .cast();
            &mut *ptr.as_ptr()
        })
    }
}

/// 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<T> {
    type_id: TypeId,
    _phantom: PhantomData<T>
}

impl<T: 'static> QueryBorrowMut<T> {
    pub fn new() -> Self {
        Self {
            type_id: TypeId::of::<T>(),
            _phantom: PhantomData,
        }
    }
}

impl<T> Query for QueryBorrowMut<T>
where
    T: 'static
{
    type Item<'a> = RefMut<'a, T>;

    type Fetch<'a> = FetchBorrowMut<'a, T>;

    fn can_visit_archetype(&self, archetype: &crate::archetype::Archetype) -> bool {
        archetype.columns.iter().any(|c| c.info.type_id == self.type_id)
    }

    unsafe fn fetch<'a>(&self, _world: &'a World, _arch_id: crate::archetype::ArchetypeId, archetype: &'a crate::archetype::Archetype) -> Self::Fetch<'a> {
        let col = archetype.columns.iter().find(|c| c.info.type_id == self.type_id)
            .expect("You ignored 'can_visit_archetype'!");

        FetchBorrowMut {
            col,
            size: col.info.layout.size(),
            _phantom: PhantomData,
        }
    }
}

impl<T: 'static> AsQuery for QueryBorrowMut<T> {
    type Query = Self;
}

impl<T: 'static> DefaultQuery for QueryBorrowMut<T> {
    fn default_query() -> Self {
        QueryBorrowMut::<T>::new()
    }
}

impl<T: 'static> AsQuery for &mut T {
    type Query = QueryBorrowMut<T>;
}

impl<T: 'static> DefaultQuery for &mut T {
    fn default_query() -> QueryBorrowMut<T> {
        QueryBorrowMut::<T>::new()
    }
}

#[cfg(test)]
mod tests {
    use std::{any::TypeId, mem::size_of, marker::PhantomData};

    use crate::{world::{World, Entity, EntityId}, archetype::{Archetype, ArchetypeId}, query::View, tests::Vec2, bundle::Bundle, Fetch};

    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::<Vec2> {
            type_id: TypeId::of::<Vec2>(),
            _phantom: std::marker::PhantomData,
        };
        let archetypes: Vec<&Archetype> = world.archetypes.values().collect();
        let v = View::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::<Vec2> {
            type_id: TypeId::of::<Vec2>(),
            _phantom: std::marker::PhantomData,
        };
        let archetypes: Vec<&Archetype> = world.archetypes.values().collect();
        let v = View::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::<Vec2> {
            type_id: TypeId::of::<Vec2>(),
            _phantom: std::marker::PhantomData,
        };
        let archetypes: Vec<&Archetype> = world.archetypes.values().collect();
        let v = View::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(),));
        }

        let col = a.columns.iter().find(|c| c.info.type_id == TypeId::of::<Vec2>()).unwrap();

        let mut bmut = FetchBorrowMut::<Vec2> {
            col,
            size: size_of::<Vec2>(),
            _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());
    }
}