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

use std::any::type_name;
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::lock::Lock;
use crate::type_name::TypeName;
use crate::util::Sortable;
use crate::EntityComponent;

#[derive(Debug, Default)]
pub struct Storage
{
    archetypes: Vec<Archetype>,
    archetype_lookup: HashMap<ArchetypeId, Vec<usize>>,
    pending_archetype_lookup_entries: Vec<Vec<ComponentMetadata>>,
}

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

        self.archetype_lookup
            .get(&ArchetypeId::from_components_metadata(components_metadata))
            .map_or_else(ArchetypeRefIter::new_empty, |archetypes_indices| {
                ArchetypeRefIter {
                    inner: archetypes_indices.iter(),
                    archetypes: &self.archetypes,
                }
            })
    }

    #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
    pub fn push_entity(
        &mut self,
        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 archetype_indices =
            self.archetype_lookup
                .entry(archetype_id)
                .or_insert_with(|| {
                    self.archetypes.push(Archetype::new(
                        components.iter().map(|component| component.id()),
                    ));

                    vec![self.archetypes.len() - 1]
                });

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

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

        let archetype = self
            .archetypes
            .get_mut(
                archetype_indices
                    .first()
                    .copied()
                    .expect("No archetype index found"),
            )
            .expect("Archetype is gone");

        let entity_uid = EntityUid::new_unique();

        archetype.entities.push(ArchetypeEntity {
            uid: entity_uid,
            components: components
                .into_iter()
                .map(|component| EntityComponent {
                    id: component.id(),
                    name: component.type_name(),
                    component: Lock::new(component),
                })
                .collect(),
        });

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

        (archetype_id, entity_uid)
    }

    pub fn add_archetype_lookup_entry(
        &mut self,
        mut components_metadata: impl IntoIterator<Item = ComponentMetadata>
            + Sortable<Item = ComponentMetadata>,
    )
    {
        components_metadata.sort_by_key_b(|component_metadata| component_metadata.id);

        self.pending_archetype_lookup_entries
            .push(components_metadata.into_iter().collect());
    }

    pub fn make_archetype_lookup_entries(&mut self)
    {
        for pending_entry in self.pending_archetype_lookup_entries.drain(..) {
            let pending_entry_ids_set = pending_entry
                .iter()
                .filter_map(|component_metadata| {
                    if component_metadata.is_optional == ComponentIsOptional::Yes {
                        return None;
                    }

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

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

                    None
                });

            let archetype_id =
                ArchetypeId::from_components_metadata(pending_entry.into_iter());

            if self.archetype_lookup.contains_key(&archetype_id) {
                continue;
            }

            self.archetype_lookup
                .insert(archetype_id, matching_archetype_indices.collect());
        }
    }
}

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

#[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()
    }
}

#[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: SliceIter<'component_storage, usize>,
    archetypes: &'component_storage [Archetype],
}

impl<'component_storage> ArchetypeRefIter<'component_storage>
{
    fn new_empty() -> Self
    {
        Self { inner: [].iter(), archetypes: &[] }
    }
}

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()
    }
}

#[cfg(test)]
mod tests
{
    use std::collections::HashMap;

    use ecs_macros::Component;

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

    #[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(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.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 lookup = component_storage
            .archetype_lookup
            .get(&ArchetypeId::from_components_metadata(components_metadata))
            .expect("Expected entry in archetype lookup map");

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

        assert_eq!(*first_archetype_index, 0);
    }

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

        let entity_uid_a = EntityUid::new_unique();
        let entity_uid_b = EntityUid::new_unique();
        let entity_uid_c = EntityUid::new_unique();

        component_storage.archetypes.push(Archetype {
            component_ids: HashMap::from([
                (ComponentId::of::<IronBoots>(), 0),
                (ComponentId::of::<HealthPotion>(), 1),
                (ComponentId::of::<Hookshot>(), 2),
            ]),
            entity_lookup: HashMap::from([
                (entity_uid_a, 0),
                (entity_uid_b, 1),
                (entity_uid_c, 2),
            ]),
            entities: vec![
                ArchetypeEntity {
                    uid: entity_uid_a,
                    components: vec![EntityComponent {
                        id: ComponentId::of::<IronBoots>(),
                        name: "Iron boots",
                        component: Lock::new(Box::new(IronBoots)),
                    }],
                },
                ArchetypeEntity {
                    uid: entity_uid_b,
                    components: vec![EntityComponent {
                        id: ComponentId::of::<HealthPotion>(),
                        name: "Health potion",
                        component: Lock::new(Box::new(HealthPotion {
                            _hp_restoration: 20,
                        })),
                    }],
                },
                ArchetypeEntity {
                    uid: entity_uid_c,
                    components: vec![EntityComponent {
                        id: ComponentId::of::<Hookshot>(),
                        name: "Hookshot",
                        component: Lock::new(Box::new(Hookshot { _range: 67 })),
                    }],
                },
            ],
        });

        let entity_uid_d = EntityUid::new_unique();
        let entity_uid_e = EntityUid::new_unique();
        let entity_uid_f = EntityUid::new_unique();
        let entity_uid_g = EntityUid::new_unique();

        component_storage.archetypes.push(Archetype {
            component_ids: HashMap::from([
                (ComponentId::of::<DekuNut>(), 0),
                (ComponentId::of::<IronBoots>(), 1),
                (ComponentId::of::<Bow>(), 2),
                (ComponentId::of::<Hookshot>(), 3),
            ]),
            entity_lookup: HashMap::from([
                (entity_uid_d, 0),
                (entity_uid_e, 1),
                (entity_uid_f, 2),
                (entity_uid_g, 3),
            ]),
            entities: vec![
                ArchetypeEntity {
                    uid: entity_uid_d,
                    components: vec![EntityComponent {
                        id: ComponentId::of::<DekuNut>(),
                        name: "Deku nut",
                        component: Lock::new(Box::new(DekuNut { _throwing_damage: 5 })),
                    }],
                },
                ArchetypeEntity {
                    uid: entity_uid_e,
                    components: vec![EntityComponent {
                        id: ComponentId::of::<IronBoots>(),
                        name: "Iron boots",
                        component: Lock::new(Box::new(IronBoots)),
                    }],
                },
                ArchetypeEntity {
                    uid: entity_uid_f,
                    components: vec![EntityComponent {
                        id: ComponentId::of::<Bow>(),
                        name: "Bow",
                        component: Lock::new(Box::new(Bow { _damage: 20 })),
                    }],
                },
                ArchetypeEntity {
                    uid: entity_uid_g,
                    components: vec![EntityComponent {
                        id: ComponentId::of::<Hookshot>(),
                        name: "Hookshot",
                        component: Lock::new(Box::new(Hookshot { _range: 67 })),
                    }],
                },
            ],
        });

        component_storage.add_archetype_lookup_entry([
            ComponentMetadata {
                id: ComponentId::of::<IronBoots>(),
                is_optional: ComponentIsOptional::No,
            },
            ComponentMetadata {
                id: ComponentId::of::<Hookshot>(),
                is_optional: ComponentIsOptional::No,
            },
        ]);

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

        component_storage.make_archetype_lookup_entries();

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

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

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

        let archetypes = component_storage
            .archetype_lookup
            .get(&ArchetypeId::from_components_metadata(comps_metadata))
            .expect(concat!(
                "Expected a archetype for IronBoots & Hookshot to be found in the ",
                "archetype lookup map"
            ));

        assert_eq!(archetypes.len(), 2);
    }
}