hyperion/

resources.rs

use std::any::{Any, TypeId};
use std::collections::HashMap;

/// A generational handle used to identify a resource stored in a `ResourceRegistry`.
///
/// # Fields
/// - `id`: A unique identifier assigned when the resource is registered.
/// - `version`: A generation number used to validate the handle's freshness.
///
/// # Safety
/// - Handles become invalid if the associated resource is removed (version is incremented).
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct ResourceHandle {
    id: u64,
    version: u32,
}

/// A type-erased registry for storing arbitrary resources using generational handles.
///
/// # Fields
/// - `next_id`: Next available numeric id for newly registered resources.
/// - `versions`: Per-id version counter (generation) for handle validation.
/// - `stores`: Map of `TypeId` -> boxed `Vec<T>` type-erased.
/// - `ids`: Map of id -> (TypeId, index in its store).
/// - `reverse_ids`: Reverse map of (TypeId, index) -> id.
/// - `unique_resources`: Map of `TypeId` -> unique resource instance.
///
/// # Safety
/// - All type-erased accesses are guarded by `TypeId` checks and handle version
///   validation to prevent type confusion and use-after-free via stale handles.
pub struct ResourceRegistry {
    next_id: u64,
    versions: HashMap<u64, u32>,
    stores: HashMap<TypeId, Box<dyn Any>>,
    ids: HashMap<u64, (TypeId, usize)>,
    reverse_ids: HashMap<(TypeId, usize), u64>,
    unique_resources: HashMap<TypeId, Box<dyn Any>>,
}

impl ResourceRegistry {
    /// Creates a new, empty `ResourceRegistry`.
    pub fn new() -> Self {
        Self {
            next_id: 0,
            versions: HashMap::new(),
            stores: HashMap::new(),
            ids: HashMap::new(),
            reverse_ids: HashMap::new(),
            unique_resources: HashMap::new(),
        }
    }

