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

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

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

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

/// 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: DynTypeId,
    _phantom: PhantomData<T>
}

impl<T: 'static> Default for QueryBorrow<T> {
    fn default() -> Self {
        Self {
            type_id: DynTypeId::of::<T>(),
            _phantom: PhantomData,
        }
    }
}

impl<T> Copy for QueryBorrow<T> {}

impl<T> Clone for QueryBorrow<T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T: 'static> QueryBorrow<T> {
    pub fn new() -> Self {
        Self::default()
    }
}

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

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

    fn new() -> Self {
        QueryBorrow::<T>::new()
    }

    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, archetype: &'a crate::archetype::Archetype, tick: crate::Tick) -> Self::Fetch<'a> {
        let _ = tick;
        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> AsQuery for &T {
    type Query = QueryBorrow<T>;
}

/// A fetcher for mutably borrowing components from archetypes.
pub struct FetchBorrowMut<'a, T> {
    col: NonNull<ComponentColumn>,
    size: usize,
    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.entity_ticks[entity.0 as usize] = self.tick;
        let ptr = 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: DynTypeId,
    _phantom: PhantomData<T>
}

impl<T: 'static> Default for QueryBorrowMut<T> {
    fn default() -> Self {
        Self {
            type_id: DynTypeId::of::<T>(),
            _phantom: PhantomData,
        }
    }
}

impl<T> Copy for QueryBorrowMut<T> {}

impl<T> Clone for QueryBorrowMut<T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T: 'static> QueryBorrowMut<T> {
    pub fn new() -> Self {
        Self::default()
    }
}

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

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

    const MUTATES: bool = true;

    fn new() -> Self {
        QueryBorrowMut::<T>::new()
    }

    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, archetype: &'a crate::archetype::Archetype, tick: crate::Tick) -> Self::Fetch<'a> {
        let col = archetype.columns.iter().find(|c| c.info.type_id == self.type_id)
            .expect("You ignored 'can_visit_archetype'!");
        let layout_size = col.info.layout.size;
        let col = NonNull::from(col);

        // 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

        FetchBorrowMut {
            col,
            size: layout_size,
            tick,
            _phantom: PhantomData,
        }
    }
}

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

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

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

    use crate::{world::{World, Entity, EntityId}, archetype::{Archetype, ArchetypeId}, query::{View, Fetch}, tests::Vec2, bundle::Bundle, DynTypeId, 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::<Vec2> {
            type_id: DynTypeId::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: DynTypeId::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: DynTypeId::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(),), &Tick::default());
        }

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

        let mut bmut = FetchBorrowMut::<Vec2> {
            col: NonNull::from(col),
            tick: Tick::default(),
            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());
    }
}