From 8c9ca764f501a9a9ab9eb9b2ac9403b1f5a3f057 Mon Sep 17 00:00:00 2001
From: Renato Caldas <renato@calgera.com>
Date: Thu, 21 Jan 2021 00:06:06 +0000
Subject: [PATCH] Add support for reading from a mapped buffer.

---
 .../renderer/headless_render_resource_context.rs  | 11 +++++++++++
 .../src/renderer/render_resource_context.rs       |  6 ++++++
 .../src/renderer/wgpu_render_resource_context.rs  | 15 +++++++++++++++
 3 files changed, 32 insertions(+)

diff --git a/crates/bevy_render/src/renderer/headless_render_resource_context.rs b/crates/bevy_render/src/renderer/headless_render_resource_context.rs
index 985c2074afa2ef..02b9edf70efe0f 100644
--- a/crates/bevy_render/src/renderer/headless_render_resource_context.rs
+++ b/crates/bevy_render/src/renderer/headless_render_resource_context.rs
@@ -66,6 +66,17 @@ impl RenderResourceContext for HeadlessRenderResourceContext {
         write(&mut buffer, self);
     }
 
+    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) {}
diff --git a/crates/bevy_render/src/renderer/render_resource_context.rs b/crates/bevy_render/src/renderer/render_resource_context.rs
index 3ee8fb5e08a126..a17d50bdc8cb00 100644
--- a/crates/bevy_render/src/renderer/render_resource_context.rs
+++ b/crates/bevy_render/src/renderer/render_resource_context.rs
@@ -24,6 +24,12 @@ pub trait RenderResourceContext: Downcast + Send + Sync + 'static {
         range: Range<u64>,
         write: &mut dyn FnMut(&mut [u8], &dyn RenderResourceContext),
     );
+    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;
diff --git a/crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs b/crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs
index 27afefff69f829..a802a8e4ec93a0 100644
--- a/crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs
+++ b/crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs
@@ -622,6 +622,21 @@ impl RenderResourceContext for WgpuRenderResourceContext {
         write(&mut data, self);
     }
 
+   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();