Skip to main content

repose_render_wgpu/
offscreen.rs

1//! Unified offscreen helper (native/android/wasm).
2
3use anyhow::Result;
4use wgpu::{PollType, TextureFormat};
5
6use crate::WgpuSceneRenderer;
7use repose_core::Scene;
8#[cfg(not(target_arch = "wasm32"))]
9use std::sync::{Mutex, OnceLock};
10
11#[cfg(not(target_arch = "wasm32"))]
12static SHARED_DEVICE: OnceLock<Mutex<Option<(wgpu::Device, wgpu::Queue)>>> = OnceLock::new();
13
14#[cfg(not(target_arch = "wasm32"))]
15fn shared_slot() -> &'static Mutex<Option<(wgpu::Device, wgpu::Queue)>> {
16    SHARED_DEVICE.get_or_init(|| Mutex::new(None))
17}
18
19#[cfg(target_arch = "wasm32")]
20thread_local! {
21    static SHARED_DEVICE_WASM: std::cell::RefCell<Option<(wgpu::Device, wgpu::Queue)>> = std::cell::RefCell::new(None);
22}
23
24/// Publish live Device/Queue for shared offscreen.
25///
26/// WASM: `thread_local`, so a worker thread never sees the main-thread
27/// device. Do NOT send the `Device`/`Queue` to a worker and render there..
28pub fn set_shared_device(device: wgpu::Device, queue: wgpu::Queue) {
29    #[cfg(not(target_arch = "wasm32"))]
30    {
31        if let Ok(mut g) = shared_slot().lock() {
32            *g = Some((device, queue));
33        }
34    }
35    #[cfg(target_arch = "wasm32")]
36    {
37        SHARED_DEVICE_WASM.with(|c| *c.borrow_mut() = Some((device, queue)));
38    }
39}
40pub fn shared_device() -> Option<(wgpu::Device, wgpu::Queue)> {
41    #[cfg(not(target_arch = "wasm32"))]
42    {
43        shared_slot().lock().ok().and_then(|g| g.clone())
44    }
45    #[cfg(target_arch = "wasm32")]
46    {
47        SHARED_DEVICE_WASM.with(|c| c.borrow().clone())
48    }
49}
50pub fn clear_shared_device() {
51    #[cfg(not(target_arch = "wasm32"))]
52    {
53        if let Ok(mut g) = shared_slot().lock() {
54            *g = None;
55        }
56    }
57    #[cfg(target_arch = "wasm32")]
58    {
59        SHARED_DEVICE_WASM.with(|c| *c.borrow_mut() = None);
60    }
61}
62
63/// Headless scene renderer: renders a [`Scene`] into CPU-readable RGBA8.
64///
65/// # Contracts
66///
67/// - Output is `width * height * 4` bytes, row-major, **premultiplied-alpha
68///   sRGB** (`Rgba8UnormSrgb`): opaque content reads back exact sRGB bytes;
69///   translucent pixels carry `rgb * a` (divide by `a` for straight alpha).
70/// - `clear = None` fills with `scene.clear_color`; `Some(c)` overrides it
71///   with linear-space `[r, g, b, a]` doubles.
72/// - All sizes clamp to a minimum of 1 (no zero-sized textures/buffers).
73/// - `render_rgba` blocks until the GPU work completes and the pixels are
74///   mapped; it is safe to call every frame (no per-frame allocations
75///   beyond the returned `Vec`).
76/// - Without a usable GPU adapter, [`new`](Self::new)/[`new_blocking`](Self::new_blocking)
77///   fail with an error (never a panic, never silent fallback).
78pub struct OffscreenRenderer {
79    renderer: WgpuSceneRenderer,
80    texture: wgpu::Texture,
81    view: wgpu::TextureView,
82    readback: wgpu::Buffer,
83    width: u32,
84    height: u32,
85    padded_bytes_per_row: u32,
86}
87
88impl OffscreenRenderer {
89    /// Create a renderer for `width`x`height` frames with `msaa` samples
90    /// (clamped to what the adapter supports via `pick_surface_msaa`).
91    /// Dimensions clamp to ≥ 1. Errors when no GPU adapter is available.
92    pub async fn new(width: u32, height: u32, msaa: u32) -> Result<Self> {
93        let width = width.max(1);
94        let height = height.max(1);
95        let instance = if cfg!(target_arch = "wasm32") {
96            let mut desc = wgpu::InstanceDescriptor::new_without_display_handle();
97            desc.backends = wgpu::Backends::BROWSER_WEBGPU | wgpu::Backends::GL;
98            wgpu::util::new_instance_with_webgpu_detection(desc).await
99        } else {
100            wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle())
101        };
102        let adapter = instance
103            .request_adapter(&wgpu::RequestAdapterOptions {
104                power_preference: wgpu::PowerPreference::HighPerformance,
105                compatible_surface: None,
106                ..Default::default()
107            })
108            .await?;
109        let format = TextureFormat::Rgba8UnormSrgb;
110        let msaa = crate::pick_surface_msaa(&adapter, format, msaa);
111        let working_space_msaa =
112            crate::pick_surface_msaa_for_mode(&adapter, TextureFormat::Rgba16Float, msaa, true);
113        let (device, queue) = adapter
114            .request_device(&wgpu::DeviceDescriptor {
115                label: Some("repose-offscreen"),
116                required_features: wgpu::Features::empty(),
117                required_limits: wgpu::Limits::default(),
118                memory_hints: wgpu::MemoryHints::Performance,
119                experimental_features: wgpu::ExperimentalFeatures::default(),
120                trace: wgpu::Trace::Off,
121            })
122            .await?;
123        let renderer = WgpuSceneRenderer::from_device_with_working_space_msaa(
124            device,
125            queue,
126            format,
127            msaa,
128            working_space_msaa,
129        );
130        Self::from_renderer(renderer, width, height)
131    }
132
133    /// Blocking [`new`](Self::new). Must not run on the wasm main thread
134    /// without block support (debug-asserted); use `render_rgba_async` there.
135    pub fn new_blocking(width: u32, height: u32, msaa: u32) -> Result<Self> {
136        pollster::block_on(Self::new(width, height, msaa))
137    }
138
139    /// Shared-device: reuse Device/Queue, no Adapter. Dimensions clamp to
140    /// ≥ 1; `msaa` clamps to ≥ 1 (no adapter-based picking without one —
141    /// prefer [`from_device_with_adapter`](Self::from_device_with_adapter)
142    /// when an adapter is handy).
143    pub fn from_device(
144        device: wgpu::Device,
145        queue: wgpu::Queue,
146        width: u32,
147        height: u32,
148        msaa: u32,
149    ) -> Result<Self> {
150        let width = width.max(1);
151        let height = height.max(1);
152        let format = TextureFormat::Rgba8UnormSrgb;
153        let renderer = WgpuSceneRenderer::from_device(device, queue, format, msaa.max(1));
154        Self::from_renderer(renderer, width, height)
155    }
156
157    /// Shared-device with adapter-based MSAA picking. Dimensions clamp to
158    /// ≥ 1.
159    pub fn from_device_with_adapter(
160        device: wgpu::Device,
161        queue: wgpu::Queue,
162        adapter: &wgpu::Adapter,
163        width: u32,
164        height: u32,
165        msaa: u32,
166    ) -> Result<Self> {
167        let width = width.max(1);
168        let height = height.max(1);
169        let format = TextureFormat::Rgba8UnormSrgb;
170        let msaa = crate::pick_surface_msaa(adapter, format, msaa);
171        let working_space_msaa =
172            crate::pick_surface_msaa_for_mode(adapter, TextureFormat::Rgba16Float, msaa, true);
173        let renderer = WgpuSceneRenderer::from_device_with_working_space_msaa(
174            device,
175            queue,
176            format,
177            msaa,
178            working_space_msaa,
179        );
180        Self::from_renderer(renderer, width, height)
181    }
182
183    /// Wrap an existing scene renderer with a fresh target + readback
184    /// buffer. Dimensions clamp to ≥ 1.
185    pub fn from_renderer(mut renderer: WgpuSceneRenderer, width: u32, height: u32) -> Result<Self> {
186        let max = renderer.device.limits().max_texture_dimension_2d;
187        let width = width.max(1).min(max);
188        let height = height.max(1).min(max);
189        renderer.resize(width, height);
190        let texture = renderer.device.create_texture(&wgpu::TextureDescriptor {
191            label: Some("repose-offscreen-tex"),
192            size: wgpu::Extent3d {
193                width,
194                height,
195                depth_or_array_layers: 1,
196            },
197            mip_level_count: 1,
198            sample_count: 1,
199            dimension: wgpu::TextureDimension::D2,
200            format: TextureFormat::Rgba8UnormSrgb,
201            usage: wgpu::TextureUsages::RENDER_ATTACHMENT
202                | wgpu::TextureUsages::TEXTURE_BINDING
203                | wgpu::TextureUsages::COPY_SRC,
204            view_formats: &[],
205        });
206        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
207        let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
208        let unpadded = width
209            .checked_mul(4)
210            .ok_or_else(|| anyhow::anyhow!("offscreen row size overflow"))?;
211        let padded = unpadded
212            .checked_add(align - 1)
213            .ok_or_else(|| anyhow::anyhow!("offscreen row alignment overflow"))?
214            / align
215            * align;
216        let buf_size = u64::from(padded)
217            .checked_mul(u64::from(height))
218            .ok_or_else(|| anyhow::anyhow!("offscreen readback size overflow"))?;
219        if buf_size > renderer.device.limits().max_buffer_size {
220            anyhow::bail!("offscreen readback exceeds device buffer limit");
221        }
222        let readback = renderer.device.create_buffer(&wgpu::BufferDescriptor {
223            label: Some("repose-offscreen-readback"),
224            size: buf_size,
225            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
226            mapped_at_creation: false,
227        });
228        Ok(Self {
229            renderer,
230            texture,
231            view,
232            readback,
233            width,
234            height,
235            padded_bytes_per_row: padded,
236        })
237    }
238
239    fn encode_rgba(&mut self, scene: &Scene, clear: Option<[f64; 4]>) -> wgpu::CommandBuffer {
240        let mut encoder =
241            self.renderer
242                .device
243                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
244                    label: Some("repose-offscreen-encoder"),
245                });
246        self.renderer.render_scene_to_encoder_with_texture(
247            scene,
248            &mut encoder,
249            &self.view,
250            Some(&self.texture),
251            clear,
252        );
253        encoder.copy_texture_to_buffer(
254            wgpu::TexelCopyTextureInfo {
255                texture: &self.texture,
256                mip_level: 0,
257                origin: wgpu::Origin3d::ZERO,
258                aspect: wgpu::TextureAspect::All,
259            },
260            wgpu::TexelCopyBufferInfo {
261                buffer: &self.readback,
262                layout: wgpu::TexelCopyBufferLayout {
263                    offset: 0,
264                    bytes_per_row: Some(self.padded_bytes_per_row),
265                    rows_per_image: Some(self.height),
266                },
267            },
268            wgpu::Extent3d {
269                width: self.width,
270                height: self.height,
271                depth_or_array_layers: 1,
272            },
273        );
274        encoder.finish()
275    }
276
277    /// Render `scene` and read back `width * height * 4` premultiplied sRGB
278    /// bytes (see the [struct contracts](Self)). `clear = None` uses
279    /// `scene.clear_color`; an empty scene therefore reads back the clear
280    /// color exactly. Blocks until mapping completes.
281    pub fn render_rgba(&mut self, scene: &Scene, clear: Option<[f64; 4]>) -> Result<Vec<u8>> {
282        // Blocking map on the wasm main thread (Window) can never complete..
283        #[cfg(all(target_family = "wasm", target_os = "unknown"))]
284        debug_assert!(
285            web_workers::web::has_block_support(),
286            "render_rgba (blocking) called on wasm main thread; use render_rgba_async"
287        );
288        self.renderer.begin_frame()?;
289        let result = (|| -> Result<Vec<u8>> {
290            let cmd = self.encode_rgba(scene, clear);
291            self.renderer.queue.submit(Some(cmd));
292            let slice = self.readback.slice(..);
293            let (tx, rx) = web_workers::sync::mpsc::channel();
294            slice.map_async(wgpu::MapMode::Read, move |r| {
295                let _ = tx.send_sync(r);
296            });
297            self.renderer.device.poll(PollType::wait_indefinitely())?;
298            rx.recv_sync()??;
299            let mapped = slice.get_mapped_range()?;
300            let result = strip_padding(&mapped, self.width, self.height, self.padded_bytes_per_row);
301            drop(mapped);
302            self.readback.unmap();
303            result
304        })();
305        self.renderer.end_frame();
306        result
307    }
308
309    #[cfg(target_arch = "wasm32")]
310    pub async fn render_rgba_async(
311        &mut self,
312        scene: &Scene,
313        clear: Option<[f64; 4]>,
314    ) -> Result<Vec<u8>> {
315        self.renderer.begin_frame()?;
316        let result = async {
317            let cmd = self.encode_rgba(scene, clear);
318            self.renderer.queue.submit(Some(cmd));
319            let slice = self.readback.slice(..);
320            let (tx, rx) = web_workers::sync::mpsc::channel();
321            slice.map_async(wgpu::MapMode::Read, {
322                let tx = tx.clone();
323                move |r| {
324                    let _ = tx.send_sync(r);
325                }
326            });
327            let start = web_workers::sync::Instant::now();
328            loop {
329                let _ = self.renderer.device.poll(PollType::Poll);
330                if let Ok(r) = rx.try_recv() {
331                    r?;
332                    break;
333                }
334                if start.elapsed().as_secs() > 30 {
335                    anyhow::bail!(
336                        "offscreen readback timed out after 30s (map callback never fired)"
337                    );
338                }
339                web_workers::web::yield_now_async(web_workers::web::YieldTime::UserVisible).await;
340            }
341            let mapped = slice.get_mapped_range()?;
342            let result =
343                strip_padding(&*mapped, self.width, self.height, self.padded_bytes_per_row);
344            drop(mapped);
345            self.readback.unmap();
346            result
347        }
348        .await;
349        self.renderer.end_frame();
350        result
351    }
352
353    pub fn renderer_mut(&mut self) -> &mut WgpuSceneRenderer {
354        &mut self.renderer
355    }
356    pub fn renderer(&self) -> &WgpuSceneRenderer {
357        &self.renderer
358    }
359
360    /// Current target width in pixels.
361    pub fn width(&self) -> u32 {
362        self.width
363    }
364
365    /// Current target height in pixels.
366    pub fn height(&self) -> u32 {
367        self.height
368    }
369
370    /// Resize the target + readback buffer. A no-op when the size already
371    /// matches (no reallocation). Dimensions clamp to ≥ 1. After this,
372    /// `render_rgba` returns `width * height * 4` bytes at the new size.
373    pub fn ensure_size(&mut self, width: u32, height: u32) -> Result<()> {
374        let max = self.renderer.device.limits().max_texture_dimension_2d;
375        let width = width.max(1).min(max);
376        let height = height.max(1).min(max);
377        if self.width == width && self.height == height {
378            return Ok(());
379        }
380        self.renderer.resize(width, height);
381        let texture = self
382            .renderer
383            .device
384            .create_texture(&wgpu::TextureDescriptor {
385                label: Some("repose-offscreen-tex"),
386                size: wgpu::Extent3d {
387                    width,
388                    height,
389                    depth_or_array_layers: 1,
390                },
391                mip_level_count: 1,
392                sample_count: 1,
393                dimension: wgpu::TextureDimension::D2,
394                format: TextureFormat::Rgba8UnormSrgb,
395                usage: wgpu::TextureUsages::RENDER_ATTACHMENT
396                    | wgpu::TextureUsages::TEXTURE_BINDING
397                    | wgpu::TextureUsages::COPY_SRC,
398                view_formats: &[],
399            });
400        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
401        let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
402        let unpadded = width
403            .checked_mul(4)
404            .ok_or_else(|| anyhow::anyhow!("offscreen row size overflow"))?;
405        let padded = unpadded
406            .checked_add(align - 1)
407            .ok_or_else(|| anyhow::anyhow!("offscreen row alignment overflow"))?
408            / align
409            * align;
410        let size = u64::from(padded)
411            .checked_mul(u64::from(height))
412            .ok_or_else(|| anyhow::anyhow!("offscreen readback size overflow"))?;
413        if size > self.renderer.device.limits().max_buffer_size {
414            anyhow::bail!("offscreen readback exceeds device buffer limit");
415        }
416        let readback = self.renderer.device.create_buffer(&wgpu::BufferDescriptor {
417            label: Some("repose-offscreen-readback"),
418            size,
419            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
420            mapped_at_creation: false,
421        });
422        self.texture = texture;
423        self.view = view;
424        self.readback = readback;
425        self.width = width;
426        self.height = height;
427        self.padded_bytes_per_row = padded;
428        Ok(())
429    }
430
431    pub async fn render_rgba_unified(
432        &mut self,
433        scene: &Scene,
434        clear: Option<[f64; 4]>,
435    ) -> Result<Vec<u8>> {
436        #[cfg(all(target_family = "wasm", target_os = "unknown"))]
437        {
438            if web_workers::web::has_block_support() {
439                return self.render_rgba(scene, clear);
440            } else {
441                return self.render_rgba_async(scene, clear).await;
442            }
443        }
444        #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
445        {
446            self.render_rgba(scene, clear)
447        }
448    }
449}
450
451fn strip_padding(mapped: &[u8], width: u32, height: u32, padded: u32) -> Result<Vec<u8>> {
452    let unpadded = usize::try_from(
453        width
454            .checked_mul(4)
455            .ok_or_else(|| anyhow::anyhow!("offscreen row size overflow"))?,
456    )
457    .map_err(|_| anyhow::anyhow!("offscreen row size exceeds addressable memory"))?;
458    let output_size = width
459        .checked_mul(height)
460        .and_then(|value| value.checked_mul(4))
461        .ok_or_else(|| anyhow::anyhow!("offscreen output size overflow"))?;
462    let capacity = usize::try_from(output_size)
463        .map_err(|_| anyhow::anyhow!("offscreen output exceeds addressable memory"))?;
464    let padded = usize::try_from(padded)
465        .map_err(|_| anyhow::anyhow!("offscreen padded row exceeds addressable memory"))?;
466    let mut out = Vec::with_capacity(capacity);
467    for y in 0..height {
468        let start = usize::try_from(y)
469            .ok()
470            .and_then(|row| row.checked_mul(padded))
471            .ok_or_else(|| anyhow::anyhow!("offscreen row offset overflow"))?;
472        let end = start
473            .checked_add(unpadded)
474            .ok_or_else(|| anyhow::anyhow!("offscreen row end overflow"))?;
475        let row = mapped
476            .get(start..end)
477            .ok_or_else(|| anyhow::anyhow!("offscreen readback row is truncated"))?;
478        out.extend_from_slice(row);
479    }
480    Ok(out)
481}
482
483pub fn map_buffer_blocking(slice: &wgpu::BufferSlice<'_>, device: &wgpu::Device) -> Result<()> {
484    let (tx, rx) = web_workers::sync::mpsc::channel();
485    slice.map_async(wgpu::MapMode::Read, move |r| {
486        let _ = tx.send_sync(r);
487    });
488    device.poll(PollType::wait_indefinitely())?;
489    rx.recv_sync()??;
490    Ok(())
491}
492
493#[cfg(target_arch = "wasm32")]
494pub async fn map_buffer_async(slice: &wgpu::BufferSlice<'_>, device: &wgpu::Device) -> Result<()> {
495    let (tx, rx) = web_workers::sync::mpsc::channel();
496    slice.map_async(wgpu::MapMode::Read, {
497        let tx = tx.clone();
498        move |r| {
499            let _ = tx.send_sync(r);
500        }
501    });
502    let start = web_workers::sync::Instant::now();
503    loop {
504        let _ = device.poll(PollType::Poll);
505        if let Ok(r) = rx.try_recv() {
506            r?;
507            break;
508        }
509        if start.elapsed().as_secs() > 30 {
510            anyhow::bail!("map_buffer timed out after 30s (map callback never fired)");
511        }
512        web_workers::web::yield_now_async(web_workers::web::YieldTime::UserVisible).await;
513    }
514    Ok(())
515}
516
517pub async fn map_buffer_unified(
518    slice: &wgpu::BufferSlice<'_>,
519    device: &wgpu::Device,
520) -> Result<()> {
521    #[cfg(all(target_family = "wasm", target_os = "unknown"))]
522    {
523        if web_workers::web::has_block_support() {
524            return map_buffer_blocking(slice, device);
525        } else {
526            return map_buffer_async(slice, device).await;
527        }
528    }
529    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
530    {
531        map_buffer_blocking(slice, device)
532    }
533}