hyperion/

component_storage.rs

use crate::entity::Entity;
use std::any::Any;
use std::collections::HashMap;

/// Marker trait for types that can be used as components in the ECS.
pub trait Component: 'static {}

/// Storage for components of type `T`. Maintains dense storage and entity-index mappings.
///
/// # Fields
/// - `components`: Dense array of components of type `T`.
/// - `entity_to_index`: Map from `Entity` to its dense array index.
/// - `index_to_entity`: Reverse map from dense array index to owning `Entity`.
///
/// # Safety
/// - Storage remains dense by using swap-remove on deletion; index maps are updated
///   accordingly to prevent dangling indices. Component references returned by
///   queries are valid for the borrow duration and follow Rust's aliasing rules.
pub struct ComponentStorage<T: Component> {
    components: Vec<T>,
    entity_to_index: HashMap<Entity, usize>,
    index_to_entity: Vec<Entity>,
}

impl<T: Component> ComponentStorage<T> {
    /// Creates a new, empty component storage.
    ///
    /// # Returns
    /// - `ComponentStorage<T>`: An empty storage ready to hold components of type `T`.
    pub fn new() -> Self {
        Self {
            components: Vec::new(),
            entity_to_index: HashMap::new(),
            index_to_entity: Vec::new(),
        }
    }

    /// Given a reference to a component in this storage, return the owning `Entity`.
    ///
    /// This performs an O(n) scan comparing pointer identity within the dense
    /// component array to locate the index, then maps that index back to the entity.
    ///
    /// # Parameters
    /// - `component`: A reference to a component stored in this storage.
    ///
    /// # Returns
    /// - `Option<Entity>`: The entity that owns the component, if found.
    pub fn get_entity_for_component(&self, component: &T) -> Option<Entity> {
        for (idx, comp_ref) in self.components.iter().enumerate() {
            if std::ptr::eq(comp_ref, component) {
                return Some(self.index_to_entity[idx]);
            }
        }
        None
    }

    /// Inserts or replaces a component for an entity.
    ///
    /// # Parameters
    /// - `entity`: The entity that will own the component.
    /// - `component`: The component instance to insert (or replace if one already exists).
    pub fn insert(&mut self, entity: Entity, component: T) {
        if let Some(&index) = self.entity_to_index.get(&entity) {
            // Replace existing component at index.
            self.components[index] = component;
        } else {
            // Insert new component at end and update mappings.
            let index = self.components.len();
            self.components.push(component);
            self.entity_to_index.insert(entity, index);
            self.index_to_entity.push(entity);
        }
    }

    /// Removes the component for an entity, returning it if present.
    /// Uses swap-remove to keep storage dense and updates mappings accordingly.
    ///
    /// # Parameters
    /// - `entity`: The entity whose component should be removed.
    ///
    /// # Returns
    /// - `Option<T>`: The removed component if it existed.
    pub fn remove(&mut self, entity: Entity) -> Option<T> {
        self.entity_to_index.remove(&entity).map(|index| {
            let last_index = self.components.len() - 1;

            // If not the last component, swap to keep storage dense
            if index != last_index {
                self.components.swap(index, last_index);
                self.index_to_entity.swap(index, last_index);

                // Update moved entity's mapping to new index
                let swapped_entity = self.index_to_entity[index];
                self.entity_to_index.insert(swapped_entity, index);
            }

            self.index_to_entity.pop();
            self.components.pop().unwrap()
        })
    }

    /// Gets an immutable reference to a component for an entity, if it exists.
    ///
    /// # Parameters
    /// - `entity`: The entity whose component to retrieve.
    ///
    /// # Returns
    /// - `Option<&T>`: A reference to the component if present.
    pub fn get(&self, entity: Entity) -> Option<&T> {
        self.entity_to_index
            .get(&entity)
            .map(|&index| &self.components[index])
    }

    /// Gets a mutable reference to a component for an entity, if it exists.
    ///
    /// # Parameters
    /// - `entity`: The entity whose component to retrieve mutably.
    ///
    /// # Returns
    /// - `Option<&mut T>`: A mutable reference to the component if present.
    pub fn get_mut(&mut self, entity: Entity) -> Option<&mut T> {
        self.entity_to_index
            .get(&entity)
            .map(|&index| &mut self.components[index])
    }

    /// Returns all components as a vector of immutable references.
    ///
    /// # Returns
    /// - `Vec<&T>`: All components currently stored.
    pub fn get_all(&self) -> Vec<&T> {
        self.components.iter().collect()
    }

    /// Returns all components as a vector of mutable references.
    ///
    /// # Returns
    /// - `Vec<&mut T>`: All components currently stored, mutably borrowed.
    pub fn get_all_mut(&mut self) -> Vec<&mut T> {
        self.components.iter_mut().collect()
    }

    /// Returns an iterator over immutable references to components.
    ///
    /// # Returns
    /// - `impl Iterator<Item = &T>`: Iterator over all components.
    pub fn iter(&self) -> impl Iterator<Item = &T> {
        self.components.iter()
    }

