path: root/ecs/src/component/storage/archetype.rs

use std::any::Any;
use std::array::IntoIter as ArrayIntoIter;
use std::hash::{DefaultHasher, Hash, Hasher};
use std::iter::{Enumerate, Filter, Map, RepeatN, Zip};
use std::option::IntoIter as OptionIntoIter;
use std::slice::Iter as SliceIter;

use hashbrown::HashMap;

use crate::lock::Lock;
use crate::uid::{Kind as UidKind, Uid};
use crate::util::{Either, HashMapExt};

#[derive(Debug)]
pub struct Archetype
{
    id: Id,
    entities: Vec<Entity>,
    entity_index_lookup: HashMap<Uid, usize>,
    component_index_lookup: HashMap<Uid, usize>,
    component_ids: Vec<Uid>,
}

impl Archetype
{
    pub fn new(id: Id, component_ids: impl AsRef<[Uid]>) -> Self
    {
        Self {
            id,
            entities: Vec::new(),
            entity_index_lookup: HashMap::new(),
            component_index_lookup: component_ids
                .as_ref()
                .iter()
                .enumerate()
                .map(|(index, id)| (*id, index))
                .collect(),
            component_ids: component_ids.as_ref().to_vec(),
        }
    }

    pub fn id(&self) -> Id
    {
        self.id
    }

    pub fn is_superset(&self, other: &Self) -> bool
    {
        self.component_index_lookup
            .keys_is_superset(&other.component_index_lookup)
    }

    pub fn is_subset(&self, other: &Self) -> bool
    {
        self.component_index_lookup
            .keys_is_subset(&other.component_index_lookup)
    }

    pub fn get_entity_by_id(&self, entity_uid: Uid) -> Option<&Entity>
    {
        let index = *self.entity_index_lookup.get(&entity_uid)?;

        Some(self.entities.get(index).unwrap_or_else(|| {
            panic!(
                "In invalid state! Index of entity with ID {entity_uid:?} is out of bounds"
            );
        }))
    }

    pub fn push_entity(&mut self, entity: Entity)
    {
        self.entity_index_lookup
            .insert(entity.uid, self.entities.len());

        self.entities.push(entity);
    }

    pub fn remove_entity(&mut self, entity_uid: Uid) -> Option<Entity>
    {
        //debug_assert_eq!(entity_uid.kind(), UidKind::Entity);

        let entity_index = self.entity_index_lookup.remove(&entity_uid)?;

        if self.entities.len() == 1 {
            return Some(self.entities.remove(entity_index));
        }

        let last_entity_uid = self
            .entities
            .last()
            .expect(concat!(
                "Invalid state. No entities in archetype but entry was ",
                "removed successfully from entity index lookup"
            ))
            .uid;

        // By using swap_remove, no memory reallocation occurs and only one index in the
        // entity lookup needs to be updated
        let removed_entity = self.entities.swap_remove(entity_index);

        self.entity_index_lookup
            .insert(last_entity_uid, entity_index);

        Some(removed_entity)
    }

    pub fn entities(&self) -> EntityIter<'_>
    {
        EntityIter { iter: self.entities.iter() }
    }

    pub fn entity_cnt(&self) -> usize
    {
        self.entities.len()
    }

    pub fn component_cnt(&self) -> usize
    {
        self.component_index_lookup.len()
    }

    pub fn get_matching_component_indices(
        &self,
        component_id: Uid,
    ) -> MatchingComponentIter
    {
        assert!(
            component_id.kind() == UidKind::Component
                || component_id.kind() == UidKind::Pair
        );

        if component_id.kind() == UidKind::Pair
            && component_id.target_component() == Uid::wildcard()
        {
            return MatchingComponentIter {
                inner: Either::A(
                    self.component_ids
                        .iter()
                        .enumerate()
                        .zip(std::iter::repeat_n(component_id, self.component_ids.len()))
                        .filter(
                            (|((_, other_comp_id), component_id)| {
                                other_comp_id.kind() == UidKind::Pair
                                    && other_comp_id.has_same_relation_as(*component_id)
                            })
                                as MatchingComponentIterFilterFn,
                        )
                        .map(|((index, other_comp_id), _)| (*other_comp_id, index)),
                ),
            };
        }

        MatchingComponentIter {
            inner: Either::B(
                [component_id]
                    .into_iter()
                    .zip(self.get_index_for_component(component_id)),
            ),
        }
    }

    pub fn get_index_for_component(&self, component_id: Uid) -> Option<usize>
    {
        assert!(
            component_id.kind() == UidKind::Component
                || (component_id.kind() == UidKind::Pair
                    && component_id.target_component() != Uid::wildcard())
        );

        self.component_index_lookup.get(&component_id).copied()
    }

    pub fn component_ids_unsorted(&self) -> impl Iterator<Item = Uid> + '_
    {
        self.component_index_lookup.keys().copied()
    }

