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::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,
        components_metadata: impl IntoIterator<Item = ComponentMetadata>,
    ) -> ArchetypeRefIter<'_>
    {
        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");

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

        let entity_uid = EntityUid::new_unique();

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

        (archetype_id, entity_uid)
    }

    pub fn add_archetype_lookup_entry(
        &mut self,
        components_metadata: impl IntoIterator<Item = ComponentMetadata>,
    )
    {
        self.pending_archetype_lookup_entries
            .push(components_metadata.into_iter().collect());
    }

    pub fn make_archetype_lookup_entries(&mut self)
    {
        // TODO: Re-doing the whole archetype lookup table is dumb and slow. Only re-do
        // the archetype lookup entries that need to be changed
        self.archetype_lookup.clear();

        for pending_entry in self.pending_archetype_lookup_entries.drain(..) {
            let ids = pending_entry.iter().filter_map(|component_metadata| {
                if component_metadata.is_optional == ComponentIsOptional::Yes {
                    return None;
                }

                Some(component_metadata.id)
            });

            let ids_set: HashSet<_> = ids.collect();

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

                    None
                });

            let archetype_indices = self
                .archetype_lookup
                .entry(ArchetypeId::from_components_metadata(
                    pending_entry.into_iter(),
                ))
                .or_default();

            archetype_indices.extend(matching_archetype_indices);
        }
    }
}

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>,
    pub components: Vec<Vec<EntityComponent>>,
}

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(),
            components: 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<&[EntityComponent]>
    {
        let entity_index = *self.entity_lookup.get(&entity_uid)?;

        Some(&self.components.get(entity_index)?)
    }

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

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

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

    use ecs_macros::Component;

    use super::{Archetype, Storage};
    use crate::archetype::Id as ArchetypeId;
    use crate::component::{
        Id as ComponentId,
        IsOptional as ComponentIsOptional,
        Metadata as ComponentMetadata,
    };
    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.components.len(), 1);

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

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

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

        let lookup = component_storage
            .archetype_lookup
            .get(&ArchetypeId::from_components_metadata([
                ComponentMetadata {
                    id: ComponentId::of::<HealthPotion>(),
                    is_optional: ComponentIsOptional::No,
                },
                ComponentMetadata {
                    id: ComponentId::of::<Hookshot>(),
                    is_optional: ComponentIsOptional::No,
                },
            ]))
            .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();

        component_storage.archetypes.push(Archetype {
            component_ids: HashSet::from([
                ComponentId::of::<IronBoots>(),
                ComponentId::of::<HealthPotion>(),
                ComponentId::of::<Hookshot>(),
            ]),
            components: vec![
                vec![EntityComponent {
                    id: ComponentId::of::<IronBoots>(),
                    name: "Iron boots",
                    component: Lock::new(Box::new(IronBoots)),
                }],
                vec![EntityComponent {
                    id: ComponentId::of::<HealthPotion>(),
                    name: "Health potion",
                    component: Lock::new(Box::new(HealthPotion { _hp_restoration: 20 })),
                }],
                vec![EntityComponent {
                    id: ComponentId::of::<Hookshot>(),
                    name: "Hookshot",
                    component: Lock::new(Box::new(Hookshot { _range: 67 })),
                }],
            ],
        });

        component_storage.archetypes.push(Archetype {
            component_ids: HashSet::from([
                ComponentId::of::<DekuNut>(),
                ComponentId::of::<IronBoots>(),
                ComponentId::of::<Bow>(),
                ComponentId::of::<Hookshot>(),
            ]),
            components: vec![
                vec![EntityComponent {
                    id: ComponentId::of::<DekuNut>(),
                    name: "Deku nut",
                    component: Lock::new(Box::new(DekuNut { _throwing_damage: 5 })),
                }],
                vec![EntityComponent {
                    id: ComponentId::of::<IronBoots>(),
                    name: "Iron boots",
                    component: Lock::new(Box::new(IronBoots)),
                }],
                vec![EntityComponent {
                    id: ComponentId::of::<Bow>(),
                    name: "Bow",
                    component: Lock::new(Box::new(Bow { _damage: 20 })),
                }],
                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 archetypes = component_storage
            .archetype_lookup
            .get(&ArchetypeId::from_components_metadata([
                ComponentMetadata {
                    id: ComponentId::of::<IronBoots>(),
                    is_optional: ComponentIsOptional::No,
                },
                ComponentMetadata {
                    id: ComponentId::of::<Hookshot>(),
                    is_optional: ComponentIsOptional::No,
                },
            ]))
            .expect(concat!(
                "Expected a archetype for IronBoots & Hookshot to be found in the ",
                "archetype lookup map"
            ));

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