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
63pub struct OffscreenRenderer {
64    renderer: WgpuSceneRenderer,
65    texture: wgpu::Texture,
66    view: wgpu::TextureView,
67    readback: wgpu::Buffer,
68    width: u32,
69    height: u32,
70    padded_bytes_per_row: u32,
71}
72
73impl OffscreenRenderer {
74    pub async fn new(width: u32, height: u32, msaa: u32) -> Result<Self> {
75        let width = width.max(1);
76        let height = height.max(1);
77        let instance = if cfg!(target_arch = "wasm32") {
78            let mut desc = wgpu::InstanceDescriptor::new_without_display_handle();
79            desc.backends = wgpu::Backends::BROWSER_WEBGPU | wgpu::Backends::GL;
80            wgpu::util::new_instance_with_webgpu_detection(desc).await
81        } else {
82            wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle())
83        };
84        let adapter = instance
85            .request_adapter(&wgpu::RequestAdapterOptions {
86                power_preference: wgpu::PowerPreference::HighPerformance,
87                compatible_surface: None,
88                ..Default::default()
89            })
90            .await?;
91        let format = TextureFormat::Rgba8UnormSrgb;
92        let msaa = crate::pick_surface_msaa(&adapter, format, msaa);
93        let (device, queue) = adapter
94            .request_device(&wgpu::DeviceDescriptor {
95                label: Some("repose-offscreen"),
96                required_features: wgpu::Features::empty(),
97                required_limits: wgpu::Limits::default(),
98                memory_hints: wgpu::MemoryHints::Performance,
99                experimental_features: wgpu::ExperimentalFeatures::default(),
100                trace: wgpu::Trace::Off,
101            })
102            .await?;
103        let renderer = WgpuSceneRenderer::from_device(device, queue, format, msaa);
104        Self::from_renderer(renderer, width, height)
105    }
106
107    pub fn new_blocking(width: u32, height: u32, msaa: u32) -> Result<Self> {
108        pollster::block_on(Self::new(width, height, msaa))
109    }
110
111    /// Shared-device: reuse Device/Queue, no Adapter.
112    pub fn from_device(
113        device: wgpu::Device,
114        queue: wgpu::Queue,
115        width: u32,
116        height: u32,
117        msaa: u32,
118    ) -> Result<Self> {
119        let width = width.max(1);
120        let height = height.max(1);
121        let format = TextureFormat::Rgba8UnormSrgb;
122        let renderer = WgpuSceneRenderer::from_device(device, queue, format, msaa.max(1));
123        Self::from_renderer(renderer, width, height)
124    }
125
126    pub fn from_device_with_adapter(
127        device: wgpu::Device,
128        queue: wgpu::Queue,
129        adapter: &wgpu::Adapter,
130        width: u32,
131        height: u32,
132        msaa: u32,
133    ) -> Result<Self> {
134        let width = width.max(1);
135        let height = height.max(1);
136        let format = TextureFormat::Rgba8UnormSrgb;
137        let msaa = crate::pick_surface_msaa(adapter, format, msaa);
138        let renderer = WgpuSceneRenderer::from_device(device, queue, format, msaa);
139        Self::from_renderer(renderer, width, height)
140    }
141
142    pub fn from_renderer(mut renderer: WgpuSceneRenderer, width: u32, height: u32) -> Result<Self> {
143        renderer.resize(width, height);
144        let texture = renderer.device.create_texture(&wgpu::TextureDescriptor {
145            label: Some("repose-offscreen-tex"),
146            size: wgpu::Extent3d {
147                width,
148                height,
149                depth_or_array_layers: 1,
150            },
151            mip_level_count: 1,
152            sample_count: 1,
153            dimension: wgpu::TextureDimension::D2,
154            format: TextureFormat::Rgba8UnormSrgb,
155            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
156            view_formats: &[],
157        });
158        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
159        let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
160        let unpadded = width * 4;
161        let padded = unpadded.div_ceil(align) * align;
162        let buf_size = (padded * height) as u64;
163        let readback = renderer.device.create_buffer(&wgpu::BufferDescriptor {
164            label: Some("repose-offscreen-readback"),
165            size: buf_size,
166            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
167            mapped_at_creation: false,
168        });
169        Ok(Self {
170            renderer,
171            texture,
172            view,
173            readback,
174            width,
175            height,
176            padded_bytes_per_row: padded,
177        })
178    }
179
180    fn encode_rgba(&mut self, scene: &Scene, clear: Option<[f64; 4]>) -> wgpu::CommandBuffer {
181        let mut encoder =
182            self.renderer
183                .device
184                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
185                    label: Some("repose-offscreen-encoder"),
186                });
187        self.renderer
188            .render_scene_to_encoder(scene, &mut encoder, &self.view, clear);
189        encoder.copy_texture_to_buffer(
190            wgpu::TexelCopyTextureInfo {
191                texture: &self.texture,
192                mip_level: 0,
193                origin: wgpu::Origin3d::ZERO,
194                aspect: wgpu::TextureAspect::All,
195            },
196            wgpu::TexelCopyBufferInfo {
197                buffer: &self.readback,
198                layout: wgpu::TexelCopyBufferLayout {
199                    offset: 0,
200                    bytes_per_row: Some(self.padded_bytes_per_row),
201                    rows_per_image: Some(self.height),
202                },
203            },
204            wgpu::Extent3d {
205                width: self.width,
206                height: self.height,
207                depth_or_array_layers: 1,
208            },
209        );
210        encoder.finish()
211    }
212
213    pub fn render_rgba(&mut self, scene: &Scene, clear: Option<[f64; 4]>) -> Result<Vec<u8>> {
214        // Blocking map on the wasm main thread (Window) can never complete..
215        #[cfg(all(target_family = "wasm", target_os = "unknown"))]
216        debug_assert!(
217            web_workers::web::has_block_support(),
218            "render_rgba (blocking) called on wasm main thread; use render_rgba_async"
219        );
220        let cmd = self.encode_rgba(scene, clear);
221        self.renderer.queue.submit(Some(cmd));
222        let slice = self.readback.slice(..);
223        let (tx, rx) = web_workers::sync::mpsc::channel();
224        slice.map_async(wgpu::MapMode::Read, move |r| {
225            let _ = tx.send_sync(r);
226        });
227        self.renderer.device.poll(PollType::wait_indefinitely())?;
228        rx.recv_sync()??;
229        let mapped = slice.get_mapped_range()?;
230        let out = strip_padding(&mapped, self.width, self.height, self.padded_bytes_per_row);
231        drop(mapped);
232        self.readback.unmap();
233        Ok(out)
234    }
235
236    #[cfg(target_arch = "wasm32")]
237    pub async fn render_rgba_async(
238        &mut self,
239        scene: &Scene,
240        clear: Option<[f64; 4]>,
241    ) -> Result<Vec<u8>> {
242        let cmd = self.encode_rgba(scene, clear);
243        self.renderer.queue.submit(Some(cmd));
244        let slice = self.readback.slice(..);
245        let (tx, rx) = web_workers::sync::mpsc::channel();
246        slice.map_async(wgpu::MapMode::Read, {
247            let tx = tx.clone();
248            move |r| {
249                let _ = tx.send_sync(r);
250            }
251        });
252        let start = web_workers::sync::Instant::now();
253        loop {
254            let _ = self.renderer.device.poll(PollType::Poll);
255            if let Ok(r) = rx.try_recv() {
256                r?;
257                break;
258            }
259            if start.elapsed().as_secs() > 30 {
260                anyhow::bail!("offscreen readback timed out after 30s (map callback never fired)");
261            }
262            web_workers::web::yield_now_async(web_workers::web::YieldTime::UserVisible).await;
263        }
264        let mapped = slice.get_mapped_range()?;
265        let out = strip_padding(&*mapped, self.width, self.height, self.padded_bytes_per_row);
266        drop(mapped);
267        self.readback.unmap();
268        Ok(out)
269    }
270
271    pub fn renderer_mut(&mut self) -> &mut WgpuSceneRenderer {
272        &mut self.renderer
273    }
274    pub fn renderer(&self) -> &WgpuSceneRenderer {
275        &self.renderer
276    }
277
278    pub fn ensure_size(&mut self, width: u32, height: u32) -> Result<()> {
279        let width = width.max(1);
280        let height = height.max(1);
281        if self.width == width && self.height == height {
282            return Ok(());
283        }
284        self.renderer.resize(width, height);
285        let texture = self
286            .renderer
287            .device
288            .create_texture(&wgpu::TextureDescriptor {
289                label: Some("repose-offscreen-tex"),
290                size: wgpu::Extent3d {
291                    width,
292                    height,
293                    depth_or_array_layers: 1,
294                },
295                mip_level_count: 1,
296                sample_count: 1,
297                dimension: wgpu::TextureDimension::D2,
298                format: TextureFormat::Rgba8UnormSrgb,
299                usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
300                view_formats: &[],
301            });
302        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
303        let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
304        let padded = (width * 4).div_ceil(align) * align;
305        let readback = self.renderer.device.create_buffer(&wgpu::BufferDescriptor {
306            label: Some("repose-offscreen-readback"),
307            size: (padded * height) as u64,
308            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
309            mapped_at_creation: false,
310        });
311        self.texture = texture;
312        self.view = view;
313        self.readback = readback;
314        self.width = width;
315        self.height = height;
316        self.padded_bytes_per_row = padded;
317        Ok(())
318    }
319
320    pub async fn render_rgba_unified(
321        &mut self,
322        scene: &Scene,
323        clear: Option<[f64; 4]>,
324    ) -> Result<Vec<u8>> {
325        #[cfg(all(target_family = "wasm", target_os = "unknown"))]
326        {
327            if web_workers::web::has_block_support() {
328                return self.render_rgba(scene, clear);
329            } else {
330                return self.render_rgba_async(scene, clear).await;
331            }
332        }
333        #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
334        {
335            self.render_rgba(scene, clear)
336        }
337    }
338}
339
340fn strip_padding(mapped: &[u8], width: u32, height: u32, padded: u32) -> Vec<u8> {
341    let unpadded = width * 4;
342    let mut out = Vec::with_capacity((width * height * 4) as usize);
343    for y in 0..height {
344        let start = (y * padded) as usize;
345        let end = start + unpadded as usize;
346        out.extend_from_slice(&mapped[start..end]);
347    }
348    out
349}
350
351pub fn map_buffer_blocking(slice: &wgpu::BufferSlice<'_>, device: &wgpu::Device) -> Result<()> {
352    let (tx, rx) = web_workers::sync::mpsc::channel();
353    slice.map_async(wgpu::MapMode::Read, move |r| {
354        let _ = tx.send_sync(r);
355    });
356    device.poll(PollType::wait_indefinitely())?;
357    rx.recv_sync()??;
358    Ok(())
359}
360
361#[cfg(target_arch = "wasm32")]
362pub async fn map_buffer_async(slice: &wgpu::BufferSlice<'_>, device: &wgpu::Device) -> Result<()> {
363    let (tx, rx) = web_workers::sync::mpsc::channel();
364    slice.map_async(wgpu::MapMode::Read, {
365        let tx = tx.clone();
366        move |r| {
367            let _ = tx.send_sync(r);
368        }
369    });
370    let start = web_workers::sync::Instant::now();
371    loop {
372        let _ = device.poll(PollType::Poll);
373        if let Ok(r) = rx.try_recv() {
374            r?;
375            break;
376        }
377        if start.elapsed().as_secs() > 30 {
378            anyhow::bail!("map_buffer timed out after 30s (map callback never fired)");
379        }
380        web_workers::web::yield_now_async(web_workers::web::YieldTime::UserVisible).await;
381    }
382    Ok(())
383}
384
385pub async fn map_buffer_unified(
386    slice: &wgpu::BufferSlice<'_>,
387    device: &wgpu::Device,
388) -> Result<()> {
389    #[cfg(all(target_family = "wasm", target_os = "unknown"))]
390    {
391        if web_workers::web::has_block_support() {
392            return map_buffer_blocking(slice, device);
393        } else {
394            return map_buffer_async(slice, device).await;
395        }
396    }
397    #[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
398    {
399        map_buffer_blocking(slice, device)
400    }
401}