    pub fn component_ids_sorted(&self) -> impl Iterator<Item = Uid> + '_
    {
        self.component_ids.iter().copied()
    }

    pub fn contains_matching_component(&self, component_id: Uid) -> bool
    {
        let component_id_kind = component_id.kind();

        debug_assert!(
            component_id_kind == UidKind::Component || component_id_kind == UidKind::Pair
        );

        if component_id.kind() == UidKind::Pair
            && component_id.target_component() == Uid::wildcard()
        {
            return self.component_ids.iter().any(|other_comp_id| {
                other_comp_id.kind() == UidKind::Pair
                    && other_comp_id.has_same_relation_as(component_id)
            });
        }

        self.contains_component_with_exact_id(component_id)
    }

    pub fn contains_component_with_exact_id(&self, component_id: Uid) -> bool
    {
        let component_id_kind = component_id.kind();

        debug_assert!(
            component_id_kind == UidKind::Component
                || (component_id_kind == UidKind::Pair
                    && component_id.target_component() != Uid::wildcard())
        );

        self.component_index_lookup.contains_key(&component_id)
    }
}

type MatchingComponentIterFilterFn = fn(&((usize, &Uid), Uid)) -> bool;

type MatchingComponentIterMapFn = fn(((usize, &Uid), Uid)) -> (Uid, usize);

type InnerMatchingComponentIterA<'archetype> = Map<
    Filter<
        Zip<Enumerate<SliceIter<'archetype, Uid>>, RepeatN<Uid>>,
        MatchingComponentIterFilterFn,
    >,
    MatchingComponentIterMapFn,
>;

type InnerMatchingComponentIterB = Zip<ArrayIntoIter<Uid, 1>, OptionIntoIter<usize>>;

#[derive(Debug)]
pub struct MatchingComponentIter<'archetype>
{
    inner: Either<InnerMatchingComponentIterA<'archetype>, InnerMatchingComponentIterB>,
}

impl Iterator for MatchingComponentIter<'_>
{
    type Item = (Uid, usize);

    fn next(&mut self) -> Option<Self::Item>
    {
        self.inner.next()
    }
}

#[derive(Debug)]
pub struct EntityIter<'archetype>
{
    iter: SliceIter<'archetype, Entity>,
}

impl<'archetype> Iterator for EntityIter<'archetype>
{
    type Item = &'archetype Entity;

    fn next(&mut self) -> Option<Self::Item>
    {
        self.iter.next()
    }
}

#[derive(Debug)]
pub struct Entity
{
    uid: Uid,
    components: Vec<EntityComponent>,
}

impl Entity
{
    pub fn new(uid: Uid, components: impl IntoIterator<Item = EntityComponent>) -> Self
    {
        Self {
            uid,
            components: components.into_iter().collect(),
        }
    }

    pub fn uid(&self) -> Uid
    {
        self.uid
    }

    pub fn components(&self) -> &[EntityComponent]
    {
        &self.components
    }

    pub fn remove_component(&mut self, component_id: Uid, archetype: &Archetype)
    {
        let index = archetype
            .get_index_for_component(component_id)
            .expect("Archetype should contain component");

        self.components.remove(index);
    }

    pub fn insert_component(
        &mut self,
        component_id: Uid,
        component: EntityComponent,
        archetype: &Archetype,
    )
    {
        let index = archetype
            .get_index_for_component(component_id)
            .expect("Archetype should contain component");

        self.components.insert(index, component);
    }
}

#[derive(Debug)]
pub struct EntityComponent
{
    id: Uid,
    component: Lock<Box<dyn Any>>,
}

impl EntityComponent
{
    pub fn new(
        component: Box<dyn Any>,
        component_id: Uid,
        component_name: &'static str,
    ) -> Self
    {
        Self {
            id: component_id,
            component: Lock::new(component, component_name),
        }
    }

    #[allow(dead_code)]
    pub fn id(&self) -> Uid
    {
        self.id
    }

    pub fn component(&self) -> &Lock<Box<dyn Any>>
    {
        &self.component
    }
}

/// Archetype ID.
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct Id
{
    hash: u64,
}

impl Id
{
    pub fn new_empty() -> Self
    {
        Self { hash: 0 }
    }

    pub fn new<'a>(component_ids: impl IntoIterator<Item = &'a Uid>) -> Self
    {
        let mut hasher = DefaultHasher::new();

        let mut prev_component_id: Option<Uid> = None;

        let mut component_id_iter = component_ids.into_iter().peekable();

        if component_id_iter.peek().is_none() {
            return Self::new_empty();
        }

        for comp_id in component_id_iter {
            if prev_component_id.is_some_and(|prev_comp_id| *comp_id < prev_comp_id) {
                panic!(
                    "Cannot create archetype ID from a unsorted component metadata list"
                );
            }

            prev_component_id = Some(*comp_id);

            comp_id.hash(&mut hasher);
        }

        Self { hash: hasher.finish() }
    }
}