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

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

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

/// Fetcher for borrowing components from archetypes.
pub struct FetchBorrow<'a, T> {
    ptr: Option<Ref<'a, NonNull<u8>>>,
    size: usize,
    _phantom: PhantomData<&'a T>
}

impl<'a, T> Fetch<'a> for FetchBorrow<'a, T>
where
    T: 'a,
{
    type Item = &'a T;

    fn dangling() -> Self {
        FetchBorrow {
            ptr: None,
            size: 0,
            _phantom: PhantomData::default()
        }
    }

    unsafe fn get_item(&mut self, entity: crate::world::ArchetypeEntityId) -> Self::Item {
        let ptr = self.ptr.as_ref().unwrap();
        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:
/// ```rust
/// for ts in world.view::<&T>() {
///     println!("Got an &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> = &'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, _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'!");
        let col_data = col.borrow_ptr();

        FetchBorrow {
            ptr: Some(col_data),
            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> {
    ptr: Option<RefMut<'a, NonNull<u8>>>,
    size: usize,
    _phantom: PhantomData<&'a T>
}

impl<'a, T> Fetch<'a> for FetchBorrowMut<'a, T>
where
    T: 'a,
{
    type Item = &'a mut T;

    fn dangling() -> Self {
        FetchBorrowMut {
            ptr: None,
            size: 0,
            _phantom: PhantomData::default()
        }
    }

    unsafe fn get_item(&mut self, entity: crate::world::ArchetypeEntityId) -> Self::Item {
        let ptr = self.ptr.as_ref().unwrap();
        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:
/// ```rust
/// for ts in world.view::<&mut T>() {
///     println!("Got an &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> = &'a mut 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, _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'!");
        let col_data = col.borrow_mut_ptr();

        FetchBorrowMut {
            ptr: Some(col_data),
            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};

    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(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(borrow, archetypes);

        let mut orig = vec![];

        for 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(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 bmut = FetchBorrowMut::<Vec2> {
            ptr: Some(col.borrow_mut_ptr()),
            size: size_of::<Vec2>(),
            _phantom: PhantomData,
        };

        assert!(col.try_borrow_mut_ptr().is_err());
        drop(bmut);
        assert!(col.try_borrow_mut_ptr().is_ok());
    }
}