allow reading only a range of the buffer

crates/bevy_render/src/gpu_readback.rs

@@ -77,7 +77,10 @@ impl Plugin for GpuReadbackPlugin {
 #[derive(Component, ExtractComponent, Clone, Debug)]
 pub enum Readback {
     Texture(Handle<Image>),
-    Buffer(Handle<ShaderStorageBuffer>),
+    Buffer {
+        buffer: Handle<ShaderStorageBuffer>,
+        start_size: Option<(u64, u64)>,
+    },
 }
 
 impl Readback {
@@ -86,9 +89,21 @@ impl Readback {
         Self::Texture(image)
     }
 
-    /// Create a readback component for a buffer using the given handle.
+    /// Create a readback component for a full buffer using the given handle.
     pub fn buffer(buffer: Handle<ShaderStorageBuffer>) -> Self {
-        Self::Buffer(buffer)
+        Self::Buffer {
+            buffer,
+            start_size: None,
+        }
+    }
+
+    /// Create a readback component for a buffer range using the given handle, a start offset in bytes
+    /// and a number of bytes to read.
+    pub fn buffer_range(buffer: Handle<ShaderStorageBuffer>, start: u64, size: u64) -> Self {
+        Self::Buffer {
+            buffer,
+            start_size: Some((start, size)),
+        }
     }
 }
 
@@ -193,9 +208,8 @@ enum ReadbackSource {
         size: Extent3d,
     },
     Buffer {
-        src_start: u64,
-        dst_start: u64,
         buffer: Buffer,
+        start_size: Option<(u64, u64)>,
     },
 }
 
@@ -266,16 +280,17 @@ fn prepare_buffers(
                     });
                 }
             }
-            Readback::Buffer(buffer) => {
+            Readback::Buffer { buffer, start_size } => {
                 if let Some(ssbo) = ssbos.get(buffer) {
-                    let size = ssbo.buffer.size();
+                    let size = start_size
+                        .map(|(_, size)| size.min(ssbo.buffer.size()))
+                        .unwrap_or(ssbo.buffer.size());
                     let buffer = buffer_pool.get(&render_device, size);
                     let (tx, rx) = async_channel::bounded(1);
                     readbacks.requested.push(GpuReadback {
                         entity: entity.id(),
                         src: ReadbackSource::Buffer {
-                            src_start: 0,
-                            dst_start: 0,
+                            start_size: *start_size,
                             buffer: ssbo.buffer.clone(),
                         },
                         buffer,
@@ -306,18 +321,9 @@ pub(crate) fn submit_readback_commands(world: &World, command_encoder: &mut Comm
                     *size,
                 );
             }
-            ReadbackSource::Buffer {
-                src_start,
-                dst_start,
-                buffer,
-            } => {
-                command_encoder.copy_buffer_to_buffer(
-                    buffer,
-                    *src_start,
-                    &readback.buffer,
-                    *dst_start,
-                    buffer.size(),
-                );
+            ReadbackSource::Buffer { buffer, start_size } => {
+                let (src_start, size) = start_size.unwrap_or((0, buffer.size()));
+                command_encoder.copy_buffer_to_buffer(buffer, src_start, &readback.buffer, 0, size);
             }
         }
     }