agg_gui_wgpu/screenshot_readback.rs
1//! Frame screenshot read-back for [`WgpuGfxCtx`].
2//!
3//! Split out of `lib.rs` (800-line guardrail). Owns the surface-texture
4//! stash and the GPU->CPU copy path used by
5//! `agg_gui::screenshot::run_frame_with_capture`: the platform shell
6//! stashes the frame texture before `begin_frame`, and after
7//! `end_frame` the capture closure copies it into a top-down RGBA8
8//! buffer via a padded staging buffer.
9
10use crate::WgpuGfxCtx;
11
12impl WgpuGfxCtx {
13 /// Stash a handle to the current frame's surface texture so a later
14 /// [`Self::read_screenshot`] call can copy from it. Called from the
15 /// platform shell with `frame.texture.clone()` BEFORE [`begin_frame`](WgpuGfxCtx::begin_frame).
16 /// `wgpu::Texture` is internally ref-counted, so the clone is cheap.
17 pub fn set_surface_texture(&mut self, tex: wgpu::Texture) {
18 self.surface_texture = Some(tex);
19 }
20
21 /// Stash captured screenshot pixels for the read-back closure to pick
22 /// up. See [`Self::take_pending_screenshot`].
23 pub fn set_pending_screenshot(&mut self, captured: (Vec<u8>, u32, u32)) {
24 self.pending_screenshot = Some(captured);
25 }
26
27 /// Consume the pending screenshot pixels — returns `(Vec::new(), 0, 0)`
28 /// when none are stashed (typical non-capture frames). Called by the
29 /// `agg_gui::screenshot::run_frame_with_capture` read-back closure.
30 pub fn take_pending_screenshot(&mut self) -> (Vec<u8>, u32, u32) {
31 self.pending_screenshot.take().unwrap_or((Vec::new(), 0, 0))
32 }
33
34 /// Read the current frame's rendered pixels back to CPU memory as a
35 /// top-down RGBA8 buffer. Returns `(pixels, width, height)`.
36 /// The first `width * 4` bytes are the TOP row (Y-down image order).
37 ///
38 /// Must be called AFTER [`Self::end_frame`] has submitted the render and
39 /// BEFORE the platform shell calls `frame.present()`. Requires the
40 /// platform shell to have called [`Self::set_surface_texture`] earlier
41 /// in the frame so we hold a handle into the surface that's still
42 /// valid post-render.
43 ///
44 /// Returns an empty buffer if no surface texture is currently stashed.
45 pub fn read_screenshot(&self) -> (Vec<u8>, u32, u32) {
46 let Some(texture) = self.surface_texture.as_ref() else {
47 return (Vec::new(), 0, 0);
48 };
49 let size = texture.size();
50 let w = size.width;
51 let h = size.height;
52 if w == 0 || h == 0 {
53 return (Vec::new(), 0, 0);
54 }
55
56 // wgpu requires `bytes_per_row` to be a multiple of
57 // COPY_BYTES_PER_ROW_ALIGNMENT (256). We allocate a padded buffer
58 // for the copy and strip the padding row-by-row when assembling the
59 // returned `Vec<u8>`.
60 const ALIGN: u32 = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
61 let unpadded_bpr = w * 4;
62 let padded_bpr = unpadded_bpr.div_ceil(ALIGN) * ALIGN;
63 let buffer_size = (padded_bpr as u64) * (h as u64);
64
65 let staging = self.device.create_buffer(&wgpu::BufferDescriptor {
66 label: Some("screenshot_staging"),
67 size: buffer_size,
68 usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
69 mapped_at_creation: false,
70 });
71
72 let mut encoder = self
73 .device
74 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
75 label: Some("screenshot_copy"),
76 });
77 encoder.copy_texture_to_buffer(
78 wgpu::TexelCopyTextureInfo {
79 texture,
80 mip_level: 0,
81 origin: wgpu::Origin3d::ZERO,
82 aspect: wgpu::TextureAspect::All,
83 },
84 wgpu::TexelCopyBufferInfo {
85 buffer: &staging,
86 layout: wgpu::TexelCopyBufferLayout {
87 offset: 0,
88 bytes_per_row: Some(padded_bpr),
89 rows_per_image: Some(h),
90 },
91 },
92 wgpu::Extent3d {
93 width: w,
94 height: h,
95 depth_or_array_layers: 1,
96 },
97 );
98 self.queue.submit(std::iter::once(encoder.finish()));
99
100 // Map the staging buffer. `map_async` is async via callback; we
101 // poll the device until the map completes. On native this is fine
102 // (synchronous from the caller's POV); on WASM the wgpu webgl
103 // backend resolves the future on the JS event-loop tick that the
104 // surrounding render loop is running on, so this still works
105 // because the JS harness drives `render()` from a microtask.
106 let slice = staging.slice(..);
107 let (sender, receiver) = std::sync::mpsc::channel();
108 slice.map_async(wgpu::MapMode::Read, move |res| {
109 let _ = sender.send(res);
110 });
111 let _ = self.device.poll(wgpu::PollType::wait_indefinitely());
112 let map_result = receiver
113 .recv()
114 .expect("map_async sender dropped before resolving");
115 if map_result.is_err() {
116 return (Vec::new(), 0, 0);
117 }
118
119 // Surface format may be Bgra8Unorm; PNG / JS expects RGBA so swap
120 // R↔B per pixel as we copy. Surface textures are Y-down, which
121 // matches the screenshot module's "TOP row first" convention, so
122 // no row flip needed.
123 let bgra = matches!(
124 self.surface_format,
125 wgpu::TextureFormat::Bgra8Unorm | wgpu::TextureFormat::Bgra8UnormSrgb
126 );
127 let mut out = Vec::with_capacity((w as usize) * (h as usize) * 4);
128 {
129 let view = slice.get_mapped_range();
130 for row in 0..h as usize {
131 let start = row * padded_bpr as usize;
132 let end = start + unpadded_bpr as usize;
133 let src = &view[start..end];
134 if bgra {
135 for px in src.chunks_exact(4) {
136 out.extend_from_slice(&[px[2], px[1], px[0], px[3]]);
137 }
138 } else {
139 out.extend_from_slice(src);
140 }
141 }
142 }
143 staging.unmap();
144 (out, w, h)
145 }
146}