From fd12726efb33d2fac18ce7914f5840b0709d3fbc Mon Sep 17 00:00:00 2001
From: mars <mars@tebibyte.media>
Date: Sat, 7 May 2022 21:49:30 -0600
Subject: [PATCH] Add GpuVec

---
 src/gpu.rs | 112 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs |   1 +
 2 files changed, 113 insertions(+)
 create mode 100644 src/gpu.rs

diff --git a/src/gpu.rs b/src/gpu.rs
new file mode 100644
index 0000000..dde3ec1
--- /dev/null
+++ b/src/gpu.rs
@@ -0,0 +1,112 @@
+use bytemuck::Pod;
+use std::ops::{Deref, DerefMut};
+use std::sync::Arc;
+
+pub struct GpuVec<T> {
+    device: Arc<wgpu::Device>,
+    data: Vec<T>,
+    buffer: wgpu::Buffer,
+    capacity: usize,
+    usage: wgpu::BufferUsages,
+    label: Option<String>,
+    realign: Option<usize>,
+}
+
+impl<T: Clone + Pod> GpuVec<T> {
+    pub fn new(
+        device: Arc<wgpu::Device>,
+        usage: wgpu::BufferUsages,
+        initial_capacity: usize,
+        label: Option<String>,
+    ) -> Self {
+        let capacity_bytes = initial_capacity * std::mem::size_of::<T>();
+        let usage = usage | wgpu::BufferUsages::COPY_DST;
+        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
+            label: label.as_ref().map(|s| s.as_str()),
+            size: capacity_bytes as wgpu::BufferAddress,
+            usage,
+            mapped_at_creation: false,
+        });
+
+        let data = Vec::with_capacity(initial_capacity);
+
+        let realign = if usage.contains(wgpu::BufferUsages::STORAGE) {
+            Some(device.limits().min_storage_buffer_offset_alignment as usize)
+        } else if usage.contains(wgpu::BufferUsages::UNIFORM) {
+            Some(device.limits().min_uniform_buffer_offset_alignment as usize)
+        } else {
+            None
+        };
+
+        Self {
+            device,
+            data,
+            buffer,
+            capacity: capacity_bytes,
+            usage,
+            label,
+            realign,
+        }
+    }
+
+    pub fn write(&mut self, queue: &wgpu::Queue) -> &wgpu::Buffer {
+        let capacity_bytes = self.buf_offset(self.data.capacity());
+        if capacity_bytes != self.capacity && capacity_bytes != 0 {
+            self.capacity = capacity_bytes;
+            self.buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
+                label: self.label.as_ref().map(|s| s.as_str()),
+                size: capacity_bytes as wgpu::BufferAddress,
+                usage: self.usage,
+                mapped_at_creation: false,
+            });
+        }
+
+        let mut realigned = Vec::new();
+        let src_buf = if let Some(realign) = self.realign {
+            realigned.resize(self.data.len() * realign, 0);
+
+            for (index, elem) in self.data.iter().enumerate() {
+                let elem = [elem.clone()];
+                let casted: &[u8] = bytemuck::cast_slice(&elem);
+                let dst_offset = index * realign;
+                realigned[dst_offset..(dst_offset + casted.len())].copy_from_slice(casted);
+            }
+
+            realigned.as_slice()
+        } else {
+            bytemuck::cast_slice(self.data.as_slice())
+        };
+
+        queue.write_buffer(&self.buffer, 0, src_buf);
+
+        &self.buffer
+    }
+
+    pub fn buf_offset(&self, index: usize) -> usize {
+        if let Some(realign) = self.realign {
+            index * realign
+        } else {
+            index * std::mem::size_of::<T>()
+        }
+    }
+}
+
+impl<T: Pod> AsRef<wgpu::Buffer> for GpuVec<T> {
+    fn as_ref(&self) -> &wgpu::Buffer {
+        &self.buffer
+    }
+}
+
+impl<T: Pod> Deref for GpuVec<T> {
+    type Target = Vec<T>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.data
+    }
+}
+
+impl<T: Pod> DerefMut for GpuVec<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.data
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index cf5a980..027b3d9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -6,6 +6,7 @@ use std::sync::Arc;
 use std::sync::Mutex;
 
 pub mod camera;
+pub mod gpu;
 pub mod mesh;
 pub mod pass;
 pub mod phase;