    /// Returns an iterator over mutable references to components.
    ///
    /// # Returns
    /// - `impl Iterator<Item = &mut T>`: Iterator over all components mutably.
    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> {
        self.components.iter_mut()
    }

    /// Returns an iterator over all entities and immutable references to their components.
    ///
    /// # Returns
    /// - `impl Iterator<Item = (Entity, &T)>`: Iterator of `(Entity, &T)` pairs.
    pub fn iter_with_entities(&self) -> impl Iterator<Item = (Entity, &T)> {
        self.index_to_entity
            .iter()
            .copied()
            .zip(self.components.iter())
    }

    /// Returns an iterator over all entities and mutable references to their components.
    ///
    /// # Returns
    /// - `impl Iterator<Item = (Entity, &mut T)>`: Iterator of `(Entity, &mut T)` pairs.
    pub fn iter_with_entities_mut(&mut self) -> impl Iterator<Item = (Entity, &mut T)> {
        let entities = self.index_to_entity.clone(); // Clone to avoid borrow checker issues
        entities.into_iter().zip(self.components.iter_mut())
    }

    /// Returns true if this storage contains a component for the specified entity.
    ///
    /// # Parameters
    /// - `entity`: The entity to check.
    ///
    /// # Returns
    /// - `bool`: `true` if a component for `entity` exists in this storage.
    pub fn has(&self, entity: Entity) -> bool {
        self.entity_to_index.contains_key(&entity)
    }
}

/// A type-erased interface over concrete `ComponentStorage<T>` implementations.
///
/// Enables the `World` to store heterogeneous component storages behind a single
/// trait object and perform operations such as removing an entity's component
/// without knowing the concrete `T` at compile time.
///
/// # Safety
/// - Downcasting is only exposed via `as_any`/`as_any_mut` and is guarded by
///   external `TypeId` checks in `World` before downcast is attempted.
/// - `remove_entity` must maintain dense storage invariants of the underlying
///   storage (e.g., swap-remove updates index maps) when implemented.
pub trait ComponentStorageTrait {
    /// Removes the component associated with `entity` from this storage, if present.
    ///
    /// # Parameters
    /// - `entity`: The entity whose component should be removed from this storage.
    ///
    /// # Effects
    /// - If a component exists for `entity`, it is removed. Implementations typically
    ///   use swap-remove to keep storage dense and update internal mappings accordingly.
    fn remove_entity(&mut self, entity: Entity);

    /// Returns a type-erased immutable view of this storage for downcasting.
    ///
    /// # Returns
    /// - `&dyn Any`: A reference that can be downcast using `Any::downcast_ref` to
    ///   access the concrete `ComponentStorage<T>` when the caller knows `T`.
    ///
    /// # Safety
    /// - Callers must only downcast to the correct type, typically ensured via
    ///   `TypeId` checks performed by `World` helpers.
    fn as_any(&self) -> &dyn Any;

    /// Returns a type-erased mutable view of this storage for downcasting.
    ///
    /// # Returns
    /// - `&mut dyn Any`: A mutable reference that can be downcast using
    ///   `Any::downcast_mut` to access the concrete `ComponentStorage<T>`.
    ///
    /// # Safety
    /// - Callers must ensure the downcast target type matches the actual storage type.
    ///   The `World` ensures correctness via `TypeId` checks before downcasting.
    fn as_any_mut(&mut self) -> &mut dyn Any;
}

impl<T: Component> ComponentStorageTrait for ComponentStorage<T> {
    /// Removes the component for the given `entity` from this concrete storage.
    ///
    /// # Parameters
    /// - `entity`: The entity whose component should be removed.
    ///
    /// # Effects
    /// - If a component exists for `entity`, it is removed using the storage's
    ///   internal removal logic (swap-remove), maintaining dense storage and
    ///   updating index mappings accordingly.
    fn remove_entity(&mut self, entity: Entity) {
        let _ = self.remove(entity);
    }

    /// Returns an immutable, type-erased view of this concrete storage for downcasting.
    ///
    /// # Returns
    /// - `&dyn Any`: A reference that can be downcast via `Any::downcast_ref`
    ///   to `ComponentStorage<T>` by callers that know the concrete `T`.
    ///
    /// # Safety
    /// - Callers must only attempt to downcast to the correct concrete type.
    ///   The `World` ensures this by checking `TypeId` prior to downcasting.
    fn as_any(&self) -> &dyn Any {
        self
    }

    /// Returns a mutable, type-erased view of this concrete storage for downcasting.
    ///
    /// # Returns
    /// - `&mut dyn Any`: A mutable reference that can be downcast via
    ///   `Any::downcast_mut` to `ComponentStorage<T>` when the caller knows `T`.
    ///
    /// # Safety
    /// - Callers must ensure the downcast target type matches the underlying
    ///   storage. `World` APIs guard this via `TypeId` checks before downcasting.
    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }
}