    /// Registers or overwrites the unique resource of type `T`.
    ///
    /// If a resource of this type already exists, it will be replaced.
    ///
    /// # Type Parameters
    /// - `T`: The type of the resource. Must be `'static`.
    ///
    /// # Parameters
    /// - `resource`: The resource instance to store as a unique singleton.
    pub fn register_unique<T: 'static>(&mut self, resource: T) {
        let tid = TypeId::of::<T>();
        self.unique_resources.insert(tid, Box::new(resource));
    }

    /// Gets an immutable reference to the unique resource of type `T`, if it exists.
    ///
    /// # Type Parameters
    /// - `T`: The type of the resource requested.
    ///
    /// # Returns
    /// - `Option<&T>`: Some reference to the resource if registered; None if not found.
    pub fn get_unique<T: 'static>(&self) -> Option<&T> {
        let tid = TypeId::of::<T>();
        self.unique_resources.get(&tid)?.downcast_ref::<T>()
    }

    /// Gets a mutable reference to the unique resource of type `T`, if it exists.
    ///
    /// # Type Parameters
    /// - `T`: The type of the resource requested.
    ///
    /// # Returns
    /// - `Option<&mut T>`: Some mutable reference to the resource if registered; None if not found.
    pub fn get_unique_mut<T: 'static>(&mut self) -> Option<&mut T> {
        let tid = TypeId::of::<T>();
        self.unique_resources.get_mut(&tid)?.downcast_mut::<T>()
    }

    /// Removes the unique resource of type `T` from storage and returns it.
    ///
    /// # Type Parameters
    /// - `T`: The type of the resource to remove.
    ///
    /// # Returns
    /// - `Option<T>`: The removed resource if it existed; None if no resource of type `T` was found.
    pub fn remove_unique<T: 'static>(&mut self) -> Option<T> {
        let tid = TypeId::of::<T>();
        self.unique_resources
            .remove(&tid)?
            .downcast::<T>()
            .ok()
            .map(|boxed| *boxed)
    }

    /// Registers a resource of type `T` and returns a unique `ResourceHandle`.
    ///
    /// # Type Parameters
    /// - `T`: The concrete type of the resource to store. Must be `'static`.
    ///
    /// # Parameters
    /// - `resource`: The resource instance to store in the registry.
    ///
    /// # Returns
    /// - `ResourceHandle`: A generational handle uniquely identifying the stored resource.
    ///
    /// # Notes
    /// - The handle can later be used to retrieve or remove the resource.
    /// - Handles are versioned to prevent use-after-free access.
    pub fn register_resource<T: 'static>(&mut self, resource: T) -> ResourceHandle {
        let tid = TypeId::of::<T>();

        let store = self
            .stores
            .entry(tid)
            .or_insert_with(|| Box::new(Vec::<T>::new()));
        let vec = store.downcast_mut::<Vec<T>>().unwrap();

        let index = vec.len();
        vec.push(resource);

        let id = self.next_id;
        self.next_id += 1;

        let version = 1;
        self.versions.insert(id, version);
        self.ids.insert(id, (tid, index));
        self.reverse_ids.insert((tid, index), id);

        ResourceHandle { id, version }
    }

    /// Retrieves an immutable reference to a resource using its handle.
    ///
    /// # Type Parameters
    /// - `T`: The expected type of the resource to retrieve.
    ///
    /// # Parameters
    /// - `handle`: A `ResourceHandle` returned from `register_resource`.
    ///
    /// # Returns
    /// - `Some(&T)`: A reference to the resource if the handle is valid and the type matches.
    /// - `None`: If the handle is invalid, outdated (version mismatch), or of the wrong type.
    ///
    /// # Safety
    /// - Ensures the handle's version matches the current version of the stored resource.
    pub fn get_resource<T: 'static>(&self, handle: ResourceHandle) -> Option<&T> {
        let tid = TypeId::of::<T>();

        let (stored_tid, index) = self.ids.get(&handle.id)?;
        if *stored_tid != tid {
            return None;
        }

        let current_version = self.versions.get(&handle.id)?;
        if *current_version != handle.version {
            return None;
        }

        let vec = self.stores.get(&tid)?.downcast_ref::<Vec<T>>()?;
        vec.get(*index)
    }

    /// Retrieves a mutable reference to a resource using its handle.
    ///
    /// # Type Parameters
    /// - `T`: The expected type of the resource to retrieve.
    ///
    /// # Parameters
    /// - `handle`: A `ResourceHandle` returned from `register_resource`.
    ///
    /// # Returns
    /// - `Some(&mut T)`: A mutable reference to the resource if the handle is valid and type matches.
    /// - `None`: If the handle is invalid, outdated, or of the wrong type.
    ///
    /// # Safety
    /// - Only succeeds if the handle is still valid and uniquely identifies a resource of type `T`.
    pub fn get_resource_mut<T: 'static>(&mut self, handle: ResourceHandle) -> Option<&mut T> {
        let tid = TypeId::of::<T>();

        let (stored_tid, index) = self.ids.get(&handle.id)?;
        if *stored_tid != tid {
            return None;
        }

        let current_version = self.versions.get(&handle.id)?;
        if *current_version != handle.version {
            return None;
        }

        let vec = self.stores.get_mut(&tid)?.downcast_mut::<Vec<T>>()?;
        vec.get_mut(*index)
    }

    /// Removes a resource from the registry, invalidating its handle.
    ///
    /// # Type Parameters
    /// - `T`: The expected type of the resource to remove.
    ///
    /// # Parameters
    /// - `handle`: The `ResourceHandle` identifying the resource to remove.
    ///
    /// # Returns
    /// - `Some(T)`: The removed resource if the handle was valid and matched the expected type.
    /// - `None`: If the handle was invalid, outdated (version mismatch), or mismatched in type.
    ///
    /// # Notes
    /// - This function also increments the resource's version, invalidating all existing handles.
    /// - If the removed element is not at the end of its internal storage, the moved resource’s
    ///   index is updated, and its handle remains valid.
    pub fn remove_resource<T: 'static>(&mut self, handle: ResourceHandle) -> Option<T> {
        let tid = TypeId::of::<T>();

        let (stored_tid, index) = self.ids.remove(&handle.id)?;
        if stored_tid != tid {
            return None;
        }

        let version = self.versions.get_mut(&handle.id)?;
        if *version != handle.version {
            return None;
        }

        // Invalidate the handle by bumping its version
        *version += 1;

        let vec = self.stores.get_mut(&tid)?.downcast_mut::<Vec<T>>()?;

        if index >= vec.len() {
            return None;
        }

        // Remove resource from vec by swapping with last, then popping
        let removed = vec.swap_remove(index);

        // Clean up reverse ID map
        self.reverse_ids.remove(&(tid, index));

        // If we swapped with the last element, update the index of the moved one
        if index < vec.len() {
            let moved_id = self.reverse_ids.remove(&(tid, vec.len()))?;
            self.ids.insert(moved_id, (tid, index));
            self.reverse_ids.insert((tid, index), moved_id);
        }

        Some(removed)
    }
}