From b4d6d6cba0eecb1faf3aa40b8aa333217cbfaf72 Mon Sep 17 00:00:00 2001
From: mars <mars@tebibyte.media>
Date: Mon, 18 Apr 2022 01:20:50 -0600
Subject: [PATCH] AttrPool keeps buffer; impl MeshPool::flush()

---
 src/mesh.rs    | 47 +++++++++++++++++++++++++++++++++++++++++++++--
 src/staging.rs | 10 +++++++---
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/src/mesh.rs b/src/mesh.rs
index 925b90d..ff76bfd 100644
--- a/src/mesh.rs
+++ b/src/mesh.rs
@@ -6,6 +6,7 @@ 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,
@@ -33,6 +34,7 @@ pub struct AttrLayout {
 #[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
 pub struct AttrInfo {
     pub layout: AttrLayout,
+    pub usages: wgpu::BufferUsages,
     pub default_pool_size: usize,
 }
 
@@ -52,6 +54,9 @@ pub trait Attribute: Sized {
         }
     }
 
+    /// 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.)
@@ -90,9 +95,11 @@ impl AttrStore {
             *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);
@@ -137,6 +144,7 @@ pub struct FreeSpace {
 
 /// A single GPU buffer containing linear arrays of attributes.
 pub struct AttrPool {
+    buffer: wgpu::Buffer,
     id: AttrId,
     pool_id: usize,
     layout: AttrLayout,
@@ -152,11 +160,20 @@ impl PartialEq for AttrPool {
 }
 
 impl AttrPool {
-    pub fn new(id: AttrId, pool_id: usize, info: AttrInfo) -> Self {
+    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,
@@ -290,6 +307,15 @@ impl AttrPool {
         let alloc = self.allocs.get(key).ok_or(PoolError::InvalidIndex)?;
         Ok(alloc.offset)
     }
+
+    /// Gets a [CopyDest] for an allocation, by key.
+    pub fn get_copy_dest(&self, key: usize) -> Result<CopyDest, PoolError> {
+        let offset = self.get_offset(key)?;
+        Ok(CopyDest {
+            buffer: &self.buffer,
+            offset,
+        })
+    }
 }
 
 /// The number of attributes a mesh can have before they're moved to the heap.
@@ -382,7 +408,7 @@ impl MeshPool {
                         None => return Err(PoolError::AttrUnregistered),
                     };
 
-                    let pool = AttrPool::new(id, 0, *info);
+                    let pool = AttrPool::new(&self.device, id, 0, *info);
                     self.pools.insert(id, vec![pool]);
                     self.pools.get_mut(&id).unwrap()
                 }
@@ -417,6 +443,23 @@ impl MeshPool {
         Ok(handle)
     }
 
+    pub fn flush(&mut self, commands: &mut wgpu::CommandEncoder) {
+        let get_dst = |target: &AttrAllocKey| {
+            self.pools
+                .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,
diff --git a/src/staging.rs b/src/staging.rs
index 32ac2cd..1f1e9b7 100644
--- a/src/staging.rs
+++ b/src/staging.rs
@@ -28,10 +28,10 @@ impl<T: Clone> StagingPool<T> {
         }
     }
 
-    pub fn flush(
+    pub fn flush<'a>(
         &mut self,
         cmd: &mut wgpu::CommandEncoder,
-        get_dst: impl Fn(&T) -> CopyDest<'_>,
+        get_dst: impl Fn(&T) -> CopyDest<'a>,
         on_complete: impl Fn(T),
     ) {
         if self.spillover.is_empty() {
@@ -115,7 +115,11 @@ impl<T: Clone> CopyBuffer<T> {
 
         if dst_size >= size {
             dst[0..size].copy_from_slice(&data);
-            let info = CopyInfo { target, offset, size };
+            let info = CopyInfo {
+                target,
+                offset,
+                size,
+            };
             (info, None)
         } else {
             let remainder = data.split_off(dst_size);