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::*;

/// 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 mesh attribute.
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub struct AttrLayout {
    /// The size (in bytes) of this attribute.
    pub size: usize,
}

/// 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,
        layout: AttrLayout,
        size: usize,
    ) -> Result<Self, PoolError> {
        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!()
    }
}