bevy_shader_plugin/
readback.rs

1use std::{
2    any::type_name,
3    marker::PhantomData,
4    sync::{
5        Arc, Mutex,
6        atomic::{AtomicBool, Ordering},
7    },
8};
9
10use bevy::{
11    core_pipeline::schedule::camera_driver,
12    prelude::*,
13    render::{
14        Render, RenderApp, RenderSystems,
15        render_asset::RenderAssets,
16        render_resource::{Buffer, BufferDescriptor, BufferUsages, MapMode},
17        renderer::{RenderContext, RenderDevice, RenderGraph, RenderGraphSystems},
18        storage::GpuShaderBuffer,
19    },
20};
21use encase::{
22    ShaderSize, ShaderType,
23    internal::{CreateFrom, WriteInto},
24};
25
26use crate::{
27    ArrayBufferChanges, ShaderPassSystems, ShaderSystems,
28    auto_array::{ArrayBufferState, array_stride},
29    bindings::AutoBufferKind,
30    cpu_buffer::CpuBuffer,
31    gpu_buffer::GpuBuffer,
32};
33
34/// When a read-back buffer is copied to the CPU.
35#[derive(Clone, Copy, Debug, PartialEq, Eq)]
36pub enum ReadbackMode {
37    /// Copy the buffer every frame (skipping frames while a previous copy is still
38    /// being mapped), so `latest()` keeps updating. Costs one buffer copy per frame.
39    EveryFrame,
40    /// Copy the buffer once per `request()` call.
41    OnRequest,
42}
43
44/// State shared between a main-world readback resource and the render world. Held in an
45/// `Arc` by both, so no extraction is needed and it works with pipelined rendering.
46#[derive(Default)]
47struct ReadbackShared {
48    /// Set by `request()`, consumed when a copy is queued.
49    requested: AtomicBool,
50    /// A copy has been queued and its bytes haven't arrived yet.
51    in_flight: AtomicBool,
52    /// Bytes delivered by the map callback, taken by the main world.
53    result: Mutex<Option<Vec<u8>>>,
54}
55
56impl ReadbackShared {
57    fn take_result(&self) -> Option<Vec<u8>> {
58        self.result.lock().ok()?.take()
59    }
60}
61
62/// The most recent CPU copy of an array buffer or GPU buffer, decoded into `T` values.
63///
64/// This resource only changes when new data arrives, so check `is_changed()` to react to
65/// fresh data. [`request`](Self::request) takes `&self`, so requesting through
66/// `Res<GpuReadback<..>>` doesn't count as a change:
67///
68/// ```rust,ignore
69/// fn request(readback: Res<GpuReadback<Particles, Particle>>) {
70///     readback.request();
71/// }
72///
73/// fn print(readback: Res<GpuReadback<Particles, Particle>>) {
74///     if !readback.is_changed() { return; }
75///     let Some(particles) = readback.latest() else { return };
76///     info!("first particle at {}", particles[0].pos);
77/// }
78/// ```
79///
80/// The data shows the buffer at the end of the frame it was copied in, and arrives
81/// 1–3 frames later. Use it for UI, stats, or logging — not for logic that needs the
82/// exact current GPU state.
83#[derive(Resource)]
84pub struct GpuReadback<Tag, T> {
85    shared: Arc<ReadbackShared>,
86    latest: Option<Vec<T>>,
87    _marker: PhantomData<Tag>,
88}
89
90impl<Tag, T> GpuReadback<Tag, T> {
91    /// The most recently read-back contents of the buffer, or `None` until the first
92    /// readback arrives.
93    pub fn latest(&self) -> Option<&[T]> {
94        self.latest.as_deref()
95    }
96
97    /// Ask for the buffer to be copied back. Only needed with
98    /// [`ReadbackMode::OnRequest`]. Requests made while a copy is in flight are merged
99    /// into one follow-up copy.
100    pub fn request(&self) {
101        self.shared.requested.store(true, Ordering::Release);
102    }
103}
104
105/// The most recent CPU copy of a storage buffer, decoded as the registered resource type
106/// `S` (which may itself contain arrays).
107///
108/// Works like [`GpuReadback`]: check `is_changed()` for fresh data, and call
109/// [`request`](Self::request) in [`ReadbackMode::OnRequest`] mode. The data arrives
110/// 1–3 frames after the copy.
111#[derive(Resource)]
112pub struct StorageReadback<S> {
113    shared: Arc<ReadbackShared>,
114    latest: Option<S>,
115}
116
117impl<S> StorageReadback<S> {
118    /// The most recently read-back contents of the buffer, or `None` until the first
119    /// readback arrives.
120    pub fn latest(&self) -> Option<&S> {
121        self.latest.as_ref()
122    }
123
124    /// Ask for the buffer to be copied back. Only needed with
125    /// [`ReadbackMode::OnRequest`]. Requests made while a copy is in flight are merged
126    /// into one follow-up copy.
127    pub fn request(&self) {
128        self.shared.requested.store(true, Ordering::Release);
129    }
130}
131
132// ---------------------------------------------------------------------------
133// Render world
134// ---------------------------------------------------------------------------
135
136/// One buffer copy recorded by [`readback_copy_pass`] this frame.
137struct QueuedCopy {
138    source: Buffer,
139    staging: Buffer,
140    size: u64,
141    shared: Arc<ReadbackShared>,
142}
143
144/// Copies queued this frame. Filled during `PrepareBindGroups`, recorded by
145/// [`readback_copy_pass`], then mapped and cleared by [`map_readbacks`].
146#[derive(Resource, Default)]
147struct ReadbackCopies(Vec<QueuedCopy>);
148
149/// Render-world state for one read-back registration. `Key` distinguishes registrations.
150#[derive(Resource)]
151struct ReadbackJob<Key> {
152    shared: Arc<ReadbackShared>,
153    mode: ReadbackMode,
154    staging: Option<Buffer>,
155    _marker: PhantomData<Key>,
156}
157
158// Keys keep registrations of different buffer kinds apart.
159struct GpuKey<Tag>(PhantomData<Tag>);
160struct ArrayKey<Tag>(PhantomData<Tag>);
161struct StorageKey<S>(PhantomData<S>);
162
163/// Whether to queue a copy this frame. Never copies while a previous copy is in flight;
164/// an on-request copy consumes the request, which otherwise stays pending.
165fn should_copy(mode: ReadbackMode, in_flight: bool, requested: &AtomicBool) -> bool {
166    if in_flight {
167        return false;
168    }
169    match mode {
170        ReadbackMode::EveryFrame => true,
171        ReadbackMode::OnRequest => requested.swap(false, Ordering::AcqRel),
172    }
173}
174
175impl<Key> ReadbackJob<Key> {
176    fn new(shared: Arc<ReadbackShared>, mode: ReadbackMode) -> Self {
177        Self {
178            shared,
179            mode,
180            staging: None,
181            _marker: PhantomData,
182        }
183    }
184
185    /// Queue a copy of the first `size` bytes of `source`, if this job wants one now.
186    fn queue_copy(
187        &mut self,
188        source: &Buffer,
189        size: u64,
190        render_device: &RenderDevice,
191        copies: &mut ReadbackCopies,
192    ) {
193        if size == 0
194            || !should_copy(
195                self.mode,
196                self.shared.in_flight.load(Ordering::Acquire),
197                &self.shared.requested,
198            )
199        {
200            return;
201        }
202
203        // Reuse the staging buffer unless the size changed. It is never still mapped
204        // here: `in_flight` is only cleared after unmapping.
205        let staging = match &self.staging {
206            Some(staging) if staging.size() == size => staging.clone(),
207            _ => {
208                let staging = render_device.create_buffer(&BufferDescriptor {
209                    label: Some("readback_staging_buffer"),
210                    size,
211                    usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
212                    mapped_at_creation: false,
213                });
214                self.staging = Some(staging.clone());
215                staging
216            }
217        };
218
219        self.shared.in_flight.store(true, Ordering::Release);
220        copies.0.push(QueuedCopy {
221            source: source.clone(),
222            staging,
223            size,
224            shared: self.shared.clone(),
225        });
226    }
227}
228
229/// Records every queued readback copy. Runs in the root
230/// [`RenderGraph`] schedule after
231/// [`camera_driver`](bevy::core_pipeline::schedule::camera_driver), so copies see
232/// everything the frame's compute passes and cameras wrote.
233fn readback_copy_pass(copies: Res<ReadbackCopies>, mut ctx: RenderContext) {
234    if copies.0.is_empty() {
235        return;
236    }
237    let encoder = ctx.command_encoder();
238    for copy in &copies.0 {
239        encoder.copy_buffer_to_buffer(&copy.source, 0, &copy.staging, 0, copy.size);
240    }
241}
242
243/// After the frame is submitted, asks wgpu to map every staging buffer copied into this
244/// frame. The callback hands the bytes to the main world.
245fn map_readbacks(mut copies: ResMut<ReadbackCopies>) {
246    for copy in copies.0.drain(..) {
247        let QueuedCopy {
248            staging,
249            size,
250            shared,
251            ..
252        } = copy;
253        let mapped = staging.clone();
254        staging
255            .slice(..size)
256            .map_async(MapMode::Read, move |result| {
257                match result {
258                    Ok(()) => {
259                        let bytes = mapped.slice(..size).get_mapped_range().to_vec();
260                        mapped.unmap();
261                        if let Ok(mut slot) = shared.result.lock() {
262                            *slot = Some(bytes);
263                        }
264                    }
265                    Err(error) => warn!("buffer readback failed: {error}"),
266                }
267                shared.in_flight.store(false, Ordering::Release);
268            });
269    }
270}
271
272/// Copy-and-map machinery shared by every readback registration.
273struct ReadbackPlugin;
274
275impl Plugin for ReadbackPlugin {
276    fn build(&self, app: &mut App) {
277        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
278            return;
279        };
280        render_app
281            .init_resource::<ReadbackCopies>()
282            .add_systems(
283                RenderGraph,
284                readback_copy_pass
285                    .in_set(ShaderPassSystems::ReadbackCopy)
286                    .in_set(RenderGraphSystems::Render)
287                    .after(camera_driver),
288            )
289            // Cleanup runs after the frame's commands are submitted, which is also where
290            // Bevy maps its own readbacks.
291            .add_systems(Render, map_readbacks.in_set(RenderSystems::Cleanup));
292    }
293}
294
295fn add_readback_plugin(app: &mut App) {
296    if !app.is_plugin_added::<ReadbackPlugin>() {
297        app.add_plugins(ReadbackPlugin);
298    }
299}
300
301/// Adds the render-world job for `Key` and a system that resolves its source buffer each
302/// frame and queues copies.
303fn add_render_job<Key, M>(
304    app: &mut App,
305    shared: Arc<ReadbackShared>,
306    mode: ReadbackMode,
307    queue_system: impl IntoScheduleConfigs<bevy::ecs::system::ScheduleSystem, M>,
308) where
309    Key: Send + Sync + 'static,
310{
311    let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
312        return;
313    };
314    render_app
315        .insert_resource(ReadbackJob::<Key>::new(shared, mode))
316        .add_systems(
317            Render,
318            queue_system
319                .in_set(RenderSystems::PrepareBindGroups)
320                .after(ShaderSystems::PrepareBindings),
321        );
322}
323
324fn decode<D: ShaderType + CreateFrom>(bytes: &[u8]) -> Option<D> {
325    match encase::StorageBuffer::new(bytes).create() {
326        Ok(value) => Some(value),
327        Err(error) => {
328            error!("could not decode read-back {}: {error}", type_name::<D>());
329            None
330        }
331    }
332}
333
334/// Main world: moves newly arrived bytes into the [`GpuReadback`] resource.
335fn deliver_slice<Tag, T>(mut readback: ResMut<GpuReadback<Tag, T>>)
336where
337    Tag: Send + Sync + 'static,
338    T: ShaderSize + CreateFrom + Send + Sync + 'static,
339{
340    if let Some(bytes) = readback.shared.take_result()
341        && let Some(values) = decode::<Vec<T>>(&bytes)
342    {
343        readback.latest = Some(values);
344    }
345}
346
347/// Main world: moves newly arrived bytes into the [`StorageReadback`] resource.
348fn deliver_storage<S>(mut readback: ResMut<StorageReadback<S>>)
349where
350    S: ShaderType + CreateFrom + Send + Sync + 'static,
351{
352    if let Some(bytes) = readback.shared.take_result()
353        && let Some(value) = decode::<S>(&bytes)
354    {
355        readback.latest = Some(value);
356    }
357}
358
359/// Inserts `GpuReadback<Tag, T>` and its delivery system, returning the shared state.
360fn insert_slice_readback<Tag, T>(app: &mut App) -> Arc<ReadbackShared>
361where
362    Tag: Send + Sync + 'static,
363    T: ShaderSize + CreateFrom + Send + Sync + 'static,
364{
365    assert!(
366        !app.world().contains_resource::<GpuReadback<Tag, T>>(),
367        "readback is already enabled for tag `{}`",
368        type_name::<Tag>()
369    );
370    let shared = Arc::new(ReadbackShared::default());
371    app.insert_resource(GpuReadback::<Tag, T> {
372        shared: shared.clone(),
373        latest: None,
374        _marker: PhantomData,
375    })
376    .add_systems(PreUpdate, deliver_slice::<Tag, T>);
377    shared
378}
379
380// ---------------------------------------------------------------------------
381// Registration
382// ---------------------------------------------------------------------------
383
384/// See [`ShaderAppExt::read_back_gpu_buffer`](crate::ShaderAppExt::read_back_gpu_buffer).
385pub(crate) fn register_gpu<Tag, T>(app: &mut App, mode: ReadbackMode)
386where
387    Tag: Send + Sync + 'static,
388    T: ShaderSize + CreateFrom + Send + Sync + 'static,
389{
390    let Some(gpu_buffer) = app.world().get_resource::<GpuBuffer<Tag>>() else {
391        panic!(
392            "read_back_gpu_buffer::<{}, _>: call register_gpu_buffer for this tag first",
393            type_name::<Tag>()
394        );
395    };
396    let asset_id = gpu_buffer.handle().id();
397
398    add_readback_plugin(app);
399    let shared = insert_slice_readback::<Tag, T>(app);
400    add_render_job::<GpuKey<Tag>, _>(
401        app,
402        shared,
403        mode,
404        move |buffers: Res<RenderAssets<GpuShaderBuffer>>,
405              mut job: ResMut<ReadbackJob<GpuKey<Tag>>>,
406              render_device: Res<RenderDevice>,
407              mut copies: ResMut<ReadbackCopies>| {
408            if let Some(gpu_buffer) = buffers.get(asset_id) {
409                let size = gpu_buffer.buffer.size();
410                job.queue_copy(&gpu_buffer.buffer, size, &render_device, &mut copies);
411            }
412        },
413    );
414}
415
416/// See [`ShaderAppExt::read_back_array_buffer`](crate::ShaderAppExt::read_back_array_buffer).
417pub(crate) fn register_array<Tag, T>(app: &mut App, mode: ReadbackMode)
418where
419    Tag: Send + Sync + 'static,
420    T: ShaderSize + CreateFrom + Send + Sync + 'static,
421{
422    assert!(
423        app.world().contains_resource::<ArrayBufferChanges<Tag>>(),
424        "read_back_array_buffer::<{}, _>: call register_array_buffer for this tag first",
425        type_name::<Tag>()
426    );
427
428    add_readback_plugin(app);
429    let shared = insert_slice_readback::<Tag, T>(app);
430    add_render_job::<ArrayKey<Tag>, _>(
431        app,
432        shared,
433        mode,
434        |state: Option<Res<ArrayBufferState<Tag>>>,
435         mut job: ResMut<ReadbackJob<ArrayKey<Tag>>>,
436         render_device: Res<RenderDevice>,
437         mut copies: ResMut<ReadbackCopies>| {
438            let Some(state) = state else { return };
439            if state.stride != array_stride::<T>() {
440                error_once!(
441                    "read_back_array_buffer::<{}, {}>: element type doesn't match the one passed to register_array_buffer",
442                    type_name::<Tag>(),
443                    type_name::<T>()
444                );
445                return;
446            }
447            let size = state.buffer.size();
448            job.queue_copy(&state.buffer, size, &render_device, &mut copies);
449        },
450    );
451}
452
453/// See [`ShaderAppExt::read_back_storage_buffer`](crate::ShaderAppExt::read_back_storage_buffer).
454pub(crate) fn register_storage<S>(app: &mut App, mode: ReadbackMode)
455where
456    S: ShaderType + WriteInto + Default + CreateFrom + Resource,
457{
458    let has_storage_slot = app
459        .get_sub_app(RenderApp)
460        .and_then(|render_app| render_app.world().get_resource::<CpuBuffer<S>>())
461        .is_some_and(|state| {
462            state
463                .slot_kinds()
464                .any(|kind| matches!(kind, AutoBufferKind::Storage { .. }))
465        });
466    assert!(
467        has_storage_slot,
468        "read_back_storage_buffer::<{}>: call register_storage_buffer for this type first",
469        type_name::<S>()
470    );
471    assert!(
472        !app.world().contains_resource::<StorageReadback<S>>(),
473        "read_back_storage_buffer::<{}>: readback is already enabled for this type",
474        type_name::<S>()
475    );
476
477    add_readback_plugin(app);
478    let shared = Arc::new(ReadbackShared::default());
479    app.insert_resource(StorageReadback::<S> {
480        shared: shared.clone(),
481        latest: None,
482    })
483    .add_systems(PreUpdate, deliver_storage::<S>);
484
485    add_render_job::<StorageKey<S>, _>(
486        app,
487        shared,
488        mode,
489        |state: Res<CpuBuffer<S>>,
490         mut job: ResMut<ReadbackJob<StorageKey<S>>>,
491         render_device: Res<RenderDevice>,
492         mut copies: ResMut<ReadbackCopies>| {
493            if let Some((buffer, size)) = state.readback_source() {
494                job.queue_copy(buffer, size, &render_device, &mut copies);
495            }
496        },
497    );
498}
499
500#[cfg(test)]
501mod tests {
502    use super::*;
503
504    #[test]
505    fn every_frame_copies_unless_in_flight() {
506        let requested = AtomicBool::new(false);
507        assert!(should_copy(ReadbackMode::EveryFrame, false, &requested));
508        assert!(!should_copy(ReadbackMode::EveryFrame, true, &requested));
509    }
510
511    #[test]
512    fn on_request_consumes_the_request() {
513        let requested = AtomicBool::new(true);
514        assert!(should_copy(ReadbackMode::OnRequest, false, &requested));
515        assert!(!requested.load(Ordering::Acquire));
516        assert!(!should_copy(ReadbackMode::OnRequest, false, &requested));
517    }
518
519    #[test]
520    fn on_request_stays_pending_while_in_flight() {
521        let requested = AtomicBool::new(true);
522        assert!(!should_copy(ReadbackMode::OnRequest, true, &requested));
523        assert!(requested.load(Ordering::Acquire));
524    }
525
526    #[test]
527    fn decodes_what_encase_wrote() {
528        #[derive(ShaderType, Debug, PartialEq)]
529        struct Sample {
530            a: Vec2,
531            b: f32,
532        }
533        let values = vec![
534            Sample {
535                a: Vec2::new(1.0, 2.0),
536                b: 3.0,
537            },
538            Sample {
539                a: Vec2::ZERO,
540                b: -1.0,
541            },
542        ];
543        let mut bytes = encase::StorageBuffer::new(Vec::new());
544        bytes.write(&values).unwrap();
545        assert_eq!(decode::<Vec<Sample>>(&bytes.into_inner()), Some(values));
546    }
547}
bevy_shader_plugin/readback.rs

bevy_shader_plugin/
readback.rs
