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

use std::any::type_name;
use std::borrow::Borrow;
use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::slice::Iter as SliceIter;

use crate::archetype::Id as ArchetypeId;
use crate::component::{
    Component,
    Id as ComponentId,
    IsOptional as ComponentIsOptional,
    Metadata as ComponentMetadata,
};
use crate::entity::Uid as EntityUid;
use crate::type_name::TypeName;
use crate::util::{RefCellRefMap, Sortable};
use crate::EntityComponent;

#[derive(Debug, Default)]
pub struct Storage
{
    archetypes: Vec<Archetype>,
    archetype_lookup: RefCell<HashMap<ArchetypeId, ArchetypeLookupEntry>>,
    entity_archetype_lookup: HashMap<EntityUid, ArchetypeId>,
}

impl Storage
{
    pub fn find_entities<CompMetadataList>(
        &self,
        mut components_metadata: CompMetadataList,
    ) -> ArchetypeRefIter<'_>
    where
        CompMetadataList: Sortable<Item = ComponentMetadata>,
        CompMetadataList: Borrow<[ComponentMetadata]>,
    {
        components_metadata.sort_by_key_b(|component_metadata| component_metadata.id);

        let archetype_id = ArchetypeId::from_components_metadata(
            components_metadata.borrow().into_iter().cloned(),
        );

        // This looks stupid but the borrow checker complains otherwise
        if self.archetype_lookup.borrow().contains_key(&archetype_id) {
            let archetype_lookup = self.archetype_lookup.borrow();

            return ArchetypeRefIter {
                inner: RefCellRefMap::new(archetype_lookup, |archetype_lookup| {
                    archetype_lookup
                        .get(&archetype_id)
                        .unwrap()
                        .archetype_indices
                        .iter()
                }),
                archetypes: &self.archetypes,
            };
        }

        let comp_ids_set = create_non_opt_component_id_set(components_metadata.borrow());

        let matching_archetype_indices = self
            .archetypes
            .iter()
            .enumerate()
            .filter_map(|(index, archetype)| {
                if archetype.component_ids_is_superset(&comp_ids_set) {
                    return Some(index);
                }

                None
            })
            .collect();

        self.archetype_lookup.borrow_mut().insert(
            archetype_id,
            ArchetypeLookupEntry {
                component_ids: comp_ids_set,
                archetype_indices: matching_archetype_indices,
            },
        );

        let archetype_lookup = self.archetype_lookup.borrow();

        ArchetypeRefIter {
            inner: RefCellRefMap::new(archetype_lookup, |archetype_lookup| {
                archetype_lookup
                    .get(&archetype_id)
                    .unwrap()
                    .archetype_indices
                    .iter()
            }),
            archetypes: &self.archetypes,
        }
    }

    #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
    pub fn push_entity(
        &mut self,
        entity_uid: EntityUid,
        mut components: Vec<Box<dyn Component>>,
    ) -> (ArchetypeId, EntityUid)
    {
        components.sort_by_key(|component| component.id());

        #[cfg(feature = "debug")]
        tracing::debug!(
            "Pushing entity with components: ({})",
            components
                .iter()
                .map(|component| component.type_name())
                .collect::<Vec<_>>()
                .join(", ")
        );

        let archetype_id = ArchetypeId::from_components_metadata(
            components
                .iter()
                .map(|component| ComponentMetadata::of(&**component)),
        );

        let comp_ids_set = create_non_opt_component_id_set(
            components
                .iter()
                .map(|component| ComponentMetadata::of(&**component)),
        );

        let archetype_index =
            self.get_or_create_archetype(archetype_id, &comp_ids_set, &components);

        self.populate_matching_archetype_lookup_entries(&comp_ids_set, archetype_index);

        let archetype = self
            .archetypes
            .get_mut(archetype_index)
            .expect("Archetype is gone");

        archetype.push_entity(entity_uid, components);

        archetype
            .entity_lookup
            .insert(entity_uid, archetype.entities.len() - 1);

        self.entity_archetype_lookup
            .insert(entity_uid, archetype_id);

        (archetype_id, entity_uid)
    }

    pub fn add_components_to_entity(
        &mut self,
        entity_uid: EntityUid,
        components: impl IntoIterator<Item = Box<dyn Component>>,
    ) -> Option<()>
    {
        let archetype_id = self.entity_archetype_lookup.get(&entity_uid)?;

        let archetype_index =
            self.find_archetype_index_with_entity(*archetype_id, entity_uid)?;

        let archetype = self.archetypes.get_mut(archetype_index)?;

        let entity = archetype.take_entity(entity_uid)?;

        self.push_entity(
            entity_uid,
            entity
                .components
                .into_iter()
                .map(|component| component.component.into_inner())
                .chain(components.into_iter().map(|component| component.into()))
                .collect(),
        );

        Some(())
    }

