.

tychedelia · tychedelia · commit 5cb139adc795 · 2026-05-02T13:21:50.000-07:00
diff --git a/crates/processing_ffi/src/lib.rs b/crates/processing_ffi/src/lib.rs
@@ -1912,11 +1912,6 @@ pub extern "C" fn processing_buffer_size(buf_id: u64) -> u64 {
 
 /// Read buffer contents into a caller-provided buffer.
 ///
-/// Returns the buffer's byte length. If the returned size is `<= out_len`, the
-/// data has been written to `out`; otherwise `out` is left untouched and the
-/// caller should reallocate and retry. Returns 0 if the buffer does not exist
-/// or the GPU readback failed (in which case the error is set).
-///
 /// # Safety
 /// - `out` must be valid for writes of `out_len` bytes (may be null if
 ///   `out_len == 0`, in which case this acts as a size query).
diff --git a/crates/processing_render/src/compute.rs b/crates/processing_render/src/compute.rs
@@ -1,4 +1,4 @@
-use std::collections::BTreeSet;
+use std::collections::{BTreeSet, HashMap};
 
 use bevy::asset::RenderAssetUsages;
 use bevy::reflect::PartialReflect;
@@ -30,6 +30,10 @@ pub struct Buffer {
     pub handle: Handle<ShaderBuffer>,
     pub readback_buffer: WgpuBuffer,
     pub size: u64,
+    /// True when `ShaderBuffer.data` reflects current GPU contents. Cleared
+    /// when a pipeline that may write to the buffer runs; the next read or
+    /// write must readback first.
+    pub synced: bool,
 }
 
 fn readback_buffer(device: &RenderDevice, size: u64) -> WgpuBuffer {
@@ -47,15 +51,16 @@ pub fn create_buffer(
     mut buffers: ResMut<Assets<ShaderBuffer>>,
     render_device: Res<RenderDevice>,
 ) -> Entity {
-    let handle = buffers.add(ShaderBuffer::with_size(
-        size as usize,
+    let handle = buffers.add(ShaderBuffer::new(
+        &vec![0u8; size as usize],
         RenderAssetUsages::all(),
     ));
     commands
         .spawn(Buffer {
             handle,
             readback_buffer: readback_buffer(&render_device, size),
             size,
+            synced: true,
         })
         .id()
 }
@@ -73,30 +78,37 @@ pub fn create_buffer_with_data(
             handle,
             readback_buffer: readback_buffer(&render_device, size),
             size,
+            synced: true,
         })
         .id()
 }
 
-pub fn write_buffer_gpu(
+/// Mutate the CPU-side data of a `ShaderBuffer` in place. Fires
+/// `AssetEvent::Modified` so Bevy's render-asset extract uploads the new
+/// contents to the GPU at the next sync point.
+pub fn write_buffer_cpu(
     In((handle, offset, data)): In<(Handle<ShaderBuffer>, u64, Vec<u8>)>,
-    gpu_buffers: Res<RenderAssets<GpuShaderBuffer>>,
-    render_queue: Res<RenderQueue>,
+    mut buffers: ResMut<Assets<ShaderBuffer>>,
 ) -> Result<()> {
-    let gpu_buffer = &gpu_buffers
-        .get(&handle)
-        .ok_or(ProcessingError::BufferNotFound)?
-        .buffer;
-    render_queue.write_buffer(gpu_buffer, offset, &data);
+    let mut asset = buffers
+        .get_mut(&handle)
+        .ok_or(ProcessingError::BufferNotFound)?;
+    let dst = asset
+        .data
+        .as_mut()
+        .ok_or(ProcessingError::BufferNotFound)?;
+    let start = offset as usize;
+    let end = start + data.len();
+    dst[start..end].copy_from_slice(&data);
     Ok(())
 }
 
+/// Copy the GPU buffer back to CPU and return its full contents. Runs in the
+/// render world; the caller is responsible for writing the bytes back into
+/// `ShaderBuffer.data` via `Assets::get_mut_untracked` (avoiding spurious
+/// `AssetEvent::Modified`s, since this is a readback, not a stage-for-upload).
 pub fn read_buffer_gpu(
-    In((handle, readback_buffer, src_offset, len)): In<(
-        Handle<ShaderBuffer>,
-        WgpuBuffer,
-        u64,
-        u64,
-    )>,
+    In((handle, readback_buffer, size)): In<(Handle<ShaderBuffer>, WgpuBuffer, u64)>,
     gpu_buffers: Res<RenderAssets<GpuShaderBuffer>>,
     render_device: Res<RenderDevice>,
     render_queue: Res<RenderQueue>,
@@ -107,10 +119,10 @@ pub fn read_buffer_gpu(
         .buffer;
 
     let mut encoder = render_device.create_command_encoder(&CommandEncoderDescriptor::default());
-    encoder.copy_buffer_to_buffer(gpu_buffer, src_offset, &readback_buffer, 0, len);
+    encoder.copy_buffer_to_buffer(gpu_buffer, 0, &readback_buffer, 0, size);
     render_queue.submit(std::iter::once(encoder.finish()));
 
-    let buffer_slice = readback_buffer.slice(0..len);
+    let buffer_slice = readback_buffer.slice(0..size);
     let (s, r) = crossbeam_channel::bounded(1);
     buffer_slice.map_async(MapMode::Read, move |result| {
         let _ = s.send(result);
@@ -122,10 +134,9 @@ pub fn read_buffer_gpu(
         .map_err(|e| ProcessingError::BufferMapError(format!("map channel closed: {e}")))?
         .map_err(|e| ProcessingError::BufferMapError(format!("map failed: {e}")))?;
 
-    let data = buffer_slice.get_mapped_range().to_vec();
+    let bytes = buffer_slice.get_mapped_range().to_vec();
     readback_buffer.unmap();
-
-    Ok(data)
+    Ok(bytes)
 }
 
 pub fn destroy_buffer(In(entity): In<Entity>, mut commands: Commands) -> Result<()> {
@@ -139,6 +150,11 @@ pub struct Compute {
     pub entry_point: String,
     pub pipeline_id: CachedComputePipelineId,
     pub bind_group_layout_descriptors: Vec<(u32, BindGroupLayoutDescriptor)>,
+    /// Buffer entities bound to this compute on a `read_write` storage param.
+    /// Their CPU view of GPU data is invalidated after each dispatch so the
+    /// next read/write does a readback. Read-only bindings don't need this
+    /// since the dispatch can't mutate them.
+    pub rw_buffers: HashMap<String, Entity>,
 }
 
 fn queue_pipeline(
@@ -240,6 +256,7 @@ pub fn create_compute(app: &mut App, shader_entity: Entity) -> Result<Entity> {
                     entry_point,
                     pipeline_id,
                     bind_group_layout_descriptors,
+                    rw_buffers: HashMap::new(),
                 })
                 .id());
         }
@@ -267,11 +284,16 @@ pub fn set_compute_property(
         .ok_or_else(|| ProcessingError::UnknownShaderProperty(name.clone()))?;
 
     match (&value, category) {
-        (ShaderValue::Buffer(buf_entity), ParameterCategory::Storage { .. }) => {
+        (ShaderValue::Buffer(buf_entity), ParameterCategory::Storage { read_only }) => {
             let buffer = p_buffers
                 .get(*buf_entity)
                 .map_err(|_| ProcessingError::BufferNotFound)?;
             compute.shader.insert(&name, buffer.handle.clone());
+            if read_only {
+                compute.rw_buffers.remove(&name);
+            } else {
+                compute.rw_buffers.insert(name.clone(), *buf_entity);
+            }
             Ok(())
         }
         (ShaderValue::Texture(img_entity), ParameterCategory::Texture)
diff --git a/crates/processing_render/src/graphics.rs b/crates/processing_render/src/graphics.rs
@@ -453,6 +453,36 @@ pub fn flush(app: &mut App, entity: Entity) -> Result<()> {
     Ok(())
 }
 
+/// Flush all graphics with pending commands and run a frame so any other
+/// pending GPU state (asset writes, etc.) is extracted and uploaded. Used as
+/// a sync boundary before operations like compute dispatch that may bind
+/// graphics targets or recently-mutated assets.
+pub fn flush_all(app: &mut App) {
+    let mut to_flush = Vec::new();
+    let world = app.world_mut();
+    let mut q = world.query::<(Entity, &CommandBuffer, &Graphics)>();
+    for (e, cb, _) in q.iter(world) {
+        if !cb.commands.is_empty() {
+            to_flush.push(e);
+        }
+    }
+
+    for e in &to_flush {
+        if let Ok(mut em) = world.get_entity_mut(*e) {
+            em.insert(Flush);
+        }
+    }
+
+    app.update();
+
+    let world = app.world_mut();
+    for e in &to_flush {
+        if let Ok(mut em) = world.get_entity_mut(*e) {
+            em.remove::<Flush>();
+        }
+    }
+}
+
 pub fn present(app: &mut App, entity: Entity) -> Result<()> {
     graphics_mut!(app, entity)
         .get_mut::<Camera>()
diff --git a/crates/processing_render/src/lib.rs b/crates/processing_render/src/lib.rs
@@ -1676,6 +1676,49 @@ pub fn buffer_write_element(entity: Entity, offset: u64, data: Vec<u8>) -> error
     buffer_write_range(entity, offset, data, false)
 }
 
+/// Ensure `ShaderBuffer.data` reflects current GPU contents by reading the
+/// buffer back into the asset if it was invalidated by a prior dispatch.
+/// Subsequent reads/writes can then operate on the in-asset bytes directly.
+fn ensure_buffer_synced(app: &mut App, entity: Entity) -> error::Result<()> {
+    let (handle, readback_buffer, size, synced) = {
+        let buf = app
+            .world()
+            .get::<compute::Buffer>(entity)
+            .ok_or(error::ProcessingError::BufferNotFound)?;
+        (
+            buf.handle.clone(),
+            buf.readback_buffer.clone(),
+            buf.size,
+            buf.synced,
+        )
+    };
+    if synced {
+        return Ok(());
+    }
+    let bytes = app
+        .sub_app_mut(bevy::render::RenderApp)
+        .world_mut()
+        .run_system_cached_with(
+            compute::read_buffer_gpu,
+            (handle.clone(), readback_buffer, size),
+        )
+        .unwrap()?;
+
+    let world = app.world_mut();
+    let mut buffers = world.resource_mut::<Assets<bevy::render::storage::ShaderBuffer>>();
+    let asset = buffers
+        .get_mut_untracked(handle.id())
+        .ok_or(error::ProcessingError::BufferNotFound)?;
+    asset.data = Some(bytes);
+    drop(buffers);
+
+    let mut buf = world
+        .get_mut::<compute::Buffer>(entity)
+        .ok_or(error::ProcessingError::BufferNotFound)?;
+    buf.synced = true;
+    Ok(())
+}
+
 fn buffer_write_range(
     entity: Entity,
     offset: u64,
@@ -1706,9 +1749,9 @@ fn buffer_write_range(
                 data.len()
             )));
         }
-        app.sub_app_mut(bevy::render::RenderApp)
-            .world_mut()
-            .run_system_cached_with(compute::write_buffer_gpu, (handle, offset, data))
+        ensure_buffer_synced(app, entity)?;
+        app.world_mut()
+            .run_system_cached_with(compute::write_buffer_cpu, (handle, offset, data))
             .unwrap()
     })
 }
@@ -1724,13 +1767,11 @@ pub fn buffer_read(entity: Entity) -> error::Result<Vec<u8>> {
 
 fn buffer_read_range(entity: Entity, offset: u64, len: u64) -> error::Result<Vec<u8>> {
     app_mut(|app| {
-        let (handle, readback_buffer, size) = {
-            let buf = app
-                .world()
-                .get::<compute::Buffer>(entity)
-                .ok_or(error::ProcessingError::BufferNotFound)?;
-            (buf.handle.clone(), buf.readback_buffer.clone(), buf.size)
-        };
+        let size = app
+            .world()
+            .get::<compute::Buffer>(entity)
+            .ok_or(error::ProcessingError::BufferNotFound)?
+            .size;
         let end = offset.checked_add(len).ok_or_else(|| {
             error::ProcessingError::InvalidArgument("offset + len overflow".to_string())
         })?;
@@ -1739,13 +1780,21 @@ fn buffer_read_range(entity: Entity, offset: u64, len: u64) -> error::Result<Vec
                 "buffer read out of bounds: offset {offset} + len {len} > size {size}"
             )));
         }
-        app.sub_app_mut(bevy::render::RenderApp)
-            .world_mut()
-            .run_system_cached_with(
-                compute::read_buffer_gpu,
-                (handle, readback_buffer, offset, len),
-            )
-            .unwrap()
+        ensure_buffer_synced(app, entity)?;
+        let handle = app
+            .world()
+            .get::<compute::Buffer>(entity)
+            .ok_or(error::ProcessingError::BufferNotFound)?
+            .handle
+            .clone();
+        let buffers = app
+            .world()
+            .resource::<Assets<bevy::render::storage::ShaderBuffer>>();
+        let data = buffers
+            .get(&handle)
+            .and_then(|a| a.data.as_ref())
+            .ok_or(error::ProcessingError::BufferNotFound)?;
+        Ok(data[offset as usize..(offset + len) as usize].to_vec())
     })
 }
 
@@ -1775,22 +1824,41 @@ pub fn compute_set(
 
 pub fn compute_dispatch(entity: Entity, x: u32, y: u32, z: u32) -> error::Result<()> {
     app_mut(|app| {
-        let c = app
-            .world()
-            .get::<compute::Compute>(entity)
-            .ok_or(error::ProcessingError::ComputeNotFound)?;
-        let args = (
-            c.pipeline_id,
-            c.bind_group_layout_descriptors.clone(),
-            c.shader.clone(),
-            x,
-            y,
-            z,
-        );
+        // Flush any pending graphics work and let Bevy's render-asset extract
+        // upload any CPU-side buffer mutations to the GPU before the dispatch
+        // runs. This is the sync boundary for compute inputs.
+        crate::graphics::flush_all(app);
+
+        let (args, rw_entities) = {
+            let c = app
+                .world()
+                .get::<compute::Compute>(entity)
+                .ok_or(error::ProcessingError::ComputeNotFound)?;
+            let args = (
+                c.pipeline_id,
+                c.bind_group_layout_descriptors.clone(),
+                c.shader.clone(),
+                x,
+                y,
+                z,
+            );
+            let rw_entities: Vec<Entity> = c.rw_buffers.values().copied().collect();
+            (args, rw_entities)
+        };
         app.sub_app_mut(bevy::render::RenderApp)
             .world_mut()
             .run_system_cached_with(compute::dispatch, args)
-            .unwrap()
+            .unwrap()?;
+
+        // Invalidate the CPU view of any buffer the dispatch could have
+        // written. The next read or write on those buffers will readback first.
+        let world = app.world_mut();
+        for e in rw_entities {
+            if let Some(mut buf) = world.get_mut::<compute::Buffer>(e) {
+                buf.synced = false;
+            }
+        }
+        Ok(())
     })
 }
 
