Add support for reading from mapped buffers

crates/bevy_pbr/src/render_graph/lights_node.rs

@@ -9,7 +9,7 @@ use bevy_ecs::{
 use bevy_render::{
     render_graph::{CommandQueue, Node, ResourceSlots, SystemNode},
     renderer::{
-        BufferId, BufferInfo, BufferUsage, RenderContext, RenderResourceBinding,
+        BufferId, BufferInfo, BufferMapMode, BufferUsage, RenderContext, RenderResourceBinding,
         RenderResourceBindings, RenderResourceContext,
     },
 };
@@ -103,7 +103,7 @@ pub fn lights_node_system(
             return;
         }
 
-        render_resource_context.map_buffer(staging_buffer);
+        render_resource_context.map_buffer(staging_buffer, BufferMapMode::Write);
     } else {
         let buffer = render_resource_context.create_buffer(BufferInfo {
             size: max_light_uniform_size,

crates/bevy_render/src/render_graph/nodes/camera_node.rs

@@ -2,7 +2,7 @@ use crate::{
     camera::{ActiveCameras, Camera},
     render_graph::{CommandQueue, Node, ResourceSlots, SystemNode},
     renderer::{
-        BufferId, BufferInfo, BufferUsage, RenderContext, RenderResourceBinding,
+        BufferId, BufferInfo, BufferMapMode, BufferUsage, RenderContext, RenderResourceBinding,
         RenderResourceBindings, RenderResourceContext,
     },
 };
@@ -87,7 +87,7 @@ pub fn camera_node_system(
     };
 
     let staging_buffer = if let Some(staging_buffer) = state.staging_buffer {
-        render_resource_context.map_buffer(staging_buffer);
+        render_resource_context.map_buffer(staging_buffer, BufferMapMode::Write);
         staging_buffer
     } else {
         let size = std::mem::size_of::<[[f32; 4]; 4]>();

crates/bevy_render/src/render_graph/nodes/render_resources_node.rs

@@ -3,7 +3,7 @@ use crate::{
     prelude::Visible,
     render_graph::{CommandQueue, Node, ResourceSlots, SystemNode},
     renderer::{
-        self, BufferInfo, BufferUsage, RenderContext, RenderResourceBinding,
+        self, BufferInfo, BufferMapMode, BufferUsage, RenderContext, RenderResourceBinding,
         RenderResourceBindings, RenderResourceContext, RenderResourceHints,
     },
     texture,
@@ -490,7 +490,7 @@ fn render_resources_node_system<T: RenderResources>(
     uniform_buffer_arrays.resize_staging_buffer(render_resource_context);
 
     if let Some(staging_buffer) = state.uniform_buffer_arrays.staging_buffer {
-        render_resource_context.map_buffer(staging_buffer);
+        render_resource_context.map_buffer(staging_buffer, BufferMapMode::Write);
         render_resource_context.write_mapped_buffer(
             staging_buffer,
             0..state.uniform_buffer_arrays.staging_buffer_size as u64,
@@ -703,7 +703,7 @@ fn asset_render_resources_node_system<T: RenderResources + Asset>(
     uniform_buffer_arrays.resize_staging_buffer(render_resource_context);
 
     if let Some(staging_buffer) = state.uniform_buffer_arrays.staging_buffer {
-        render_resource_context.map_buffer(staging_buffer);
+        render_resource_context.map_buffer(staging_buffer, BufferMapMode::Write);
         render_resource_context.write_mapped_buffer(
             staging_buffer,
             0..state.uniform_buffer_arrays.staging_buffer_size as u64,

crates/bevy_render/src/renderer/headless_render_resource_context.rs

@@ -1,7 +1,9 @@
 use super::RenderResourceContext;
 use crate::{
     pipeline::{BindGroupDescriptorId, PipelineDescriptor},
-    renderer::{BindGroup, BufferId, BufferInfo, RenderResourceId, SamplerId, TextureId},
+    renderer::{
+        BindGroup, BufferId, BufferInfo, BufferMapMode, RenderResourceId, SamplerId, TextureId,
+    },
     shader::{Shader, ShaderError},
     texture::{SamplerDescriptor, TextureDescriptor},
 };
@@ -66,7 +68,18 @@ impl RenderResourceContext for HeadlessRenderResourceContext {
         write(&mut buffer, self);
     }
 
-    fn map_buffer(&self, _id: BufferId) {}
+    fn read_mapped_buffer(
+        &self,
+        id: BufferId,
+        _range: Range<u64>,
+        read: &dyn Fn(&[u8], &dyn RenderResourceContext),
+    ) {
+        let size = self.buffer_info.read().get(&id).unwrap().size;
+        let buffer = vec![0; size];
+        read(&buffer, self);
+    }
+
+    fn map_buffer(&self, _id: BufferId, _mode: BufferMapMode) {}
 
     fn unmap_buffer(&self, _id: BufferId) {}
 

crates/bevy_render/src/renderer/render_resource/buffer.rs

@@ -43,3 +43,9 @@ bitflags::bitflags! {
         const INDIRECT = 256;
     }
 }
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum BufferMapMode {
+    Read,
+    Write,
+}

crates/bevy_render/src/renderer/render_resource/shared_buffers.rs

@@ -1,7 +1,7 @@
 use super::{BufferId, BufferInfo, RenderResource, RenderResourceBinding};
 use crate::{
     render_graph::CommandQueue,
-    renderer::{BufferUsage, RenderContext, RenderResourceContext},
+    renderer::{BufferMapMode, BufferUsage, RenderContext, RenderResourceContext},
 };
 use bevy_ecs::{Res, ResMut};
 
@@ -115,7 +115,7 @@ impl SharedBuffers {
         }
 
         if let Some(staging_buffer) = self.staging_buffer {
-            render_resource_context.map_buffer(staging_buffer);
+            render_resource_context.map_buffer(staging_buffer, BufferMapMode::Write);
         }
     }
 

crates/bevy_render/src/renderer/render_resource_context.rs

@@ -1,6 +1,8 @@
 use crate::{
     pipeline::{BindGroupDescriptorId, PipelineDescriptor, PipelineLayout},
-    renderer::{BindGroup, BufferId, BufferInfo, RenderResourceId, SamplerId, TextureId},
+    renderer::{
+        BindGroup, BufferId, BufferInfo, BufferMapMode, RenderResourceId, SamplerId, TextureId,
+    },
     shader::{Shader, ShaderError, ShaderLayout, ShaderStages},
     texture::{SamplerDescriptor, TextureDescriptor},
 };
@@ -24,7 +26,13 @@ pub trait RenderResourceContext: Downcast + Send + Sync + 'static {
         range: Range<u64>,
         write: &mut dyn FnMut(&mut [u8], &dyn RenderResourceContext),
     );
-    fn map_buffer(&self, id: BufferId);
+    fn read_mapped_buffer(
+        &self,
+        id: BufferId,
+        range: Range<u64>,
+        read: &dyn Fn(&[u8], &dyn RenderResourceContext),
+    );
+    fn map_buffer(&self, id: BufferId, mode: BufferMapMode);
     fn unmap_buffer(&self, id: BufferId);
     fn create_buffer_with_data(&self, buffer_info: BufferInfo, data: &[u8]) -> BufferId;
     fn create_shader_module(&self, shader_handle: &Handle<Shader>, shaders: &Assets<Shader>);

crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs

@@ -9,8 +9,8 @@ use bevy_render::{
         BindGroupDescriptor, BindGroupDescriptorId, BindingShaderStage, PipelineDescriptor,
     },
     renderer::{
-        BindGroup, BufferId, BufferInfo, RenderResourceBinding, RenderResourceContext,
-        RenderResourceId, SamplerId, TextureId,
+        BindGroup, BufferId, BufferInfo, BufferMapMode, RenderResourceBinding,
+        RenderResourceContext, RenderResourceId, SamplerId, TextureId,
     },
     shader::{glsl_to_spirv, Shader, ShaderError, ShaderSource},
     texture::{Extent3d, SamplerDescriptor, TextureDescriptor},
@@ -622,11 +622,30 @@ impl RenderResourceContext for WgpuRenderResourceContext {
         write(&mut data, self);
     }
 
-    fn map_buffer(&self, id: BufferId) {
+    fn read_mapped_buffer(
+        &self,
+        id: BufferId,
+        range: Range<u64>,
+        read: &dyn Fn(&[u8], &dyn RenderResourceContext),
+    ) {
+        let buffer = {
+            let buffers = self.resources.buffers.read();
+            buffers.get(&id).unwrap().clone()
+        };
+        let buffer_slice = buffer.slice(range);
+        let data = buffer_slice.get_mapped_range();
+        read(&data, self);
+    }
+
+    fn map_buffer(&self, id: BufferId, mode: BufferMapMode) {
         let buffers = self.resources.buffers.read();
         let buffer = buffers.get(&id).unwrap();
         let buffer_slice = buffer.slice(..);
-        let data = buffer_slice.map_async(wgpu::MapMode::Write);
+        let wgpu_mode = match mode {
+            BufferMapMode::Read => wgpu::MapMode::Read,
+            BufferMapMode::Write => wgpu::MapMode::Write,
+        };
+        let data = buffer_slice.map_async(wgpu_mode);
         self.device.poll(wgpu::Maintain::Wait);
         if future::block_on(data).is_err() {
             panic!("Failed to map buffer to host.");