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

use std::any::type_name;
use std::collections::{HashMap, HashSet};
use std::hash::{DefaultHasher, Hash, Hasher};
use std::slice::Iter as SliceIter;

use crate::component::{Component, Id as ComponentId, IsOptional as ComponentIsOptional};
use crate::lock::Lock;
use crate::type_name::TypeName;
use crate::EntityComponent;

#[derive(Debug, Default)]
pub struct ComponentStorage
{
    archetypes: Vec<Archetype>,
    archetype_lookup: HashMap<ArchetypeComponentsHash, Vec<usize>>,
    pending_archetype_lookup_entries: Vec<Vec<ComponentId>>,
}

#[derive(Debug)]
pub struct ArchetypeRefIter<'component_storage>
{
    inner: SliceIter<'component_storage, usize>,
    archetypes: &'component_storage Vec<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"),
        )
    }
}

impl ComponentStorage
{
    pub fn find_entities(
        &self,
        component_ids: &[(ComponentId, ComponentIsOptional)],
    ) -> Option<ArchetypeRefIter<'_>>
    {
        let ids = component_ids
            .iter()
            .filter_map(|(component_id, is_optional)| {
                if *is_optional == ComponentIsOptional::Yes {
                    return None;
                }

                Some(*component_id)
            });

        self.archetype_lookup
            .get(&ArchetypeComponentsHash::new(ids))
            .map(|archetypes_indices| ArchetypeRefIter {
                inner: archetypes_indices.iter(),
                archetypes: &self.archetypes,
            })
    }

    #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
    pub fn push_entity(&mut self, components: Vec<Box<dyn Component>>)
    {
        #[cfg(feature = "debug")]
        tracing::debug!(
            "Pushing entity with components: ({})",
            components
                .iter()
                .map(|component| component.type_name())
                .collect::<Vec<_>>()
                .join(", ")
        );

        let archetype_indices = self
            .archetype_lookup
            .entry(ArchetypeComponentsHash::new(
                components
                    .iter()
                    .filter(|component| !component.is_optional())
                    .map(|component| component.id()),
            ))
            .or_insert_with(|| {
                self.archetypes.push(Archetype::default());

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

        if archetype_indices.is_empty() {
            self.archetypes.push(Archetype::default());

            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
            .component_ids
            .extend(components.iter().map(|component| component.id()));

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

    pub fn add_archetype_lookup_entry(&mut self, component_ids: &[ComponentId])
    {
        self.pending_archetype_lookup_entries
            .push(component_ids.to_vec());
    }

    pub fn make_archetype_lookup_entries(&mut self)
    {
        for pending_entry in &self.pending_archetype_lookup_entries {
            let components_set: HashSet<_> = pending_entry
                .into_iter()
                .map(|component_id| *component_id)
                .collect();

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

                    None
                });

            let archetype_indices = self
                .archetype_lookup
                .entry(ArchetypeComponentsHash::new(
                    pending_entry.into_iter().copied().clone(),
                ))
                .or_default();

            archetype_indices.extend(matching_archetype_indices);
        }
    }
}

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

#[derive(Debug, Default)]
pub struct Archetype
{
    component_ids: HashSet<ComponentId>,
    pub components: Vec<Vec<EntityComponent>>,
}

#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
struct ArchetypeComponentsHash
{
    hash: u64,
}

impl ArchetypeComponentsHash
{
    fn new(component_ids: impl IntoIterator<Item = ComponentId>) -> Self
    {
        let mut hasher = DefaultHasher::new();

        for component_id in component_ids {
            component_id.hash(&mut hasher);
        }

        let hash = hasher.finish();

        Self { hash }
    }
}

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

    use ecs_macros::Component;

    use super::{Archetype, ArchetypeComponentsHash, ComponentStorage};
    use crate::component::Id as ComponentId;
    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 = ComponentStorage::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(&ArchetypeComponentsHash::new([
                ComponentId::of::<HealthPotion>(),
                ComponentId::of::<Hookshot>(),
            ]))
            .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 = ComponentStorage::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(&[
            ComponentId::of::<IronBoots>(),
            ComponentId::of::<Hookshot>(),
        ]);

        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(&ArchetypeComponentsHash::new([
                ComponentId::of::<IronBoots>(),
                ComponentId::of::<Hookshot>(),
            ]))
            .expect(concat!(
                "Expected a archetype for IronBoots & Hookshot to be found in the ",
                "archetype lookup map"
            ));

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