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//!
10//! Also owns the release half of that stash: [`WgpuGfxCtx::present`] /
11//! [`WgpuGfxCtx::release_frame_texture`] drop the ctx's frame handles before
12//! the swap chain gets the frame back (required on DX12 — see `present`).
13
14use crate::WgpuGfxCtx;
15
16impl WgpuGfxCtx {
17 /// Stash a handle to the current frame's surface texture so a later
18 /// [`Self::read_screenshot`] / `capture_screenshot` call can copy from
19 /// it. Called from the platform shell with `frame.texture.clone()`
20 /// BEFORE [`begin_frame`](WgpuGfxCtx::begin_frame). `wgpu::Texture` is
21 /// internally ref-counted, so the clone is cheap — but it is a live
22 /// reference to the swap-chain back buffer, so the frame MUST then be
23 /// presented through [`Self::present`] (or the stash dropped with
24 /// [`Self::release_frame_texture`]), never via a bare
25 /// `SurfaceTexture::present()`.
26 ///
27 /// If the previous frame's texture is still stashed (the shell never
28 /// released it), this logs a one-time `log::warn!` per ctx, since that
29 /// is the pattern that breaks DX12 swap-chain resize.
30 pub fn set_surface_texture(&mut self, tex: wgpu::Texture) {
31 if self.surface_texture.is_some() && !self.warned_unreleased_frame {
32 self.warned_unreleased_frame = true;
33 log::warn!(
34 "agg-gui-wgpu: set_surface_texture called while the previous frame's \
35 texture is still stashed. Present frames via WgpuGfxCtx::present (or call \
36 release_frame_texture) — holding a back-buffer handle past present makes \
37 DX12 swap-chain resizes fail (\"Invalid surface\")."
38 );
39 }
40 self.surface_texture = Some(tex);
41 }
42
43 /// Release the ctx's per-frame handles: the [`Self::set_surface_texture`]
44 /// stash and the [`begin_frame`](WgpuGfxCtx::begin_frame) view if
45 /// `end_frame` did not already consume it.
46 ///
47 /// [`Self::present`] calls this for you — including when the stash is an
48 /// SSAA / scene texture rather than the surface itself. Call it directly
49 /// only from shells that do not present a frame at all (e.g. a frame
50 /// abandoned mid-way or an offscreen-only render). Screenshot read-back
51 /// must happen before this — afterwards [`Self::read_screenshot`]
52 /// returns an empty buffer.
53 pub fn release_frame_texture(&mut self) {
54 self.surface_texture = None;
55 self.surface_view = None;
56 }
57
58 /// Present `frame`, first releasing the ctx's per-frame handles
59 /// (see [`Self::release_frame_texture`]).
60 ///
61 /// This is THE way to present whenever the ctx was given the frame's
62 /// surface texture (via [`Self::set_surface_texture`] or
63 /// [`begin_frame`](WgpuGfxCtx::begin_frame)). On Windows/DX12 a clone of
64 /// the back buffer that outlives `present()` keeps the swap chain
65 /// referenced, so the next `Surface::configure` (DXGI `ResizeBuffers`)
66 /// fails with "Invalid surface" and recreating the swap chain fails with
67 /// `E_ACCESSDENIED`. Do any [`Self::read_screenshot`] before calling.
68 pub fn present(&mut self, frame: wgpu::SurfaceTexture) {
69 self.release_frame_texture();
70 frame.present();
71 }
72
73 /// Stash captured screenshot pixels for the read-back closure to pick
74 /// up. See [`Self::take_pending_screenshot`].
75 pub fn set_pending_screenshot(&mut self, captured: (Vec<u8>, u32, u32)) {
76 self.pending_screenshot = Some(captured);
77 }
78
79 /// Consume the pending screenshot pixels — returns `(Vec::new(), 0, 0)`
80 /// when none are stashed (typical non-capture frames). Called by the
81 /// `agg_gui::screenshot::run_frame_with_capture` read-back closure.
82 pub fn take_pending_screenshot(&mut self) -> (Vec<u8>, u32, u32) {
83 self.pending_screenshot.take().unwrap_or((Vec::new(), 0, 0))
84 }
85
86 /// Read the current frame's rendered pixels back to CPU memory as a
87 /// top-down RGBA8 buffer. Returns `(pixels, width, height)`.
88 /// The first `width * 4` bytes are the TOP row (Y-down image order).
89 ///
90 /// Must be called AFTER [`Self::end_frame`] has submitted the render and
91 /// BEFORE the platform shell calls [`Self::present`] (which releases the
92 /// stashed texture). Requires the platform shell to have called
93 /// [`Self::set_surface_texture`] earlier in the frame so we hold a handle
94 /// into the surface that's still valid post-render.
95 ///
96 /// Returns an empty buffer if no surface texture is currently stashed
97 /// (including after [`Self::present`] / [`Self::release_frame_texture`]).
98 pub fn read_screenshot(&self) -> (Vec<u8>, u32, u32) {
99 let Some(texture) = self.surface_texture.as_ref() else {
100 return (Vec::new(), 0, 0);
101 };
102 let size = texture.size();
103 let w = size.width;
104 let h = size.height;
105 if w == 0 || h == 0 {
106 return (Vec::new(), 0, 0);
107 }
108
109 // wgpu requires `bytes_per_row` to be a multiple of
110 // COPY_BYTES_PER_ROW_ALIGNMENT (256). We allocate a padded buffer
111 // for the copy and strip the padding row-by-row when assembling the
112 // returned `Vec<u8>`.
113 const ALIGN: u32 = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
114 let unpadded_bpr = w * 4;
115 let padded_bpr = unpadded_bpr.div_ceil(ALIGN) * ALIGN;
116 let buffer_size = (padded_bpr as u64) * (h as u64);
117
118 let staging = self.device.create_buffer(&wgpu::BufferDescriptor {
119 label: Some("screenshot_staging"),
120 size: buffer_size,
121 usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
122 mapped_at_creation: false,
123 });
124
125 let mut encoder = self
126 .device
127 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
128 label: Some("screenshot_copy"),
129 });
130 encoder.copy_texture_to_buffer(
131 wgpu::TexelCopyTextureInfo {
132 texture,
133 mip_level: 0,
134 origin: wgpu::Origin3d::ZERO,
135 aspect: wgpu::TextureAspect::All,
136 },
137 wgpu::TexelCopyBufferInfo {
138 buffer: &staging,
139 layout: wgpu::TexelCopyBufferLayout {
140 offset: 0,
141 bytes_per_row: Some(padded_bpr),
142 rows_per_image: Some(h),
143 },
144 },
145 wgpu::Extent3d {
146 width: w,
147 height: h,
148 depth_or_array_layers: 1,
149 },
150 );
151 self.queue.submit(std::iter::once(encoder.finish()));
152
153 // Map the staging buffer. `map_async` is async via callback; we
154 // poll the device until the map completes. On native this is fine
155 // (synchronous from the caller's POV); on WASM the wgpu webgl
156 // backend resolves the future on the JS event-loop tick that the
157 // surrounding render loop is running on, so this still works
158 // because the JS harness drives `render()` from a microtask.
159 let slice = staging.slice(..);
160 let (sender, receiver) = std::sync::mpsc::channel();
161 slice.map_async(wgpu::MapMode::Read, move |res| {
162 let _ = sender.send(res);
163 });
164 let _ = self.device.poll(wgpu::PollType::wait_indefinitely());
165 let map_result = receiver
166 .recv()
167 .expect("map_async sender dropped before resolving");
168 if map_result.is_err() {
169 return (Vec::new(), 0, 0);
170 }
171
172 // Surface format may be Bgra8Unorm; PNG / JS expects RGBA so swap
173 // R↔B per pixel as we copy. Surface textures are Y-down, which
174 // matches the screenshot module's "TOP row first" convention, so
175 // no row flip needed.
176 let bgra = matches!(
177 self.surface_format,
178 wgpu::TextureFormat::Bgra8Unorm | wgpu::TextureFormat::Bgra8UnormSrgb
179 );
180 let mut out = Vec::with_capacity((w as usize) * (h as usize) * 4);
181 {
182 let view = slice.get_mapped_range();
183 for row in 0..h as usize {
184 let start = row * padded_bpr as usize;
185 let end = start + unpadded_bpr as usize;
186 let src = &view[start..end];
187 if bgra {
188 for px in src.chunks_exact(4) {
189 out.extend_from_slice(&[px[2], px[1], px[0], px[3]]);
190 }
191 } else {
192 out.extend_from_slice(src);
193 }
194 }
195 }
196 staging.unmap();
197 (out, w, h)
198 }
199}