path: root/ecs/src/relationship.rs

use std::any::{type_name, Any};
use std::marker::PhantomData;

use crate::component::storage::Storage as ComponentStorage;
use crate::component::{
    Component,
    FromOptional as ComponentFromOptional,
    Id as ComponentId,
};
use crate::entity::Uid as EntityUid;
use crate::lock::ReadGuard;
use crate::system::{ComponentRefMut, Input as SystemInput};
use crate::type_name::TypeName;
use crate::World;

/// A relationship to one or more targets.
#[derive(Debug)]
pub struct Relationship<Kind, ComponentT: Component>
{
    entity_uid: SingleOrMultiple<EntityUid>,
    _pd: PhantomData<(Kind, ComponentT)>,
}

impl<Kind, ComponentT> Relationship<Kind, ComponentT>
where
    ComponentT: Component,
{
    /// Creates a new `Relationship` with a single target.
    pub fn new(entity_uid: EntityUid) -> Self
    {
        Self {
            entity_uid: SingleOrMultiple::Single(entity_uid),
            _pd: PhantomData,
        }
    }

    /// Creates a new `Relationship` with multiple targets.
    pub fn new_multiple(entity_uids: impl IntoIterator<Item = EntityUid>) -> Self
    {
        Self {
            entity_uid: SingleOrMultiple::Multiple(entity_uids.into_iter().collect()),
            _pd: PhantomData,
        }
    }
}

impl<Kind, ComponentT> Component for Relationship<Kind, ComponentT>
where
    Kind: 'static,
    ComponentT: Component,
{
    type Component = Self;
    type RefMut<'component> = Relation<'component, Kind, ComponentT>;

    fn id(&self) -> ComponentId
    {
        ComponentId::of::<Self>()
    }

    fn as_any_mut(&mut self) -> &mut dyn Any
    {
        self
    }

    fn as_any(&self) -> &dyn Any
    {
        self
    }
}

impl<Kind, ComponentT> SystemInput for Relationship<Kind, ComponentT>
where
    Kind: 'static,
    ComponentT: Component,
{
}

impl<Kind, ComponentT: Component> TypeName for Relationship<Kind, ComponentT>
where
    ComponentT: Component,
{
    fn type_name(&self) -> &'static str
    {
        type_name::<Self>()
    }
}

pub struct Relation<'rel_comp, Kind, ComponentT>
where
    Kind: 'static,
    ComponentT: Component,
{
    component_storage_lock: ReadGuard<'static, ComponentStorage>,
    relationship_comp: ComponentRefMut<'rel_comp, Relationship<Kind, ComponentT>>,
}

impl<'rel_comp, Kind, ComponentT> ComponentFromOptional<'rel_comp>
    for Relation<'rel_comp, Kind, ComponentT>
where
    ComponentT: Component,
{
    fn from_optional_component(
        optional_component: Option<
            crate::lock::WriteGuard<'rel_comp, Box<dyn Component>>,
        >,
        world: &'rel_comp World,
    ) -> Self
    {
        let relationship_comp =
            ComponentRefMut::<Relationship<Kind, ComponentT>>::from_optional_component(
                optional_component,
                world,
            );

        let component_storage_lock = world
            .data
            .component_storage
            .read_nonblock()
            .expect("Failed to aquire read-only component storage lock");

        Self {
            relationship_comp,
            // SAFETY: The component lock is not used for longer than the original
            // lifetime
            component_storage_lock: unsafe { component_storage_lock.upgrade_lifetime() },
        }
    }
}

