Add buffer reuse and copy on resize. (#22559)

Makes `ShaderStorageBuffer` have parity with `Image` vis-a-vis
#22552.

We apply the same pattern used there here.

Author: tychedelia

crates/bevy_render/src/storage.rs

@@ -1,7 +1,7 @@
 use crate::{
     render_asset::{AssetExtractionError, PrepareAssetError, RenderAsset, RenderAssetPlugin},
     render_resource::{Buffer, BufferUsages},
-    renderer::RenderDevice,
+    renderer::{RenderDevice, RenderQueue},
 };
 use bevy_app::{App, Plugin};
 use bevy_asset::{Asset, AssetApp, AssetId, RenderAssetUsages};
@@ -34,6 +34,8 @@ pub struct ShaderBuffer {
     pub buffer_description: wgpu::BufferDescriptor<'static>,
     /// The asset usage of the storage buffer.
     pub asset_usage: RenderAssetUsages,
+    /// Whether this buffer should be copied on the GPU when resized.
+    pub copy_on_resize: bool,
 }
 
 impl Default for ShaderBuffer {
@@ -43,10 +45,11 @@ impl Default for ShaderBuffer {
             buffer_description: wgpu::BufferDescriptor {
                 label: None,
                 size: 0,
-                usage: BufferUsages::STORAGE,
+                usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC | BufferUsages::COPY_DST,
                 mapped_at_creation: false,
             },
             asset_usage: RenderAssetUsages::default(),
+            copy_on_resize: false,
         }
     }
 }
@@ -84,6 +87,30 @@ impl ShaderBuffer {
         wrapper.write(&value).unwrap();
         self.data = Some(wrapper.into_inner());
     }
+
+    /// Resizes the buffer to the new size.
+    ///
+    /// If CPU data is present, it will be truncated or zero-extended.
+    /// Does not preserve GPU data when the descriptor changes.
+    pub fn resize(&mut self, size: u64) {
+        self.buffer_description.size = size;
+        if let Some(ref mut data) = self.data {
+            data.resize(size as usize, 0);
+        }
+    }
+
+    /// Resizes the buffer to the new size, preserving existing data.
+    ///
+    /// If CPU data is present, it will be truncated or zero-extended.
+    /// If no CPU data is present, sets `copy_on_resize` to preserve GPU data.
+    pub fn resize_in_place(&mut self, size: u64) {
+        self.buffer_description.size = size;
+        if let Some(ref mut data) = self.data {
+            data.resize(size as usize, 0);
+        } else {
+            self.copy_on_resize = true;
+        }
+    }
 }
 
 impl<T> From<T> for ShaderBuffer
@@ -101,12 +128,13 @@ where
 /// A storage buffer that is prepared as a [`RenderAsset`] and uploaded to the GPU.
 pub struct GpuShaderBuffer {
     pub buffer: Buffer,
+    pub buffer_descriptor: wgpu::BufferDescriptor<'static>,
     pub had_data: bool,
 }
 
 impl RenderAsset for GpuShaderBuffer {
     type SourceAsset = ShaderBuffer;
-    type Param = SRes<RenderDevice>;
+    type Param = (SRes<RenderDevice>, SRes<RenderQueue>);
 
     fn asset_usage(source_asset: &Self::SourceAsset) -> RenderAssetUsages {
         source_asset.asset_usage
@@ -131,28 +159,59 @@ impl RenderAsset for GpuShaderBuffer {
     fn prepare_asset(
         source_asset: Self::SourceAsset,
         _: AssetId<Self::SourceAsset>,
-        render_device: &mut SystemParamItem<Self::Param>,
-        _: Option<&Self>,
+        (render_device, render_queue): &mut SystemParamItem<Self::Param>,
+        previous_asset: Option<&Self>,
     ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
-        match source_asset.data {
-            Some(data) => {
-                let buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
-                    label: source_asset.buffer_description.label,
-                    contents: &data,
-                    usage: source_asset.buffer_description.usage,
-                });
-                Ok(GpuShaderBuffer {
-                    buffer,
-                    had_data: true,
-                })
+        let had_data = source_asset.data.is_some();
+
+        let buffer = if let Some(prev) = previous_asset
+            && prev.buffer_descriptor == source_asset.buffer_description
+            && source_asset
+                .buffer_description
+                .usage
+                .contains(BufferUsages::COPY_DST)
+        {
+            if let Some(ref data) = source_asset.data {
+                render_queue.write_buffer(&prev.buffer, 0, data);
             }
-            None => {
-                let buffer = render_device.create_buffer(&source_asset.buffer_description);
-                Ok(GpuShaderBuffer {
-                    buffer,
-                    had_data: false,
-                })
+            prev.buffer.clone()
+        } else if let Some(ref data) = source_asset.data {
+            render_device.create_buffer_with_data(&BufferInitDescriptor {
+                label: source_asset.buffer_description.label,
+                contents: data,
+                usage: source_asset.buffer_description.usage,
+            })
+        } else {
+            let new_buffer = render_device.create_buffer(&source_asset.buffer_description);
+            if source_asset.copy_on_resize
+                && let Some(previous) = previous_asset
+                && previous
+                    .buffer_descriptor
+                    .usage
+                    .contains(BufferUsages::COPY_SRC)
+                && source_asset
+                    .buffer_description
+                    .usage
+                    .contains(BufferUsages::COPY_DST)
+            {
+                let copy_size = source_asset
+                    .buffer_description
+                    .size
+                    .min(previous.buffer_descriptor.size);
+                let mut encoder =
+                    render_device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
+                        label: Some("copy_buffer_on_resize"),
+                    });
+                encoder.copy_buffer_to_buffer(&previous.buffer, 0, &new_buffer, 0, copy_size);
+                render_queue.submit([encoder.finish()]);
             }
-        }
+            new_buffer
+        };
+
+        Ok(GpuShaderBuffer {
+            buffer,
+            buffer_descriptor: source_asset.buffer_description,
+            had_data,
+        })
     }
 }