use bytemuck::NoUninit;
use std::marker::PhantomData;
use std::mem;
use std::num::{NonZero, NonZeroU64};
use std::ops::Range;
use vek::az::UnwrappedAs;
use wgpu::util::{BufferInitDescriptor, DeviceExt};
use wgpu::{
BindGroupEntry, BindGroupLayoutEntry, BindingResource, BindingType, Buffer, BufferBinding,
BufferBindingType, BufferDescriptor, BufferSlice, BufferUsages, Device, Queue, ShaderStages,
};
pub struct GpuBuffer<T> {
buffer: Buffer,
_phantom: PhantomData<T>,
}
pub struct GpuArrayBuffer<T> {
buffer: Buffer,
capacity: usize,
length: usize,
_phantom: PhantomData<T>,
}
pub struct GpuVecBuffer<T> {
label: Option<&'static str>,
usage: BufferUsages,
buffer: Buffer,
length: u64,
size_limit: u64,
_phantom: PhantomData<T>,
}
pub trait GpuDeviceExt {
fn buffer<T: bytemuck::NoUninit>(
&self,
data: &T,
label: Option<&'static str>,
usage: BufferUsages,
) -> GpuBuffer<T>;
fn array_buffer<T: bytemuck::NoUninit>(
&self,
data: &[T],
label: Option<&'static str>,
usage: BufferUsages,
) -> GpuArrayBuffer<T>;
fn vec_buffer<T: bytemuck::NoUninit>(
&self,
label: Option<&'static str>,
usage: BufferUsages,
) -> GpuVecBuffer<T>;
}
impl GpuDeviceExt for Device {
fn buffer<T: bytemuck::NoUninit>(
&self,
data: &T,
label: Option<&'static str>,
usage: BufferUsages,
) -> GpuBuffer<T> {
GpuBuffer {
buffer: self.create_buffer_init(&BufferInitDescriptor {
label,
contents: bytemuck::bytes_of(data),
usage: usage | BufferUsages::COPY_DST,
}),
_phantom: PhantomData,
}
}
fn array_buffer<T: bytemuck::NoUninit>(
&self,
data: &[T],
label: Option<&'static str>,
usage: BufferUsages,
) -> GpuArrayBuffer<T> {
GpuArrayBuffer {
buffer: self.create_buffer_init(&BufferInitDescriptor {
label,
contents: bytemuck::cast_slice(data),
usage: usage | BufferUsages::COPY_DST,
}),
capacity: data.len(),
length: data.len(),
_phantom: PhantomData,
}
}
fn vec_buffer<T: NoUninit>(
&self,
label: Option<&'static str>,
usage: BufferUsages,
) -> GpuVecBuffer<T> {
const INIT_CAPACITY: usize = 8;
let init_size = (INIT_CAPACITY * mem::size_of::<T>()).unwrapped_as();
let usage = usage | BufferUsages::COPY_DST;
let size_limit = self.limits().max_buffer_size;
GpuVecBuffer {
label,
usage,
buffer: self.create_buffer(&BufferDescriptor {
label,
size: init_size,
usage,
mapped_at_creation: false,
}),
length: 0,
size_limit,
_phantom: PhantomData,
}
}
}
impl<T: bytemuck::NoUninit> GpuBuffer<T> {
pub fn update(&self, data: &T, queue: &Queue) {
queue.write_buffer(&self.buffer, 0, bytemuck::bytes_of(data))
}
pub fn slice(&self) -> BufferSlice<'_> {
self.buffer.slice(..)
}
pub fn full_binding(&self) -> BufferBinding<'_> {
BufferBinding {
buffer: &self.buffer,
offset: 0,
size: None,
}
}
pub fn bind_group_entry(&self, binding: u32) -> BindGroupEntry<'_> {
BindGroupEntry {
binding,
resource: BindingResource::Buffer(self.full_binding()),
}
}
pub fn bind_group_layout_entry(binding: u32, has_dynamic_offset: bool) -> BindGroupLayoutEntry {
BindGroupLayoutEntry {
binding,
visibility: ShaderStages::VERTEX,
ty: BindingType::Buffer {
ty: BufferBindingType::Uniform,
has_dynamic_offset,
min_binding_size: Some(
NonZeroU64::new(mem::size_of::<T>().unwrapped_as()).unwrap(),
),
},
count: None,
}
}
}
impl<T: bytemuck::NoUninit> GpuArrayBuffer<T> {
const ELEMENT_SIZE: usize = mem::size_of::<T>();
pub fn update(&mut self, data: &[T], queue: &Queue) {
assert!(data.len() <= self.capacity);
queue.write_buffer(&self.buffer, 0, bytemuck::cast_slice(data));
self.length = data.len();
}
pub fn capacity(&self) -> usize {
self.capacity
}
pub fn len(&self) -> usize {
self.length
}
pub fn full_slice(&self) -> BufferSlice<'_> {
self.buffer
.slice(..(self.length * Self::ELEMENT_SIZE) as u64)
}
pub fn slice(&self, range: Range<usize>) -> BufferSlice<'_> {
let size = mem::size_of::<T>() as u64;
self.buffer
.slice(range.start as u64 * size..range.end as u64 * size)
}
pub fn first_element_binding(&self) -> BufferBinding<'_> {
BufferBinding {
buffer: &self.buffer,
offset: 0,
size: Some(NonZero::new(mem::size_of::<T>() as u64).unwrap()),
}
}
pub fn bind_group_entry(&self, binding: u32) -> BindGroupEntry<'_> {
BindGroupEntry {
binding,
resource: BindingResource::Buffer(BufferBinding {
buffer: &self.buffer,
offset: 0,
size: None,
}),
}
}
pub fn bind_group_layout_entry(&self, binding: u32) -> BindGroupLayoutEntry {
BindGroupLayoutEntry {
binding,
visibility: ShaderStages::VERTEX,
ty: BindingType::Buffer {
ty: BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: Some(NonZeroU64::new(0).unwrap()),
},
count: None,
}
}
}
impl<T: bytemuck::NoUninit> GpuVecBuffer<T> {
pub fn clear(&mut self) {
self.length = 0;
}
pub fn is_empty(&self) -> bool {
self.length == 0
}
pub fn len(&self) -> usize {
self.length as usize / mem::size_of::<T>()
}
pub fn upload<'a>(
&mut self,
data: &'a [T],
device: &Device,
queue: &Queue,
) -> Result<bool, &'a [T]> {
let required_size: u64 = mem::size_of_val(data).unwrapped_as();
let reallocate = required_size > self.buffer.size();
if reallocate {
let reallocate_size = required_size.next_power_of_two();
if reallocate_size > self.size_limit {
let max_allocate_size = self.size_limit.next_power_of_two() / 2;
if self.buffer.size() != max_allocate_size {
self.buffer = device.create_buffer(&BufferDescriptor {
label: self.label,
size: max_allocate_size,
usage: self.usage,
mapped_at_creation: false,
});
}
let max_elements = self.buffer.size() as usize / mem::size_of::<T>();
let (fit, leftover) = data.split_at(max_elements);
queue.write_buffer(&self.buffer, 0, bytemuck::cast_slice(fit));
self.length = mem::size_of_val(fit).unwrapped_as();
return Err(leftover);
}
self.buffer = device.create_buffer(&BufferDescriptor {
label: self.label,
size: required_size.next_power_of_two(),
usage: self.usage,
mapped_at_creation: false,
});
}
queue.write_buffer(&self.buffer, 0, bytemuck::cast_slice(data));
self.length = required_size;
Ok(reallocate)
}
pub fn full_slice(&self) -> BufferSlice<'_> {
self.buffer.slice(..self.length)
}
pub fn slice(&self, range: Range<usize>) -> BufferSlice<'_> {
let size = mem::size_of::<T>() as u64;
self.buffer
.slice(range.start as u64 * size..range.end as u64 * size)
}
pub fn binding(&self, offset: u64, size: Option<u64>) -> BufferBinding<'_> {
let t_size = mem::size_of::<T>() as u64;
BufferBinding {
buffer: &self.buffer,
offset: offset * t_size,
size: size.map(|size| NonZeroU64::new(size * t_size).unwrap()),
}
}
}