1- use std:: collections:: BTreeSet ;
1+ use std:: collections:: { BTreeSet , HashMap } ;
22
33use bevy:: asset:: RenderAssetUsages ;
44use bevy:: reflect:: PartialReflect ;
@@ -30,6 +30,10 @@ pub struct Buffer {
3030 pub handle : Handle < ShaderBuffer > ,
3131 pub readback_buffer : WgpuBuffer ,
3232 pub size : u64 ,
33+ /// True when `ShaderBuffer.data` reflects current GPU contents. Cleared
34+ /// when a pipeline that may write to the buffer runs; the next read or
35+ /// write must readback first.
36+ pub synced : bool ,
3337}
3438
3539fn readback_buffer ( device : & RenderDevice , size : u64 ) -> WgpuBuffer {
@@ -47,15 +51,16 @@ pub fn create_buffer(
4751 mut buffers : ResMut < Assets < ShaderBuffer > > ,
4852 render_device : Res < RenderDevice > ,
4953) -> Entity {
50- let handle = buffers. add ( ShaderBuffer :: with_size (
51- size as usize ,
54+ let handle = buffers. add ( ShaderBuffer :: new (
55+ & vec ! [ 0u8 ; size as usize ] ,
5256 RenderAssetUsages :: all ( ) ,
5357 ) ) ;
5458 commands
5559 . spawn ( Buffer {
5660 handle,
5761 readback_buffer : readback_buffer ( & render_device, size) ,
5862 size,
63+ synced : true ,
5964 } )
6065 . id ( )
6166}
@@ -73,30 +78,37 @@ pub fn create_buffer_with_data(
7378 handle,
7479 readback_buffer : readback_buffer ( & render_device, size) ,
7580 size,
81+ synced : true ,
7682 } )
7783 . id ( )
7884}
7985
80- pub fn write_buffer_gpu (
86+ /// Mutate the CPU-side data of a `ShaderBuffer` in place. Fires
87+ /// `AssetEvent::Modified` so Bevy's render-asset extract uploads the new
88+ /// contents to the GPU at the next sync point.
89+ pub fn write_buffer_cpu (
8190 In ( ( handle, offset, data) ) : In < ( Handle < ShaderBuffer > , u64 , Vec < u8 > ) > ,
82- gpu_buffers : Res < RenderAssets < GpuShaderBuffer > > ,
83- render_queue : Res < RenderQueue > ,
91+ mut buffers : ResMut < Assets < ShaderBuffer > > ,
8492) -> Result < ( ) > {
85- let gpu_buffer = & gpu_buffers
86- . get ( & handle)
87- . ok_or ( ProcessingError :: BufferNotFound ) ?
88- . buffer ;
89- render_queue. write_buffer ( gpu_buffer, offset, & data) ;
93+ let mut asset = buffers
94+ . get_mut ( & handle)
95+ . ok_or ( ProcessingError :: BufferNotFound ) ?;
96+ let dst = asset
97+ . data
98+ . as_mut ( )
99+ . ok_or ( ProcessingError :: BufferNotFound ) ?;
100+ let start = offset as usize ;
101+ let end = start + data. len ( ) ;
102+ dst[ start..end] . copy_from_slice ( & data) ;
90103 Ok ( ( ) )
91104}
92105
106+ /// Copy the GPU buffer back to CPU and return its full contents. Runs in the
107+ /// render world; the caller is responsible for writing the bytes back into
108+ /// `ShaderBuffer.data` via `Assets::get_mut_untracked` (avoiding spurious
109+ /// `AssetEvent::Modified`s, since this is a readback, not a stage-for-upload).
93110pub fn read_buffer_gpu (
94- In ( ( handle, readback_buffer, src_offset, len) ) : In < (
95- Handle < ShaderBuffer > ,
96- WgpuBuffer ,
97- u64 ,
98- u64 ,
99- ) > ,
111+ In ( ( handle, readback_buffer, size) ) : In < ( Handle < ShaderBuffer > , WgpuBuffer , u64 ) > ,
100112 gpu_buffers : Res < RenderAssets < GpuShaderBuffer > > ,
101113 render_device : Res < RenderDevice > ,
102114 render_queue : Res < RenderQueue > ,
@@ -107,10 +119,10 @@ pub fn read_buffer_gpu(
107119 . buffer ;
108120
109121 let mut encoder = render_device. create_command_encoder ( & CommandEncoderDescriptor :: default ( ) ) ;
110- encoder. copy_buffer_to_buffer ( gpu_buffer, src_offset , & readback_buffer, 0 , len ) ;
122+ encoder. copy_buffer_to_buffer ( gpu_buffer, 0 , & readback_buffer, 0 , size ) ;
111123 render_queue. submit ( std:: iter:: once ( encoder. finish ( ) ) ) ;
112124
113- let buffer_slice = readback_buffer. slice ( 0 ..len ) ;
125+ let buffer_slice = readback_buffer. slice ( 0 ..size ) ;
114126 let ( s, r) = crossbeam_channel:: bounded ( 1 ) ;
115127 buffer_slice. map_async ( MapMode :: Read , move |result| {
116128 let _ = s. send ( result) ;
@@ -122,10 +134,9 @@ pub fn read_buffer_gpu(
122134 . map_err ( |e| ProcessingError :: BufferMapError ( format ! ( "map channel closed: {e}" ) ) ) ?
123135 . map_err ( |e| ProcessingError :: BufferMapError ( format ! ( "map failed: {e}" ) ) ) ?;
124136
125- let data = buffer_slice. get_mapped_range ( ) . to_vec ( ) ;
137+ let bytes = buffer_slice. get_mapped_range ( ) . to_vec ( ) ;
126138 readback_buffer. unmap ( ) ;
127-
128- Ok ( data)
139+ Ok ( bytes)
129140}
130141
131142pub fn destroy_buffer ( In ( entity) : In < Entity > , mut commands : Commands ) -> Result < ( ) > {
@@ -139,6 +150,11 @@ pub struct Compute {
139150 pub entry_point : String ,
140151 pub pipeline_id : CachedComputePipelineId ,
141152 pub bind_group_layout_descriptors : Vec < ( u32 , BindGroupLayoutDescriptor ) > ,
153+ /// Buffer entities bound to this compute on a `read_write` storage param.
154+ /// Their CPU view of GPU data is invalidated after each dispatch so the
155+ /// next read/write does a readback. Read-only bindings don't need this
156+ /// since the dispatch can't mutate them.
157+ pub rw_buffers : HashMap < String , Entity > ,
142158}
143159
144160fn queue_pipeline (
@@ -240,6 +256,7 @@ pub fn create_compute(app: &mut App, shader_entity: Entity) -> Result<Entity> {
240256 entry_point,
241257 pipeline_id,
242258 bind_group_layout_descriptors,
259+ rw_buffers : HashMap :: new ( ) ,
243260 } )
244261 . id ( ) ) ;
245262 }
@@ -267,11 +284,16 @@ pub fn set_compute_property(
267284 . ok_or_else ( || ProcessingError :: UnknownShaderProperty ( name. clone ( ) ) ) ?;
268285
269286 match ( & value, category) {
270- ( ShaderValue :: Buffer ( buf_entity) , ParameterCategory :: Storage { .. } ) => {
287+ ( ShaderValue :: Buffer ( buf_entity) , ParameterCategory :: Storage { read_only } ) => {
271288 let buffer = p_buffers
272289 . get ( * buf_entity)
273290 . map_err ( |_| ProcessingError :: BufferNotFound ) ?;
274291 compute. shader . insert ( & name, buffer. handle . clone ( ) ) ;
292+ if read_only {
293+ compute. rw_buffers . remove ( & name) ;
294+ } else {
295+ compute. rw_buffers . insert ( name. clone ( ) , * buf_entity) ;
296+ }
275297 Ok ( ( ) )
276298 }
277299 ( ShaderValue :: Texture ( img_entity) , ParameterCategory :: Texture )
0 commit comments