572 lines
17 KiB
Rust
572 lines
17 KiB
Rust
use crate::staging::*;
|
|
use parking_lot::{RwLock, RwLockReadGuard};
|
|
use slab::Slab;
|
|
use smallvec::SmallVec;
|
|
use std::any::TypeId;
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
/// An error that can be returned when allocating a mesh.
|
|
#[derive(Debug)]
|
|
pub enum PoolError {
|
|
TooBig,
|
|
NoMoreRoom,
|
|
InvalidIndex,
|
|
AttrTaken,
|
|
AttrUnregistered,
|
|
LayoutUnregistered,
|
|
MismatchedId,
|
|
MismatchedLayout,
|
|
}
|
|
|
|
/// An identifier for a mesh attribute.
|
|
#[repr(transparent)]
|
|
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
|
|
pub struct AttrId(usize);
|
|
|
|
/// A description of an 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 usages: wgpu::BufferUsages,
|
|
pub default_pool_size: usize,
|
|
}
|
|
|
|
/// The data single mesh attribute.
|
|
#[derive(Clone)]
|
|
pub struct AttrBuffer {
|
|
pub id: AttrId,
|
|
pub count: usize,
|
|
pub data: Vec<u8>,
|
|
}
|
|
|
|
/// 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 [wgpu::BufferUsages] of this data.
|
|
fn get_usages() -> wgpu::BufferUsages;
|
|
|
|
/// 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: RwLock<Slab<AttrInfo>>,
|
|
types: RwLock<HashMap<TypeId, AttrId>>,
|
|
}
|
|
|
|
impl AttrStore {
|
|
pub fn new() -> Arc<Self> {
|
|
Arc::new(Self {
|
|
attributes: Default::default(),
|
|
types: Default::default(),
|
|
})
|
|
}
|
|
|
|
/// Dynamically creates a new [AttrId].
|
|
pub fn add(&self, info: AttrInfo) -> AttrId {
|
|
let index = self.attributes.write().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>(&self) -> AttrId {
|
|
let type_id = TypeId::of::<T>();
|
|
let existing_id = self.types.read().get(&type_id).copied();
|
|
|
|
if let Some(id) = existing_id {
|
|
id
|
|
} else {
|
|
let layout = T::get_layout();
|
|
let usages = T::get_usages();
|
|
let default_pool_size = T::get_default_pool_size();
|
|
let info = AttrInfo {
|
|
layout,
|
|
usages,
|
|
default_pool_size,
|
|
};
|
|
let id = self.add(info);
|
|
self.types.write().insert(type_id, id);
|
|
id
|
|
}
|
|
}
|
|
|
|
/// Gets the [AttrInfo] for an [AttrId].
|
|
pub fn get_info(&self, id: &AttrId) -> Option<AttrInfo> {
|
|
self.attributes.read().get(id.0).copied()
|
|
}
|
|
}
|
|
|
|
/// Attribute pool allocation location.
|
|
#[derive(Clone)]
|
|
pub struct AttrAllocKey {
|
|
/// The target attribute.
|
|
pub attr: AttrId,
|
|
|
|
/// The index of the attribute pool.
|
|
pub pool: usize,
|
|
|
|
/// The allocation within the attribute pool.
|
|
pub alloc: usize,
|
|
}
|
|
|
|
/// Info about an array of attributes that has been allocated in an [AttrPool].
|
|
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
|
|
pub struct AttrAlloc {
|
|
pub offset: usize,
|
|
pub count: usize,
|
|
pub offset_bytes: usize,
|
|
pub count_bytes: 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 {
|
|
buffer: wgpu::Buffer,
|
|
id: AttrId,
|
|
pool_id: usize,
|
|
layout: AttrLayout,
|
|
size: usize,
|
|
allocs: Slab<AttrAlloc>,
|
|
free_space: Vec<FreeSpace>,
|
|
}
|
|
|
|
impl PartialEq for AttrPool {
|
|
fn eq(&self, other: &Self) -> bool {
|
|
self.pool_id == other.pool_id && self.id == other.id
|
|
}
|
|
}
|
|
|
|
impl AttrPool {
|
|
pub fn new(device: &wgpu::Device, id: AttrId, pool_id: usize, info: AttrInfo) -> Self {
|
|
let layout = info.layout;
|
|
let size = info.default_pool_size;
|
|
|
|
// TODO debug strings for attributes + pools + buffers
|
|
let buffer = device.create_buffer(&wgpu::BufferDescriptor {
|
|
label: None,
|
|
usage: wgpu::BufferUsages::COPY_DST | info.usages,
|
|
mapped_at_creation: false,
|
|
size: (size * layout.size) as wgpu::BufferAddress,
|
|
});
|
|
|
|
Self {
|
|
buffer,
|
|
id,
|
|
pool_id,
|
|
layout,
|
|
size,
|
|
free_space: vec![FreeSpace {
|
|
offset: 0,
|
|
count: size,
|
|
}],
|
|
allocs: Default::default(),
|
|
}
|
|
}
|
|
|
|
/// Loads an [AttrBuffer].
|
|
///
|
|
/// Returns the new [AttrAllocKey].
|
|
pub fn load(&mut self, buf: &AttrBuffer) -> Result<AttrAllocKey, PoolError> {
|
|
let best_index = self.can_load(buf)?;
|
|
self.alloc_at(best_index, buf.count)
|
|
}
|
|
|
|
/// Allocates an array of attributes.
|
|
///
|
|
/// Returns the new [AttrAllocKey].
|
|
pub fn alloc(&mut self, count: usize) -> Result<AttrAllocKey, PoolError> {
|
|
let best_index = self.can_alloc(count)?;
|
|
self.alloc_at(best_index, count)
|
|
}
|
|
|
|
/// 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.count * self.layout.size != buf.data.len() {
|
|
Err(PoolError::MismatchedLayout)
|
|
} else {
|
|
self.can_alloc(buf.count)
|
|
}
|
|
}
|
|
|
|
/// Tests if an array of 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)
|
|
}
|
|
}
|
|
|
|
/// 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
|
|
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 an array of attribute at a specific free space index.
|
|
///
|
|
/// Returns the new [AttrAllocKey].
|
|
fn alloc_at(&mut self, index: usize, count: usize) -> Result<AttrAllocKey, PoolError> {
|
|
let free_space = match self.free_space.get_mut(index) {
|
|
Some(index) => index,
|
|
None => return Err(PoolError::InvalidIndex),
|
|
};
|
|
|
|
let offset = free_space.offset;
|
|
|
|
let alloc = AttrAlloc {
|
|
offset,
|
|
count,
|
|
offset_bytes: offset * self.layout.size,
|
|
count_bytes: count * self.layout.size,
|
|
};
|
|
|
|
let index = self.allocs.insert(alloc);
|
|
|
|
let key = AttrAllocKey {
|
|
attr: self.id,
|
|
pool: self.pool_id,
|
|
alloc: index,
|
|
};
|
|
|
|
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(key)
|
|
}
|
|
|
|
/// Frees an allocation (by key) from the pool.
|
|
pub fn free(&mut self, key: usize) -> Result<(), PoolError> {
|
|
let alloc = self.allocs.try_remove(key).ok_or(PoolError::InvalidIndex)?;
|
|
|
|
let free_space = FreeSpace {
|
|
offset: alloc.offset,
|
|
count: alloc.count,
|
|
};
|
|
|
|
// TODO merge free spaces
|
|
self.free_space.push(free_space);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Retrieves an [AttrAlloc] by key.
|
|
pub fn get(&self, key: usize) -> Option<AttrAlloc> {
|
|
self.allocs.get(key).copied()
|
|
}
|
|
|
|
/// Gets this pool's internal GPU buffer.
|
|
pub fn get_buffer(&self) -> &wgpu::Buffer {
|
|
&self.buffer
|
|
}
|
|
|
|
/// Gets a [CopyDest] for an allocation, by key.
|
|
pub fn get_copy_dest(&self, key: usize) -> Result<CopyDest, PoolError> {
|
|
let offset = self.get(key).ok_or(PoolError::InvalidIndex)?.offset_bytes;
|
|
Ok(CopyDest {
|
|
buffer: self.get_buffer(),
|
|
offset,
|
|
})
|
|
}
|
|
}
|
|
|
|
/// The number of attributes a mesh can have before they're moved to the heap.
|
|
pub const MAX_MESH_INLINE_ATTRIBUTES: usize = 16;
|
|
|
|
/// A mesh and all of its attributes.
|
|
#[derive(Clone, Default)]
|
|
pub struct MeshBuffer {
|
|
pub attributes: SmallVec<[AttrBuffer; MAX_MESH_INLINE_ATTRIBUTES]>,
|
|
}
|
|
|
|
/// An allocated mesh.
|
|
pub struct MeshAlloc {
|
|
pub attributes: SmallVec<[AttrAllocKey; MAX_MESH_INLINE_ATTRIBUTES]>,
|
|
}
|
|
|
|
/// A handle to an allocated mesh.
|
|
#[repr(transparent)]
|
|
#[derive(Copy, Clone, Debug)]
|
|
pub struct MeshHandle(usize);
|
|
|
|
/// A reusable set of [AttrIds][AttrId], for use with querying compatible meshes.
|
|
pub type MeshLayoutDesc = smallmap::Set<AttrId>;
|
|
|
|
/// The ID of a [MeshLayoutDesc] registered in a [MeshPool].
|
|
#[repr(transparent)]
|
|
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
|
|
pub struct MeshLayoutId(usize);
|
|
|
|
/// Mappings of the attributes in a [MeshLayout] to specific pool indices.
|
|
pub type MeshLayoutBindingIndices = smallmap::Map<AttrId, usize>;
|
|
|
|
/// Mappings of the attributes in a [MeshLayout] to specific pools.
|
|
pub struct MeshLayoutBindings<'a> {
|
|
lock: RwLockReadGuard<'a, HashMap<AttrId, Vec<AttrPool>>>,
|
|
indices: MeshLayoutBindingIndices,
|
|
}
|
|
|
|
impl<'a> MeshLayoutBindings<'a> {
|
|
pub fn get(&self, attr: AttrId) -> Option<&AttrPool> {
|
|
let pool_id = self.indices.get(&attr)?;
|
|
let pools = self.lock.get(&attr)?;
|
|
pools.get(*pool_id)
|
|
}
|
|
}
|
|
|
|
/// Mappings of the attributes in a [MeshAlloc] to specific pool offsets.
|
|
pub type MeshAllocInfos = SmallVec<[(AttrId, AttrAlloc); MAX_MESH_INLINE_ATTRIBUTES]>;
|
|
|
|
/// A set of mesh instances fitting a common [MeshLayoutDesc].
|
|
#[derive(Default)]
|
|
pub struct MeshLayoutInstances<T> {
|
|
pub bindings: MeshLayoutBindingIndices,
|
|
pub instances: Vec<(T, MeshAllocInfos)>,
|
|
}
|
|
|
|
/// A mesh data pool.
|
|
pub struct MeshPool {
|
|
device: Arc<wgpu::Device>,
|
|
staging: StagingPool<AttrAllocKey>,
|
|
attr_store: Arc<AttrStore>,
|
|
allocs: RwLock<Slab<MeshAlloc>>,
|
|
mesh_layouts: RwLock<Slab<MeshLayoutDesc>>,
|
|
pools: RwLock<HashMap<AttrId, Vec<AttrPool>>>,
|
|
}
|
|
|
|
impl MeshPool {
|
|
pub fn new(device: Arc<wgpu::Device>, attr_store: Arc<AttrStore>) -> Arc<Self> {
|
|
Arc::new(Self {
|
|
device: device.clone(),
|
|
staging: StagingPool::new(device, 1024 * 1024),
|
|
attr_store,
|
|
mesh_layouts: Default::default(),
|
|
allocs: Default::default(),
|
|
pools: Default::default(),
|
|
})
|
|
}
|
|
|
|
/// Registers a [MeshLayoutDesc] with this pool.
|
|
///
|
|
/// TODO: keep track of mesh allocations that fit each layout
|
|
pub fn add_layout(&self, layout: MeshLayoutDesc) -> Result<MeshLayoutId, PoolError> {
|
|
let idx = self.mesh_layouts.write().insert(layout);
|
|
Ok(MeshLayoutId(idx))
|
|
}
|
|
|
|
/// Loads a [MeshBuffer].
|
|
pub fn load(&self, buf: MeshBuffer) -> Result<MeshHandle, PoolError> {
|
|
let mut attrs = HashMap::new();
|
|
|
|
for attr in buf.attributes.into_iter() {
|
|
if let Some(_) = attrs.insert(attr.id, attr) {
|
|
// TODO: support for loading duplicate attribute IDs
|
|
return Err(PoolError::AttrTaken);
|
|
}
|
|
}
|
|
|
|
let mut attr_allocs = SmallVec::new();
|
|
let mut copies = Vec::new();
|
|
for (id, buf) in attrs.drain() {
|
|
let mut pools_write = self.pools.write();
|
|
|
|
let pools = match pools_write.get_mut(&id) {
|
|
Some(pools) => pools,
|
|
None => {
|
|
let info = match self.attr_store.get_info(&id) {
|
|
Some(info) => info,
|
|
None => return Err(PoolError::AttrUnregistered),
|
|
};
|
|
|
|
let pool = AttrPool::new(&self.device, id, 0, info);
|
|
pools_write.insert(id, vec![pool]);
|
|
pools_write.get_mut(&id).unwrap()
|
|
}
|
|
};
|
|
|
|
for pool in pools.iter_mut() {
|
|
match pool.load(&buf) {
|
|
Ok(alloc) => {
|
|
let copy = CopyBuffer {
|
|
target: alloc.clone(),
|
|
offset: 0,
|
|
data: buf.data,
|
|
};
|
|
|
|
attr_allocs.push(alloc);
|
|
copies.push(copy);
|
|
break;
|
|
}
|
|
Err(PoolError::NoMoreRoom) => {}
|
|
Err(e) => return Err(e),
|
|
}
|
|
}
|
|
|
|
// TODO create a new pool if no available pool was found
|
|
}
|
|
|
|
self.staging.queue_copies(copies);
|
|
|
|
let alloc = MeshAlloc {
|
|
attributes: attr_allocs,
|
|
};
|
|
let key = self.allocs.write().insert(alloc);
|
|
let handle = MeshHandle(key);
|
|
Ok(handle)
|
|
}
|
|
|
|
pub fn flush(&self, commands: &mut wgpu::CommandEncoder) {
|
|
let pools_read = self.pools.read();
|
|
|
|
let get_dst = |target: &AttrAllocKey| {
|
|
pools_read
|
|
.get(&target.attr)
|
|
.unwrap()
|
|
.get(target.pool)
|
|
.unwrap()
|
|
.get_copy_dest(target.alloc)
|
|
.unwrap()
|
|
};
|
|
|
|
// TODO: keep track of loaded/unloaded meshes
|
|
let on_complete = |_target| {};
|
|
|
|
self.staging.flush(commands, get_dst, on_complete);
|
|
}
|
|
|
|
pub fn iter_meshes<T, I, F>(
|
|
&self,
|
|
layout: MeshLayoutId,
|
|
meshes: I,
|
|
get_handle: F,
|
|
) -> Result<Vec<MeshLayoutInstances<T>>, PoolError>
|
|
where
|
|
I: Iterator<Item = T>,
|
|
F: Fn(&T) -> &MeshHandle,
|
|
{
|
|
let layout = self
|
|
.mesh_layouts
|
|
.read()
|
|
.get(layout.0)
|
|
.ok_or(PoolError::LayoutUnregistered)?
|
|
.clone();
|
|
|
|
let layout_attrs: Vec<AttrId> = layout.iter().map(|(id, _)| *id).collect();
|
|
let mut layouts = Vec::<MeshLayoutInstances<T>>::new();
|
|
|
|
let allocs_read = self.allocs.read();
|
|
let mut attr_allocs: Vec<AttrAllocKey> = Vec::with_capacity(layout_attrs.len());
|
|
for mesh in meshes {
|
|
let handle = get_handle(&mesh);
|
|
let alloc = allocs_read.get(handle.0).ok_or(PoolError::InvalidIndex)?;
|
|
|
|
attr_allocs.clear();
|
|
for layout_attr in layout_attrs.iter() {
|
|
match alloc
|
|
.attributes
|
|
.iter()
|
|
.find(|attr_alloc| attr_alloc.attr == *layout_attr)
|
|
{
|
|
Some(alloc) => attr_allocs.push(alloc.clone()),
|
|
None => break,
|
|
}
|
|
}
|
|
|
|
if attr_allocs.len() != layout_attrs.len() {
|
|
continue;
|
|
}
|
|
|
|
let mut layout_bindings = MeshLayoutBindingIndices::default();
|
|
let mut alloc_infos = MeshAllocInfos::default();
|
|
for alloc in attr_allocs.iter() {
|
|
let pools_read = self.pools.read();
|
|
let pools = pools_read
|
|
.get(&alloc.attr)
|
|
.ok_or(PoolError::AttrUnregistered)?;
|
|
|
|
let pool = pools.get(alloc.pool).ok_or(PoolError::InvalidIndex)?;
|
|
let alloc_info = pool.get(alloc.alloc).ok_or(PoolError::InvalidIndex)?;
|
|
|
|
layout_bindings.insert(alloc.attr, alloc.pool);
|
|
alloc_infos.push((alloc.attr, alloc_info));
|
|
}
|
|
|
|
let instance = (mesh, alloc_infos);
|
|
|
|
match layouts
|
|
.iter_mut()
|
|
.find(|layout| layout.bindings == layout_bindings)
|
|
{
|
|
Some(layout) => layout.instances.push(instance),
|
|
None => layouts.push(MeshLayoutInstances {
|
|
bindings: layout_bindings,
|
|
instances: vec![instance],
|
|
}),
|
|
}
|
|
}
|
|
|
|
Ok(layouts)
|
|
}
|
|
|
|
pub fn get_bindings<'a>(&'a self, indices: MeshLayoutBindingIndices) -> MeshLayoutBindings<'a> {
|
|
let lock = self.pools.read();
|
|
MeshLayoutBindings { lock, indices }
|
|
}
|
|
}
|