cyborg/src/mesh/attr.rs

//! Mesh storage pooling for a single attribute.
//!
//! Attribute pools have a fixed size, and once created cannot be expanded to
//! fit more data.

use super::*;
use std::any::TypeId;
use std::collections::HashMap;

/// An externally-defined identifier for a mesh attribute.
#[repr(transparent)]
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub struct AttrId(pub usize);

/// A description of a attribute's layout in memory.
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub struct AttrLayout {
    /// The size (in bytes) of this attribute.
    pub size: usize,
}

/// Information about an [Attribute] registered in [AttrStore].
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub struct AttrInfo {
    pub layout: AttrLayout,
    pub default_pool_size: usize,
}

/// A compile-time attribute data type.
pub trait Attribute: Sized {
    /// The memory layout of this data.
    fn get_layout() -> AttrLayout {
        AttrLayout {
            size: std::mem::size_of::<Self>(),
        }
    }

    /// The default size for new pools of this attribute.
    ///
    /// Defaults to 1024 * 1024. (Around one million.)
    fn get_default_pool_size() -> usize {
        1024 * 1024
    }
}

/// A store of [AttrIds][AttrId].
pub struct AttrStore {
    attributes: Slab<AttrInfo>,
    types: HashMap<TypeId, AttrId>,
}

impl AttrStore {
    pub fn new() -> Self {
        Self {
            attributes: Default::default(),
            types: Default::default(),
        }
    }

    /// Dynamically creates a new [AttrId].
    pub fn add(&mut self, info: AttrInfo) -> AttrId {
        let index = self.attributes.insert(info);
        AttrId(index)
    }

    /// Gets the [AttrId] for a type implementing [Attribute].
    ///
    /// Creates a new [AttrId] for unrecognized types, otherwise reuses an
    /// existing [AttrId].
    pub fn get_type<T: 'static + Attribute>(&mut self) -> AttrId {
        let type_id = TypeId::of::<T>();
        if let Some(id) = self.types.get(&type_id) {
            *id
        } else {
            let layout = T::get_layout();
            let default_pool_size = T::get_default_pool_size();
            let info = AttrInfo {
                layout,
                default_pool_size,
            };
            let id = self.add(info);
            self.types.insert(type_id, id);
            id
        }
    }

    /// Gets the [AttrInfo] for an [AttrId].
    pub fn get_info(&self, id: &AttrId) -> Option<&AttrInfo> {
        self.attributes.get(id.0)
    }
}

/// An attribute buffer that has been allocated in an [AttrPool].
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub struct AttrAlloc {
    offset: usize,
    count: usize,
}

/// An unused space range in an [AttrPool].
pub struct FreeSpace {
    offset: usize,
    count: usize,
}

/// A single GPU buffer containing linear arrays of attributes.
pub struct AttrPool {
    group: usize,
    id: AttrId,
    layout: AttrLayout,
    size: usize,
    allocs: Slab<AttrAlloc>,
    free_space: Vec<FreeSpace>,
}

impl AttrPool {
    pub fn new(group: usize, id: AttrId, info: AttrInfo) -> Result<Self, PoolError> {
        let layout = info.layout;
        let size = info.default_pool_size;

        Ok(Self {
            group,
            id,
            layout,
            size,
            free_space: vec![FreeSpace {
                offset: 0,
                count: size,
            }],
            allocs: Default::default(),
        })
    }

    /// Tests if attributes can be allocated.
    ///
    /// Returns the result of [Self::best_fit].
    pub fn can_alloc(&self, count: usize) -> Result<usize, PoolError> {
        if count > self.size {
            Err(PoolError::TooBig)
        } else if let Some(best_index) = self.best_fit(count) {
            Ok(best_index)
        } else {
            Err(PoolError::NoMoreRoom)
        }
    }

    /// Tests if an [AttrBuffer] can be loaded.
    ///
    /// Returns the result of [Self::best_fit].
    pub fn can_load(&self, buf: &AttrBuffer) -> Result<usize, PoolError> {
        if buf.id != self.id {
            Err(PoolError::MismatchedId)
        } else if buf.layout != self.layout {
            Err(PoolError::MismatchedLayout)
        } else {
            self.can_alloc(buf.count)
        }
    }

    /// Finds the index of the best-fit free space for an array of attributes.
    ///
    /// TODO: use a binary tree to find best-fit free space in logarithmic time
    pub fn best_fit(&self, count: usize) -> Option<usize> {
        let mut best_index = None;
        let mut best_count = usize::MAX;
        for (index, space) in self.free_space.iter().enumerate() {
            if space.count >= count && space.count < best_count {
                best_index = Some(index);
                best_count = space.count;
            }
        }

        best_index
    }

    /// Allocates room for attributes at a specific free space index.
    ///
    /// Returns the new [AttrAlloc] and its key.
    pub fn alloc_at(
        &mut self,
        index: usize,
        count: usize,
    ) -> Result<(AttrAlloc, usize), PoolError> {
        let free_space = match self.free_space.get_mut(index) {
            Some(index) => index,
            None => return Err(PoolError::InvalidIndex),
        };

        let alloc = AttrAlloc {
            offset: free_space.offset,
            count,
        };

        let key = self.allocs.insert(alloc);

        use std::cmp::Ordering;
        match free_space.count.cmp(&count) {
            Ordering::Less => {
                return Err(PoolError::TooBig);
            }
            Ordering::Equal => {
                self.free_space.remove(index);
            }
            Ordering::Greater => {
                free_space.count -= count;
                free_space.offset += count;
            }
        }

        Ok((alloc, key))
    }

    /// Allocates room for attributes.
    ///
    /// Returns the new [AttrAlloc] and its key.
    pub fn alloc(&mut self, count: usize) -> Result<(AttrAlloc, usize), PoolError> {
        let best_index = self.can_alloc(count)?;
        self.alloc_at(best_index, count)
    }

    /// Loads an [AttrBuffer].
    ///
    /// Returns the key for the allocation, as well as [CopyInfo] that can be
    /// queued into a [StagingPool].
    pub fn load(&mut self, buf: &AttrBuffer) -> Result<(usize, CopyInfo), PoolError> {
        let best_index = self.can_load(buf)?;
        let (alloc, key) = self.alloc_at(best_index, buf.count)?;

        let copy = CopyInfo {
            group: self.group,
            target: self.id,
            offset: alloc.offset * self.layout.size,
            size: alloc.count * self.layout.size,
        };

        Ok((key, copy))
    }

    /// Frees an allocation (by key) from the pool.
    pub fn free(&mut self, alloc: usize) -> Result<(), PoolError> {
        todo!()
    }
}