    fn populate_matching_archetype_lookup_entries(
        &mut self,
        comp_ids_set: &HashSet<ComponentId>,
        archetype_index: usize,
    )
    {
        let mut archetype_lookup = self.archetype_lookup.borrow_mut();

        for (_, lookup_entry) in archetype_lookup.iter_mut() {
            if &lookup_entry.component_ids == comp_ids_set {
                continue;
            }

            if lookup_entry.component_ids.is_subset(&comp_ids_set) {
                lookup_entry.archetype_indices.push(archetype_index);
            }
        }
    }

    fn get_or_create_archetype(
        &mut self,
        archetype_id: ArchetypeId,
        comp_ids_set: &HashSet<ComponentId>,
        components: &[Box<dyn Component>],
    ) -> usize
    {
        let mut archetype_lookup = self.archetype_lookup.borrow_mut();

        let lookup_entry = archetype_lookup.entry(archetype_id).or_insert_with(|| {
            ArchetypeLookupEntry {
                component_ids: comp_ids_set.clone(),
                archetype_indices: Vec::with_capacity(1),
            }
        });

        if lookup_entry.archetype_indices.is_empty() {
            self.archetypes.push(Archetype::new(
                components.iter().map(|component| component.id()),
            ));

            lookup_entry
                .archetype_indices
                .push(self.archetypes.len() - 1);
        }

        // SAFETY: Above, we push a archetype index if archetype_indices is empty so this
        // cannot fail
        unsafe { *lookup_entry.archetype_indices.first().unwrap_unchecked() }
    }

    fn find_archetype_index_with_entity(
        &self,
        archetype_id: ArchetypeId,
        entity_uid: EntityUid,
    ) -> Option<usize>
    {
        let archetype_lookup = self.archetype_lookup.borrow_mut();

        let archetype_lookup_entry = archetype_lookup.get(&archetype_id)?;

        // TODO: Also have a hashmap for precise archetype ID -> archetype index lookup.
        // This way is slow
        archetype_lookup_entry
            .archetype_indices
            .iter()
            .find(|archetype_index| {
                let Some(archetype) = self.archetypes.get(**archetype_index) else {
                    return false;
                };

                archetype.has_entity(entity_uid)
            })
            .map(|archetype_index| *archetype_index)
    }
}

impl TypeName for Storage
{
    fn type_name(&self) -> &'static str
    {
        type_name::<Self>()
    }
}

#[derive(Debug)]
struct ArchetypeLookupEntry
{
    component_ids: HashSet<ComponentId>,
    archetype_indices: Vec<usize>,
}

#[derive(Debug)]
pub struct Archetype
{
    component_ids: HashMap<ComponentId, usize>,
    entity_lookup: HashMap<EntityUid, usize>,
    entities: Vec<ArchetypeEntity>,
}

impl Archetype
{
    fn new(component_ids: impl IntoIterator<Item = ComponentId>) -> Self
    {
        Self {
            component_ids: component_ids
                .into_iter()
                .enumerate()
                .map(|(index, component_id)| (component_id, index))
                .collect(),
            entity_lookup: HashMap::new(),
            entities: Vec::new(),
        }
    }

    pub fn component_ids_is_superset(
        &self,
        other_component_ids: &HashSet<ComponentId>,
    ) -> bool
    {
        if other_component_ids.len() <= self.component_ids.len() {
            other_component_ids
                .iter()
                .all(|v| self.component_ids.contains_key(v))
        } else {
            false
        }
    }

    pub fn get_entity(&self, entity_uid: EntityUid) -> Option<&ArchetypeEntity>
    {
        let entity_index = *self.entity_lookup.get(&entity_uid)?;

        self.entities.get(entity_index)
    }

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

    pub fn get_index_for_component(&self, component_id: &ComponentId) -> Option<usize>
    {
        self.component_ids.get(component_id).copied()
    }

    fn push_entity(
        &mut self,
        entity_uid: EntityUid,
        components: impl IntoIterator<Item = Box<dyn Component>>,
    )
    {
        self.entities.push(ArchetypeEntity {
            uid: entity_uid,
            components: components
                .into_iter()
                .map(|component| component.into())
                .collect(),
        });
    }

