Skip to main content

repose_render_wgpu/
depth_composite.rs

1//! Depth-capable scene composite for [`WgpuCallback`](super::WgpuCallback)s.
2//!
3//! The shared UI pass this crate executes carries stencil but no depth ops
4//! (`depth_ops: None`), so pipelines that write depth are rejected inside
5//! it by validation. This module is the supported path for depth-tested
6//! content (3D viewports): render the scene into a caller-owned offscreen
7//! target (with its own depth texture) during `prepare`, then draw the
8//! target back as a fullscreen triangle in `paint`.
9//!
10//! One [`DepthComposite`] per viewport id (same ownership split as the
11//! sprite batch in `repame-sprite`): each id owns its target, so
12//! overlapping viewports never share depth. The blit honors the UI
13//! stencil contract (`LessEqual`) and touches no depth, so clips keep
14//! working while depth stays viewport-local. The blit composites over the
15//! UI instead of replacing it, so a transparent scene leaves the interface
16//! beneath it visible.
17
18use std::collections::HashMap;
19
20use super::{CallbackRenderPass, CallbackResources, ScreenDescriptor};
21
22const MAX_DEPTH_RESOURCE_BYTES: u64 = 256 * 1024 * 1024;
23const MAX_DEPTH_TARGETS: usize = 64;
24/// Frames a target survives without an `ensure`. Viewports are intermittent
25/// (animation, collapse, occlusion); dropping the textures every frame an
26/// embedder skips one forces a reallocation and a bind-group rebuild on the
27/// next paint, and used to throw away a shared blit pipeline with them.
28const TARGET_IDLE_FRAMES: u64 = 3;
29const DEPTH_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Depth24PlusStencil8;
30
31pub trait BlitRenderPass {
32    fn set_blit_pipeline(&mut self, pipeline: &wgpu::RenderPipeline);
33    fn set_blit_bind_group(&mut self, index: u32, bind_group: &wgpu::BindGroup, offsets: &[u32]);
34    fn draw_blit_triangle(&mut self);
35}
36
37impl BlitRenderPass for wgpu::RenderPass<'_> {
38    fn set_blit_pipeline(&mut self, pipeline: &wgpu::RenderPipeline) {
39        self.set_pipeline(pipeline);
40    }
41
42    fn set_blit_bind_group(&mut self, index: u32, bind_group: &wgpu::BindGroup, offsets: &[u32]) {
43        self.set_bind_group(index, bind_group, offsets);
44    }
45
46    fn draw_blit_triangle(&mut self) {
47        self.draw(0..3, 0..1);
48    }
49}
50
51impl BlitRenderPass for CallbackRenderPass<'_, '_> {
52    fn set_blit_pipeline(&mut self, pipeline: &wgpu::RenderPipeline) {
53        self.set_pipeline(pipeline);
54    }
55
56    fn set_blit_bind_group(&mut self, index: u32, bind_group: &wgpu::BindGroup, offsets: &[u32]) {
57        self.set_bind_group(index, bind_group, offsets);
58    }
59
60    fn draw_blit_triangle(&mut self) {
61        self.draw(0..3, 0..1);
62    }
63}
64
65/// Fullscreen textured triangle: samples the offscreen scene 1:1.
66///
67/// The scene target holds straight (un-premultiplied) alpha, and the blit
68/// composites over UI already painted this frame, so the fragment stage
69/// premultiplies and the pipeline uses `PREMULTIPLIED_ALPHA_BLENDING`.
70/// Opaque scenes (`a == 1`) pass through unchanged.
71const BLIT_WGSL: &str = r#"
72@group(0) @binding(0) var scene_tex: texture_2d<f32>;
73@group(0) @binding(1) var scene_smp: sampler;
74struct VsOut {
75    @builtin(position) pos: vec4<f32>,
76    @location(0) uv: vec2<f32>,
77};
78@vertex
79fn vs_main(@builtin(vertex_index) i: u32) -> VsOut {
80    let x = f32(i / 2u) * 4.0 - 1.0;
81    let y = f32(i % 2u) * 4.0 - 1.0;
82    var out: VsOut;
83    out.pos = vec4<f32>(x, y, 0.0, 1.0);
84    out.uv = vec2<f32>((x + 1.0) * 0.5, (1.0 - y) * 0.5);
85    return out;
86}
87@fragment
88fn fs_main(in: VsOut) -> @location(0) vec4<f32> {
89    let c = textureSample(scene_tex, scene_smp, in.uv);
90    return vec4<f32>(c.rgb * c.a, c.a);
91}
92"#;
93
94/// Viewport-owned offscreen scene target + depth buffer + blit pipeline.
95///
96/// Created through [`DepthComposite::ensure`] (per viewport id), drawn into
97/// with [`DepthComposite::begin_scene`], composited back with
98/// [`DepthComposite::blit`]. Textures recreate on format/sample/size
99/// change; buffer contents are per-frame and never retained.
100#[derive(Default)]
101pub struct DepthComposite {
102    targets: HashMap<String, Target>,
103    shared: Option<BlitShared>,
104    /// Blit pipelines depend only on `(target_format, sample_count)`, never
105    /// on the viewport size, so they outlive the textures they sample.
106    pipelines: HashMap<(wgpu::TextureFormat, u32), wgpu::RenderPipeline>,
107    next_tick: u64,
108    frame_index: u64,
109    bytes_total: u64,
110}
111
112/// Attachment behaviour of an offscreen scene pass
113/// ([`DepthComposite::begin_scene_ops`]).
114#[derive(Clone, Copy, Debug, PartialEq)]
115pub struct ScenePassOps {
116    /// Clear color when `Some`, else keep the target's contents instead (a
117    /// pass that draws over what an earlier pass left).
118    pub clear: Option<[f32; 4]>,
119    /// Keep the depth buffer's contents instead of clearing it.
120    pub load_depth: bool,
121    /// Keep the depth buffer after the pass. `false` discards it (the
122    /// single-pass default); a later pass that loads depth needs `true` on
123    /// the earlier one.
124    pub store_depth: bool,
125}
126
127struct BlitShared {
128    shader: wgpu::ShaderModule,
129    layout: wgpu::BindGroupLayout,
130    sampler: wgpu::Sampler,
131    pipeline_layout: wgpu::PipelineLayout,
132}
133
134struct Target {
135    key: (wgpu::TextureFormat, u32, u32, u32),
136    #[allow(dead_code)]
137    scene: wgpu::Texture,
138    scene_view: wgpu::TextureView,
139    #[allow(dead_code)]
140    depth: wgpu::Texture,
141    depth_view: wgpu::TextureView,
142    blit_bind: wgpu::BindGroup,
143    last_used_tick: u64,
144    last_used_frame: u64,
145    bytes: u64,
146}
147
148impl DepthComposite {
149    /// Fetch the composite store from callback resources (creating it on
150    /// first use). One store per render pass; ids disambiguate viewports.
151    pub fn get(resources: &mut CallbackResources) -> &mut Self {
152        resources.get_or_insert_with::<Self>()
153    }
154
155    pub fn begin_frame(&mut self) {
156        self.frame_index = self.frame_index.wrapping_add(1);
157    }
158
159    pub fn end_frame(&mut self) {
160        let frame_index = self.frame_index;
161        let mut removed = 0u64;
162        self.targets.retain(|_, target| {
163            if frame_index.saturating_sub(target.last_used_frame) <= TARGET_IDLE_FRAMES {
164                true
165            } else {
166                removed = removed.saturating_add(target.bytes);
167                false
168            }
169        });
170        self.bytes_total = self.bytes_total.saturating_sub(removed);
171    }
172
173    fn shared(&mut self, device: &wgpu::Device) -> &BlitShared {
174        if self.shared.is_none() {
175            let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
176                label: Some("depth_composite_blit"),
177                source: wgpu::ShaderSource::Wgsl(BLIT_WGSL.into()),
178            });
179            let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
180                label: Some("depth_composite_blit_bgl"),
181                entries: &[
182                    wgpu::BindGroupLayoutEntry {
183                        binding: 0,
184                        visibility: wgpu::ShaderStages::FRAGMENT,
185                        ty: wgpu::BindingType::Texture {
186                            sample_type: wgpu::TextureSampleType::Float { filterable: true },
187                            view_dimension: wgpu::TextureViewDimension::D2,
188                            multisampled: false,
189                        },
190                        count: None,
191                    },
192                    wgpu::BindGroupLayoutEntry {
193                        binding: 1,
194                        visibility: wgpu::ShaderStages::FRAGMENT,
195                        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
196                        count: None,
197                    },
198                ],
199            });
200            let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
201                label: Some("depth_composite_blit_sampler"),
202                address_mode_u: wgpu::AddressMode::ClampToEdge,
203                address_mode_v: wgpu::AddressMode::ClampToEdge,
204                address_mode_w: wgpu::AddressMode::ClampToEdge,
205                mag_filter: wgpu::FilterMode::Linear,
206                min_filter: wgpu::FilterMode::Linear,
207                mipmap_filter: wgpu::MipmapFilterMode::Nearest,
208                lod_min_clamp: 0.0,
209                lod_max_clamp: 0.0,
210                compare: None,
211                anisotropy_clamp: 1,
212                border_color: None,
213            });
214            let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
215                label: Some("depth_composite_blit_pl"),
216                bind_group_layouts: &[Some(&layout)],
217                immediate_size: 0,
218            });
219            self.shared = Some(BlitShared {
220                shader,
221                layout,
222                sampler,
223                pipeline_layout,
224            });
225        }
226        self.shared.as_ref().expect("shared blit resources")
227    }
228
229    fn pipeline(
230        &mut self,
231        device: &wgpu::Device,
232        target_format: wgpu::TextureFormat,
233        sample_count: u32,
234    ) -> &wgpu::RenderPipeline {
235        let pipeline_key = (target_format, sample_count);
236        if !self.pipelines.contains_key(&pipeline_key) {
237            let shared = self.shared(device);
238            let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
239                label: Some("depth_composite_blit"),
240                layout: Some(&shared.pipeline_layout),
241                vertex: wgpu::VertexState {
242                    module: &shared.shader,
243                    entry_point: Some("vs_main"),
244                    buffers: &[],
245                    compilation_options: Default::default(),
246                },
247                fragment: Some(wgpu::FragmentState {
248                    module: &shared.shader,
249                    entry_point: Some("fs_main"),
250                    targets: &[Some(wgpu::ColorTargetState {
251                        format: target_format,
252                        // Premultiplied source over the existing UI: a
253                        // transparent scene composites instead of erasing
254                        // whatever the UI pass already drew beneath it.
255                        blend: Some(wgpu::BlendState::PREMULTIPLIED_ALPHA_BLENDING),
256                        write_mask: wgpu::ColorWrites::ALL,
257                    })],
258                    compilation_options: Default::default(),
259                }),
260                primitive: wgpu::PrimitiveState {
261                    topology: wgpu::PrimitiveTopology::TriangleList,
262                    ..Default::default()
263                },
264                // depth ops disabled, so the blit never disturbs UI depth/stencil.
265                depth_stencil: Some(wgpu::DepthStencilState {
266                    format: DEPTH_FORMAT,
267                    depth_write_enabled: Some(false),
268                    depth_compare: Some(wgpu::CompareFunction::Always),
269                    stencil: wgpu::StencilState {
270                        front: wgpu::StencilFaceState {
271                            compare: wgpu::CompareFunction::LessEqual,
272                            ..Default::default()
273                        },
274                        back: wgpu::StencilFaceState {
275                            compare: wgpu::CompareFunction::LessEqual,
276                            ..Default::default()
277                        },
278                        read_mask: 0xFF,
279                        write_mask: 0,
280                    },
281                    bias: wgpu::DepthBiasState::default(),
282                }),
283                multisample: wgpu::MultisampleState {
284                    count: sample_count,
285                    mask: !0,
286                    alpha_to_coverage_enabled: false,
287                },
288                multiview_mask: None,
289                cache: None,
290            });
291            self.pipelines.insert(pipeline_key, pipeline);
292        }
293        &self.pipelines[&pipeline_key]
294    }
295
296    /// Ensure the offscreen target for `id` at `w`x`h` (recreates on
297    /// format/sample/size change, like every other viewport-owned target).
298    /// Dimensions clamp to >= 1. Errors name the reason the target could
299    /// not be allocated; the other viewports are left untouched.
300    #[allow(clippy::too_many_arguments)] // (device, screen, id, w, h), mirrors ensure_resources conventions
301    pub fn ensure(
302        &mut self,
303        device: &wgpu::Device,
304        screen: &ScreenDescriptor,
305        id: &str,
306        w: u32,
307        h: u32,
308    ) -> anyhow::Result<()> {
309        let w = w.max(1);
310        let h = h.max(1);
311        let key = (screen.target_format, screen.sample_count, w, h);
312        self.next_tick = self.next_tick.wrapping_add(1);
313        let tick = self.next_tick;
314        if let Some(target) = self.targets.get_mut(id)
315            && target.key == key
316        {
317            target.last_used_tick = tick;
318            target.last_used_frame = self.frame_index;
319            return Ok(());
320        }
321        let scene_bytes = screen
322            .target_format
323            .theoretical_memory_footprint(wgpu::Extent3d {
324                width: w,
325                height: h,
326                depth_or_array_layers: 1,
327            });
328        let depth_bytes = DEPTH_FORMAT.theoretical_memory_footprint(wgpu::Extent3d {
329            width: w,
330            height: h,
331            depth_or_array_layers: 1,
332        });
333        let bytes = scene_bytes.saturating_add(depth_bytes);
334        let max_dimension = device.limits().max_texture_dimension_2d;
335        if w > max_dimension || h > max_dimension {
336            anyhow::bail!(
337                "depth composite {id}: {w}x{h} exceeds the device texture limit {max_dimension}"
338            );
339        }
340        if bytes > MAX_DEPTH_RESOURCE_BYTES {
341            anyhow::bail!(
342                "depth composite {id}: {w}x{h} needs {bytes} bytes, over the {MAX_DEPTH_RESOURCE_BYTES} byte limit"
343            );
344        }
345        // `id`'s own bytes are replaced, not added, so the budget only ever
346        // evicts siblings, and never leaves `id` without a target.
347        let replaced = self.targets.get(id).map_or(0, |old| old.bytes);
348        while self
349            .bytes_total
350            .saturating_sub(replaced)
351            .saturating_add(bytes)
352            > MAX_DEPTH_RESOURCE_BYTES
353        {
354            let Some(victim) = self
355                .targets
356                .iter()
357                .filter(|(key, _)| key.as_str() != id)
358                .min_by_key(|(_, target)| target.last_used_tick)
359                .map(|(key, _)| key.clone())
360            else {
361                anyhow::bail!(
362                    "depth composite {id}: no room in the {MAX_DEPTH_RESOURCE_BYTES} byte budget"
363                );
364            };
365            if let Some(old) = self.targets.remove(&victim) {
366                self.bytes_total = self.bytes_total.saturating_sub(old.bytes);
367            }
368        }
369        let shared = self.shared(device);
370        let scene = device.create_texture(&wgpu::TextureDescriptor {
371            label: Some("depth_composite_scene"),
372            size: wgpu::Extent3d {
373                width: w,
374                height: h,
375                depth_or_array_layers: 1,
376            },
377            mip_level_count: 1,
378            sample_count: 1,
379            dimension: wgpu::TextureDimension::D2,
380            format: screen.target_format,
381            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
382            view_formats: &[],
383        });
384        let scene_view = scene.create_view(&wgpu::TextureViewDescriptor::default());
385        let depth = device.create_texture(&wgpu::TextureDescriptor {
386            label: Some("depth_composite_depth"),
387            size: wgpu::Extent3d {
388                width: w,
389                height: h,
390                depth_or_array_layers: 1,
391            },
392            mip_level_count: 1,
393            sample_count: 1,
394            dimension: wgpu::TextureDimension::D2,
395            format: DEPTH_FORMAT,
396            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
397            view_formats: &[],
398        });
399        let depth_view = depth.create_view(&wgpu::TextureViewDescriptor::default());
400        let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {
401            label: Some("depth_composite_blit_bg"),
402            layout: &shared.layout,
403            entries: &[
404                wgpu::BindGroupEntry {
405                    binding: 0,
406                    resource: wgpu::BindingResource::TextureView(&scene_view),
407                },
408                wgpu::BindGroupEntry {
409                    binding: 1,
410                    resource: wgpu::BindingResource::Sampler(&shared.sampler),
411                },
412            ],
413        });
414        self.pipeline(device, screen.target_format, screen.sample_count);
415        let target = Target {
416            key,
417            scene,
418            scene_view,
419            depth,
420            depth_view,
421            blit_bind: bind,
422            last_used_tick: tick,
423            last_used_frame: self.frame_index,
424            bytes,
425        };
426        if let Some(old) = self.targets.insert(id.to_string(), target) {
427            self.bytes_total = self.bytes_total.saturating_sub(old.bytes);
428        }
429        self.bytes_total = self.bytes_total.saturating_add(bytes);
430        while self.targets.len() > MAX_DEPTH_TARGETS {
431            let Some(victim) = self
432                .targets
433                .iter()
434                .filter(|(key, _)| key.as_str() != id)
435                .min_by_key(|(_, target)| target.last_used_tick)
436                .map(|(key, _)| key.clone())
437            else {
438                break;
439            };
440            if let Some(old) = self.targets.remove(&victim) {
441                self.bytes_total = self.bytes_total.saturating_sub(old.bytes);
442            }
443        }
444        Ok(())
445    }
446
447    /// Begin the offscreen scene pass for `id` (clearing color to `clear`
448    /// and depth to 1.0). The caller issues its depth-tested draws inside
449    /// the returned pass, then ends it; [`blit`](Self::blit) composites in
450    /// `paint`. Returns `false` when `id` has no target (call [`ensure`](Self::ensure) first).
451    ///
452    /// The depth and stencil attachments are discarded when the pass ends:
453    /// [`blit`](Self::blit) only samples the color target, so nothing in
454    /// this crate reads them afterwards. A caller that needs the depth
455    /// buffer afterwards uses [`begin_scene_ops`](Self::begin_scene_ops)
456    /// with `store_depth`.
457    ///
458    /// `clear` uses straight (un-premultiplied) alpha, which the blit
459    /// premultiplies before compositing.
460    pub fn begin_scene<'a>(
461        &'a self,
462        id: &str,
463        encoder: &'a mut wgpu::CommandEncoder,
464        clear: [f32; 4],
465    ) -> Option<wgpu::RenderPass<'a>> {
466        self.begin_scene_ops(
467            id,
468            encoder,
469            ScenePassOps {
470                clear: Some(clear),
471                load_depth: false,
472                store_depth: false,
473            },
474        )
475    }
476
477    /// [`begin_scene`](Self::begin_scene) with explicit attachment ops, for
478    /// a pass that draws over what an earlier pass left (see
479    /// [`ScenePassOps`]). The target keeps its contents only if the earlier
480    /// pass stored them.
481    pub fn begin_scene_ops<'a>(
482        &'a self,
483        id: &str,
484        encoder: &'a mut wgpu::CommandEncoder,
485        ops: ScenePassOps,
486    ) -> Option<wgpu::RenderPass<'a>> {
487        let t = self.targets.get(id)?;
488        let (w, h) = (t.key.2 as f32, t.key.3 as f32);
489        let color_load = match ops.clear {
490            Some(clear) => wgpu::LoadOp::Clear(wgpu::Color {
491                r: clear[0] as f64,
492                g: clear[1] as f64,
493                b: clear[2] as f64,
494                a: clear[3] as f64,
495            }),
496            None => wgpu::LoadOp::Load,
497        };
498        let depth_ops = wgpu::Operations {
499            load: if ops.load_depth {
500                wgpu::LoadOp::Load
501            } else {
502                wgpu::LoadOp::Clear(1.0)
503            },
504            // Nothing in this crate reads depth/stencil after a pass; a
505            // caller that loads depth asks for it to be kept too.
506            store: if ops.store_depth {
507                wgpu::StoreOp::Store
508            } else {
509                wgpu::StoreOp::Discard
510            },
511        };
512        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
513            label: Some("depth_composite_scene"),
514            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
515                view: &t.scene_view,
516                depth_slice: None,
517                resolve_target: None,
518                ops: wgpu::Operations {
519                    load: color_load,
520                    store: wgpu::StoreOp::Store,
521                },
522            })],
523            depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
524                view: &t.depth_view,
525                depth_ops: Some(depth_ops),
526                stencil_ops: Some(wgpu::Operations {
527                    load: wgpu::LoadOp::Clear(0),
528                    store: wgpu::StoreOp::Discard,
529                }),
530            }),
531            timestamp_writes: None,
532            occlusion_query_set: None,
533            multiview_mask: None,
534        });
535        pass.set_viewport(0.0, 0.0, w, h, 0.0, 1.0);
536        Some(pass)
537    }
538
539    /// Depth format backing the scene target (for caller pipelines).
540    pub fn depth_format(&self, id: &str) -> Option<wgpu::TextureFormat> {
541        self.targets.get(id).map(|_| DEPTH_FORMAT)
542    }
543
544    /// Composite the offscreen scene for `id` into the main pass. The
545    /// renderer has already set the viewport to the callback rect, which
546    /// matches the offscreen texture 1:1 (both come from the painted frame
547    /// geometry). Composites over the UI already drawn this frame, so a
548    /// transparent scene keeps it visible. No-op when `id` has no target.
549    pub fn blit<P: BlitRenderPass>(&self, id: &str, rpass: &mut P) {
550        let Some(t) = self.targets.get(id) else {
551            return;
552        };
553        let Some(pipeline) = self.pipelines.get(&(t.key.0, t.key.1)) else {
554            log::warn!("depth composite {id}: blit pipeline missing");
555            return;
556        };
557        rpass.set_blit_pipeline(pipeline);
558        rpass.set_blit_bind_group(0, &t.blit_bind, &[]);
559        rpass.draw_blit_triangle();
560    }
561
562    pub fn blit_callback(&self, id: &str, rpass: &mut CallbackRenderPass<'_, '_>) {
563        self.blit(id, rpass);
564    }
565}