impl<'rel_comp, Kind, ComponentT> Relation<'rel_comp, Kind, ComponentT>
where
    ComponentT: Component,
{
    /// Returns the component of the target at the specified index.
    pub fn get(&self, index: usize) -> Option<ComponentRefMut<'_, ComponentT>>
    {
        let target = self.get_target(index)?;

        let archetype = self.component_storage_lock.get_entity_archetype(*target)?;

        let entity = archetype
            .get_entity(*target)
            .expect("Target entity is gone from archetype");

        let component_index =
            archetype.get_index_for_component(&ComponentId::of::<ComponentT>())?;

        let component = ComponentRefMut::new(
            entity
                .get_component(component_index)?
                .component
                .write_nonblock()
                .unwrap_or_else(|_| {
                    panic!(
                        "Failed to aquire read-write lock of component {}",
                        type_name::<ComponentT>()
                    )
                }),
        );

        Some(component)
    }

    /// Returns a reference to the target at the specified index.
    pub fn get_target(&self, index: usize) -> Option<&EntityUid>
    {
        match &self.relationship_comp.entity_uid {
            SingleOrMultiple::Single(entity_uid) if index == 0 => Some(entity_uid),
            SingleOrMultiple::Multiple(entity_uids) => entity_uids.get(index),
            _ => None,
        }
    }

    /// Returns a mutable reference to the target at the specified index.
    pub fn get_target_mut(&mut self, index: usize) -> Option<&mut EntityUid>
    {
        match &mut self.relationship_comp.entity_uid {
            SingleOrMultiple::Single(entity_uid) if index == 0 => Some(entity_uid),
            SingleOrMultiple::Multiple(entity_uids) => entity_uids.get_mut(index),
            _ => None,
        }
    }

    /// Adds a target to the relationship.
    pub fn add_target(&mut self, entity_uid: EntityUid)
    {
        match &mut self.relationship_comp.entity_uid {
            SingleOrMultiple::Single(prev_entity_uid) => {
                self.relationship_comp.entity_uid =
                    SingleOrMultiple::Multiple(vec![*prev_entity_uid, entity_uid]);
            }
            SingleOrMultiple::Multiple(entity_uids) => entity_uids.push(entity_uid),
        }
    }

    /// Removes a target to the relationship, returning it.
    pub fn remove_target(&mut self, index: usize) -> Option<EntityUid>
    {
        match &mut self.relationship_comp.entity_uid {
            SingleOrMultiple::Single(entity_uid) => {
                let prev_entity_uid = *entity_uid;

                self.relationship_comp.entity_uid =
                    SingleOrMultiple::Multiple(Vec::new());

                Some(prev_entity_uid)
            }
            SingleOrMultiple::Multiple(entity_uids) => {
                if index >= entity_uids.len() {
                    return None;
                }

                Some(entity_uids.remove(index))
            }
        }
    }

    pub fn target_count(&self) -> usize
    {
        match &self.relationship_comp.entity_uid {
            SingleOrMultiple::Single(_) => 1,
            SingleOrMultiple::Multiple(entity_uids) => entity_uids.len(),
        }
    }

    /// Returns a iterator of the components of the targets of this relationship.
    pub fn iter(&self) -> TargetComponentIter<'_, 'rel_comp, Kind, ComponentT>
    {
        TargetComponentIter { relation: self, index: 0 }
    }
}

/// Iterator of the components of the targets of a relationship.
pub struct TargetComponentIter<'relationship, 'rel_comp, Kind, ComponentT>
where
    Kind: 'static,
    ComponentT: Component,
{
    relation: &'relationship Relation<'rel_comp, Kind, ComponentT>,
    index: usize,
}

impl<'relationship, 'rel_comp, Kind, ComponentT> Iterator
    for TargetComponentIter<'relationship, 'rel_comp, Kind, ComponentT>
where
    'relationship: 'rel_comp,
    Kind: 'static,
    ComponentT: Component,
{
    type Item = ComponentRefMut<'rel_comp, ComponentT>;

    fn next(&mut self) -> Option<Self::Item>
    {
        let index = self.index;

        self.index += 1;

        self.relation.get(index)
    }
}

#[derive(Debug)]
enum SingleOrMultiple<Value>
{
    Single(Value),
    Multiple(Vec<Value>),
}