    pub fn take_entity(&mut self, entity_uid: EntityUid) -> Option<ArchetypeEntity>
    {
        let entity_index = self.entity_lookup.remove(&entity_uid)?;

        let entity = self.entities.remove(entity_index);

        for index in self.entity_lookup.values_mut() {
            if *index > entity_index {
                *index -= 1;
            }
        }

        Some(entity)
    }

    fn has_entity(&self, entity_uid: EntityUid) -> bool
    {
        self.entity_lookup.contains_key(&entity_uid)
    }
}

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

impl ArchetypeEntity
{
    pub fn uid(&self) -> EntityUid
    {
        self.uid
    }

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

    pub fn get_component(&self, index: usize) -> Option<&EntityComponent>
    {
        self.components.get(index)
    }
}

#[derive(Debug)]
pub struct ArchetypeRefIter<'component_storage>
{
    inner: RefCellRefMap<
        'component_storage,
        HashMap<ArchetypeId, ArchetypeLookupEntry>,
        SliceIter<'component_storage, usize>,
    >,
    archetypes: &'component_storage [Archetype],
}

impl<'component_storage> Iterator for ArchetypeRefIter<'component_storage>
{
    type Item = &'component_storage Archetype;

    fn next(&mut self) -> Option<Self::Item>
    {
        let archetype_index = *self.inner.next()?;

        Some(
            self.archetypes
                .get(archetype_index)
                .expect("Archetype index in archetype lookup entry was not found"),
        )
    }
}

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

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

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

fn create_non_opt_component_id_set<'a, Item>(
    component_metadata_iter: impl IntoIterator<Item = Item>,
) -> HashSet<ComponentId>
where
    Item: Borrow<ComponentMetadata>,
{
    component_metadata_iter
        .into_iter()
        .filter_map(|item| {
            let component_metadata = item.borrow();

            if component_metadata.is_optional == ComponentIsOptional::Yes {
                return None;
            }

            Some(component_metadata.id)
        })
        .collect::<HashSet<_>>()
}

#[cfg(test)]
mod tests
{

    use ecs_macros::Component;

    use super::Storage;
    use crate::archetype::Id as ArchetypeId;
    use crate::component::{
        Id as ComponentId,
        IsOptional as ComponentIsOptional,
        Metadata as ComponentMetadata,
    };
    use crate::entity::Uid as EntityUid;
    use crate::{self as ecs};

    #[derive(Debug, Component)]
    struct HealthPotion
    {
        _hp_restoration: u32,
    }

    #[derive(Debug, Component)]
    struct Hookshot
    {
        _range: u32,
    }

    #[derive(Debug, Component)]
    struct DekuNut
    {
        _throwing_damage: u32,
    }

    #[derive(Debug, Component)]
    struct Bow
    {
        _damage: u32,
    }

    #[derive(Debug, Component)]
    struct IronBoots;

    #[test]
    fn push_entity_works()
    {
        let mut component_storage = Storage::default();

        component_storage.push_entity(
            EntityUid::new_unique(),
            vec![
                Box::new(HealthPotion { _hp_restoration: 12 }),
                Box::new(Hookshot { _range: 50 }),
            ],
        );

        assert_eq!(component_storage.archetypes.len(), 1);

        let archetype = component_storage
            .archetypes
            .first()
            .expect("Expected a archetype in archetypes Vec");

        assert_eq!(archetype.component_ids.len(), 2);

        // One entity
        assert_eq!(archetype.entities.len(), 1);

        let entity_components = archetype
            .entities
            .first()
            .expect("Expected a entity in archetype");

        assert_eq!(entity_components.components.len(), 2);

        assert_eq!(component_storage.archetype_lookup.borrow().len(), 1);

        let mut components_metadata = [
            ComponentMetadata {
                id: ComponentId::of::<HealthPotion>(),
                is_optional: ComponentIsOptional::No,
            },
            ComponentMetadata {
                id: ComponentId::of::<Hookshot>(),
                is_optional: ComponentIsOptional::No,
            },
        ];

        components_metadata.sort_by_key(|comp_metadata| comp_metadata.id);

        let archetype_lookup = component_storage.archetype_lookup.borrow();

        let lookup_entry = archetype_lookup
            .get(&ArchetypeId::from_components_metadata(components_metadata))
            .expect("Expected entry in archetype lookup map");

        let first_archetype_index = lookup_entry
            .archetype_indices
            .first()
            .expect("Expected archetype lookup to contain a archetype reference");

        assert_eq!(*first_archetype_index, 0);
    }
}