path: root/ecs/src/query.rs

use std::any::{Any, TypeId};
use std::collections::HashSet;
use std::marker::PhantomData;
use std::sync::{Arc, Weak};

use crate::component::Sequence as ComponentSequence;
use crate::lock::{Lock, ReadGuard};
use crate::system::{
    NoInitParamFlag as NoInitSystemParamFlag,
    Param as SystemParam,
    System,
};
use crate::tuple::FilterExclude as TupleFilterExclude;
use crate::{ComponentStorage, WorldData};

#[derive(Debug)]
pub struct Query<'world, Comps>
where
    Comps: ComponentSequence,
{
    component_storage: ReadGuard<'world, ComponentStorage>,
    component_storage_lock: Weak<Lock<ComponentStorage>>,
    comps_pd: PhantomData<Comps>,
}

impl<'world, Comps> Query<'world, Comps>
where
    Comps: ComponentSequence,
{
    #[must_use]
    pub fn iter<'this>(&'this self) -> ComponentIter<'world, Comps>
    where
        'this: 'world,
    {
        ComponentIter {
            component_storage: &self.component_storage,
            current_entity_index: 0,
            component_type_ids: Comps::type_ids(),
            comps_pd: PhantomData,
        }
    }

    /// Returns a weak reference query to the same components.
    #[must_use]
    pub fn to_weak_ref(&self) -> WeakRef<Comps>
    {
        WeakRef {
            component_storage: self.component_storage_lock.clone(),
            comps_pd: PhantomData,
        }
    }

    pub(crate) fn new(component_storage: &'world Arc<Lock<ComponentStorage>>) -> Self
    {
        Self {
            component_storage: component_storage
                .read_nonblock()
                .expect("Failed to acquire read-only component storage lock"),
            component_storage_lock: Arc::downgrade(component_storage),
            comps_pd: PhantomData,
        }
    }
}

impl<'world, Comps> IntoIterator for &'world Query<'world, Comps>
where
    Comps: ComponentSequence,
{
    type IntoIter = ComponentIter<'world, Comps>;
    type Item = Comps::Refs<'world>;

    fn into_iter(self) -> Self::IntoIter
    {
        self.iter()
    }
}

unsafe impl<'world, Comps> SystemParam<'world> for Query<'world, Comps>
where
    Comps: ComponentSequence,
{
    type Flags = NoInitSystemParamFlag;
    type Input = TupleFilterExclude;

    fn initialize<SystemImpl>(
        _system: &mut impl System<'world, SystemImpl>,
        _input: Self::Input,
    )
    {
    }

    fn new<SystemImpl>(
        _system: &'world impl System<'world, SystemImpl>,
        world_data: &'world WorldData,
    ) -> Self
    {
        Self::new(&world_data.component_storage)
    }

    fn is_compatible<Other: SystemParam<'world>>() -> bool
    {
        let other_comparable = Other::get_comparable();

        let Some(other_query_component_ids) =
            other_comparable.downcast_ref::<QueryComponentIds>()
        else {
            return true;
        };

        !other_query_component_ids.contains_component_in::<Comps>()
    }

    fn get_comparable() -> Box<dyn Any>
    {
        Box::new(QueryComponentIds {
            component_type_ids: Comps::type_ids(),
        })
    }
}

/// A entity query containing a weak reference to the world.
#[derive(Debug)]
pub struct WeakRef<Comps>
where
    Comps: ComponentSequence,
{
    component_storage: Weak<Lock<ComponentStorage>>,
    comps_pd: PhantomData<Comps>,
}

impl<Comps> WeakRef<Comps>
where
    Comps: ComponentSequence,
{
    /// Returns a struct which can be used to retrieve a [`Query`].
    ///
    /// Returns [`None`] if the [`World`] has been dropped.
    #[must_use]
    pub fn access(&self) -> Option<Ref<'_, Comps>>
    {
        Some(Ref {
            component_storage: self.component_storage.upgrade()?,
            _pd: PhantomData,
        })
    }
}

impl<Comps> Clone for WeakRef<Comps>
where
    Comps: ComponentSequence,
{
    fn clone(&self) -> Self
    {
        Self {
            component_storage: self.component_storage.clone(),
            comps_pd: PhantomData,
        }
    }
}

/// Intermediate between [`Query`] and [`WeakRefQuery`]. Contains a strong reference to
/// the world which is not allowed direct access to.
#[derive(Debug, Clone)]
pub struct Ref<'weak_ref, Comps>
where
    Comps: ComponentSequence,
{
    component_storage: Arc<Lock<ComponentStorage>>,
    _pd: PhantomData<&'weak_ref Comps>,
}

impl<'weak_ref, Comps> Ref<'weak_ref, Comps>
where
    Comps: ComponentSequence,
{
    #[must_use]
    pub fn to_query(&self) -> Query<'_, Comps>
    {
        Query::new(&self.component_storage)
    }
}

pub struct ComponentIter<'world, Comps>
{
    component_storage: &'world ComponentStorage,
    current_entity_index: usize,
    component_type_ids: Vec<TypeId>,
    comps_pd: PhantomData<Comps>,
}

impl<'world, Comps> Iterator for ComponentIter<'world, Comps>
where
    Comps: ComponentSequence + 'world,
{
    type Item = Comps::Refs<'world>;

    fn next(&mut self) -> Option<Self::Item>
    {
        let (matching_entity_index, matching_entity) =
            self.component_storage.find_entity_with_components(
                self.current_entity_index,
                &self.component_type_ids,
            )?;

        self.current_entity_index = matching_entity_index + 1;

        Some(Comps::from_components(
            matching_entity
                .components
                .iter()
                .map(|component| &component.component),
        ))
    }
}

#[derive(Debug)]
struct QueryComponentIds
{
    component_type_ids: Vec<TypeId>,
}

impl QueryComponentIds
{
    fn contains_component_in<OtherComps>(&self) -> bool
    where
        OtherComps: ComponentSequence,
    {
        let other_component_type_ids = OtherComps::type_ids()
            .into_iter()
            .collect::<HashSet<TypeId>>();

        // TODO: Make this a bit smarter. Queries with a same component can be compatible
        // if one of the queries have a component the other one does not have
        self.component_type_ids
            .iter()
            .any(|component_type_id| other_component_type_ids.contains(component_type_id))
    }
}