From 21e1893c4f39bfd33967145810479d56b03486ce Mon Sep 17 00:00:00 2001
From: James Liu <contact@jamessliu.com>
Date: Tue, 4 Apr 2023 13:12:31 -0700
Subject: [PATCH] Remove capacity fields from all Buffer wrapper types (#8301)

# Objective
While working on #8299, I noticed that we're using a `capacity` field,
even though `wgpu::Buffer` exposes a `size` accessor that does the same
thing.

## Solution
Remove it from all buffer wrappers. Use `wgpu::Buffer::size` instead.
Default to 0 if no buffer has been allocated yet.
---
 .../src/render_resource/storage_buffer.rs       | 17 ++++++-----------
 .../src/render_resource/uniform_buffer.rs       |  8 +++-----
 2 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/crates/bevy_render/src/render_resource/storage_buffer.rs b/crates/bevy_render/src/render_resource/storage_buffer.rs
index 00f531c230d57..26d3797eded5c 100644
--- a/crates/bevy_render/src/render_resource/storage_buffer.rs
+++ b/crates/bevy_render/src/render_resource/storage_buffer.rs
@@ -33,7 +33,6 @@ pub struct StorageBuffer<T: ShaderType> {
     value: T,
     scratch: StorageBufferWrapper<Vec<u8>>,
     buffer: Option<Buffer>,
-    capacity: usize,
     label: Option<String>,
     changed: bool,
     buffer_usage: BufferUsages,
@@ -45,7 +44,6 @@ impl<T: ShaderType> From<T> for StorageBuffer<T> {
             value,
             scratch: StorageBufferWrapper::new(Vec::new()),
             buffer: None,
-            capacity: 0,
             label: None,
             changed: false,
             buffer_usage: BufferUsages::COPY_DST | BufferUsages::STORAGE,
@@ -59,7 +57,6 @@ impl<T: ShaderType + Default> Default for StorageBuffer<T> {
             value: T::default(),
             scratch: StorageBufferWrapper::new(Vec::new()),
             buffer: None,
-            capacity: 0,
             label: None,
             changed: false,
             buffer_usage: BufferUsages::COPY_DST | BufferUsages::STORAGE,
@@ -124,15 +121,15 @@ impl<T: ShaderType + WriteInto> StorageBuffer<T> {
     pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) {
         self.scratch.write(&self.value).unwrap();
 
-        let size = self.scratch.as_ref().len();
+        let capacity = self.buffer.as_deref().map(wgpu::Buffer::size).unwrap_or(0);
+        let size = self.scratch.as_ref().len() as u64;
 
-        if self.capacity < size || self.changed {
+        if capacity < size || self.changed {
             self.buffer = Some(device.create_buffer_with_data(&BufferInitDescriptor {
                 label: self.label.as_deref(),
                 usage: self.buffer_usage,
                 contents: self.scratch.as_ref(),
             }));
-            self.capacity = size;
             self.changed = false;
         } else if let Some(buffer) = &self.buffer {
             queue.write_buffer(buffer, 0, self.scratch.as_ref());
@@ -164,7 +161,6 @@ impl<T: ShaderType + WriteInto> StorageBuffer<T> {
 pub struct DynamicStorageBuffer<T: ShaderType> {
     scratch: DynamicStorageBufferWrapper<Vec<u8>>,
     buffer: Option<Buffer>,
-    capacity: usize,
     label: Option<String>,
     changed: bool,
     buffer_usage: BufferUsages,
@@ -176,7 +172,6 @@ impl<T: ShaderType> Default for DynamicStorageBuffer<T> {
         Self {
             scratch: DynamicStorageBufferWrapper::new(Vec::new()),
             buffer: None,
-            capacity: 0,
             label: None,
             changed: false,
             buffer_usage: BufferUsages::COPY_DST | BufferUsages::STORAGE,
@@ -236,15 +231,15 @@ impl<T: ShaderType + WriteInto> DynamicStorageBuffer<T> {
 
     #[inline]
     pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) {
-        let size = self.scratch.as_ref().len();
+        let capacity = self.buffer.as_deref().map(wgpu::Buffer::size).unwrap_or(0);
+        let size = self.scratch.as_ref().len() as u64;
 
-        if self.capacity < size || self.changed {
+        if capacity < size || self.changed {
             self.buffer = Some(device.create_buffer_with_data(&BufferInitDescriptor {
                 label: self.label.as_deref(),
                 usage: self.buffer_usage,
                 contents: self.scratch.as_ref(),
             }));
-            self.capacity = size;
             self.changed = false;
         } else if let Some(buffer) = &self.buffer {
             queue.write_buffer(buffer, 0, self.scratch.as_ref());
diff --git a/crates/bevy_render/src/render_resource/uniform_buffer.rs b/crates/bevy_render/src/render_resource/uniform_buffer.rs
index ac79433a752e5..137432c062bfb 100644
--- a/crates/bevy_render/src/render_resource/uniform_buffer.rs
+++ b/crates/bevy_render/src/render_resource/uniform_buffer.rs
@@ -157,7 +157,6 @@ impl<T: ShaderType + WriteInto> UniformBuffer<T> {
 pub struct DynamicUniformBuffer<T: ShaderType> {
     scratch: DynamicUniformBufferWrapper<Vec<u8>>,
     buffer: Option<Buffer>,
-    capacity: usize,
     label: Option<String>,
     changed: bool,
     buffer_usage: BufferUsages,
@@ -169,7 +168,6 @@ impl<T: ShaderType> Default for DynamicUniformBuffer<T> {
         Self {
             scratch: DynamicUniformBufferWrapper::new(Vec::new()),
             buffer: None,
-            capacity: 0,
             label: None,
             changed: false,
             buffer_usage: BufferUsages::COPY_DST | BufferUsages::UNIFORM,
@@ -235,15 +233,15 @@ impl<T: ShaderType + WriteInto> DynamicUniformBuffer<T> {
     /// allocated does not have enough capacity, a new GPU-side buffer is created.
     #[inline]
     pub fn write_buffer(&mut self, device: &RenderDevice, queue: &RenderQueue) {
-        let size = self.scratch.as_ref().len();
+        let capacity = self.buffer.as_deref().map(wgpu::Buffer::size).unwrap_or(0);
+        let size = self.scratch.as_ref().len() as u64;
 
-        if self.capacity < size || self.changed {
+        if capacity < size || self.changed {
             self.buffer = Some(device.create_buffer_with_data(&BufferInitDescriptor {
                 label: self.label.as_deref(),
                 usage: self.buffer_usage,
                 contents: self.scratch.as_ref(),
             }));
-            self.capacity = size;
             self.changed = false;
         } else if let Some(buffer) = &self.buffer {
             queue.write_buffer(buffer, 0, self.scratch.as_ref());