MeshLayouts and iter_meshes()

Cargo.toml

@@ -6,6 +6,7 @@ edition = "2021"
 [dependencies]
 rayon = "1"
 slab = "^0.4"
+smallmap = { git = "https://github.com/marceline-cramer/smallmap" }
 smallvec = "^1.0"
 strum = { version = "0.24", features = ["derive"] }
 wgpu = "^0.12"

src/mesh.rs

@@ -12,6 +12,7 @@ pub enum PoolError {
     InvalidIndex,
     AttrTaken,
     AttrUnregistered,
+    LayoutUnregistered,
     MismatchedId,
     MismatchedLayout,
 }
@@ -144,6 +145,12 @@ pub struct AttrPool {
     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(id: AttrId, pool_id: usize, info: AttrInfo) -> Self {
         let layout = info.layout;
@@ -223,11 +230,7 @@ impl AttrPool {
     /// 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> {
+    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),
@@ -281,6 +284,12 @@ impl AttrPool {
 
         Ok(())
     }
+
+    /// Gets the offset (in elements) of an allocation, passed by key.
+    pub fn get_offset(&self, key: usize) -> Result<usize, PoolError> {
+        let alloc = self.allocs.get(key).ok_or(PoolError::InvalidIndex)?;
+        Ok(alloc.offset)
+    }
 }
 
 /// The number of attributes a mesh can have before they're moved to the heap.
@@ -300,16 +309,25 @@ pub struct MeshAlloc {
 #[repr(transparent)]
 pub struct MeshHandle(usize);
 
-impl HasMeshHandle for MeshHandle {
-    fn get_mesh_handle(&self) -> &Self {
-        self
-    }
-}
+/// A reusable set of [AttrIds][AttrId], for use with querying compatible meshes.
+pub type MeshLayoutDesc = smallmap::Set<AttrId>;
 
-/// A trait for structs containing [MeshHandles][MeshHandle] that are not
-/// themselves handles. Used for iteration.
-pub trait HasMeshHandle {
-    fn get_mesh_handle(&self) -> &MeshHandle;
+/// 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 pools.
+pub type MeshLayoutBindings<'a> = smallmap::Map<AttrId, &'a AttrPool>;
+
+/// Mappings of the attributes in a [MeshAlloc] to specific pool offsets.
+pub type MeshAllocOffsets = SmallVec<[(AttrId, usize); MAX_MESH_INLINE_ATTRIBUTES]>;
+
+/// A set of mesh instances fitting a common [MeshLayoutDesc].
+#[derive(Default)]
+pub struct MeshLayoutInstances<'a, T> {
+    pub bindings: MeshLayoutBindings<'a>,
+    pub instances: Vec<(T, MeshAllocOffsets)>,
 }
 
 /// A mesh data pool.
@@ -318,6 +336,7 @@ pub struct MeshPool {
     staging: StagingPool<AttrAllocKey>,
     attr_store: Arc<AttrStore>,
     allocs: Slab<MeshAlloc>,
+    mesh_layouts: Slab<MeshLayoutDesc>,
     pools: HashMap<AttrId, Vec<AttrPool>>,
 }
 
@@ -327,11 +346,20 @@ impl MeshPool {
             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(&mut self, layout: MeshLayoutDesc) -> Result<MeshLayoutId, PoolError> {
+        let idx = self.mesh_layouts.insert(layout);
+        Ok(MeshLayoutId(idx))
+    }
+
     /// Loads a [MeshBuffer].
     pub fn load(&mut self, buf: MeshBuffer) -> Result<MeshHandle, PoolError> {
         let mut attrs = HashMap::new();
@@ -388,4 +416,75 @@ impl MeshPool {
         let handle = MeshHandle(key);
         Ok(handle)
     }
+
+    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) -> &'static MeshHandle,
+    {
+        let layout = self
+            .mesh_layouts
+            .get(layout.0)
+            .ok_or(PoolError::LayoutUnregistered)?;
+
+        let layout_attrs: Vec<AttrId> = layout.iter().map(|(id, _)| *id).collect();
+        let mut layouts = Vec::<MeshLayoutInstances<T>>::new();
+
+        let mut attr_allocs: Vec<AttrAllocKey> = Vec::with_capacity(layout_attrs.len());
+        for mesh in meshes {
+            let handle = get_handle(&mesh);
+            let alloc = self.allocs.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 = MeshLayoutBindings::default();
+            let mut alloc_offsets = MeshAllocOffsets::default();
+            for alloc in attr_allocs.iter() {
+                let pools = self
+                    .pools
+                    .get(&alloc.attr)
+                    .ok_or(PoolError::AttrUnregistered)?;
+
+                let pool = pools.get(alloc.pool).ok_or(PoolError::InvalidIndex)?;
+                let alloc_offset = pool.get_offset(alloc.alloc)?;
+
+                layout_bindings.insert(alloc.attr, pool);
+                alloc_offsets.push((alloc.attr, alloc_offset));
+            }
+
+            let instance = (mesh, alloc_offsets);
+
+            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)
+    }
 }