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
112struct BlitShared {
113    shader: wgpu::ShaderModule,
114    layout: wgpu::BindGroupLayout,
115    sampler: wgpu::Sampler,
116    pipeline_layout: wgpu::PipelineLayout,
117}
118
119struct Target {
120    key: (wgpu::TextureFormat, u32, u32, u32),
121    #[allow(dead_code)]
122    scene: wgpu::Texture,
123    scene_view: wgpu::TextureView,
124    #[allow(dead_code)]
125    depth: wgpu::Texture,
126    depth_view: wgpu::TextureView,
127    blit_bind: wgpu::BindGroup,
128    last_used_tick: u64,
129    last_used_frame: u64,
130    bytes: u64,
131}
132
133impl DepthComposite {
134    /// Fetch the composite store from callback resources (creating it on
135    /// first use). One store per render pass; ids disambiguate viewports.
136    pub fn get(resources: &mut CallbackResources) -> &mut Self {
137        resources.get_or_insert_with::<Self>()
138    }
139
140    pub fn begin_frame(&mut self) {
141        self.frame_index = self.frame_index.wrapping_add(1);
142    }
143
144    pub fn end_frame(&mut self) {
145        let frame_index = self.frame_index;
146        let mut removed = 0u64;
147        self.targets.retain(|_, target| {
148            if frame_index.saturating_sub(target.last_used_frame) <= TARGET_IDLE_FRAMES {
149                true
150            } else {
151                removed = removed.saturating_add(target.bytes);
152                false
153            }
154        });
155        self.bytes_total = self.bytes_total.saturating_sub(removed);
156    }
157
158    fn shared(&mut self, device: &wgpu::Device) -> &BlitShared {
159        if self.shared.is_none() {
160            let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
161                label: Some("depth_composite_blit"),
162                source: wgpu::ShaderSource::Wgsl(BLIT_WGSL.into()),
163            });
164            let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
165                label: Some("depth_composite_blit_bgl"),
166                entries: &[
167                    wgpu::BindGroupLayoutEntry {
168                        binding: 0,
169                        visibility: wgpu::ShaderStages::FRAGMENT,
170                        ty: wgpu::BindingType::Texture {
171                            sample_type: wgpu::TextureSampleType::Float { filterable: true },
172                            view_dimension: wgpu::TextureViewDimension::D2,
173                            multisampled: false,
174                        },
175                        count: None,
176                    },
177                    wgpu::BindGroupLayoutEntry {
178                        binding: 1,
179                        visibility: wgpu::ShaderStages::FRAGMENT,
180                        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
181                        count: None,
182                    },
183                ],
184            });
185            let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
186                label: Some("depth_composite_blit_sampler"),
187                address_mode_u: wgpu::AddressMode::ClampToEdge,
188                address_mode_v: wgpu::AddressMode::ClampToEdge,
189                address_mode_w: wgpu::AddressMode::ClampToEdge,
190                mag_filter: wgpu::FilterMode::Linear,
191                min_filter: wgpu::FilterMode::Linear,
192                mipmap_filter: wgpu::MipmapFilterMode::Nearest,
193                lod_min_clamp: 0.0,
194                lod_max_clamp: 0.0,
195                compare: None,
196                anisotropy_clamp: 1,
197                border_color: None,
198            });
199            let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
200                label: Some("depth_composite_blit_pl"),
201                bind_group_layouts: &[Some(&layout)],
202                immediate_size: 0,
203            });
204            self.shared = Some(BlitShared {
205                shader,
206                layout,
207                sampler,
208                pipeline_layout,
209            });
210        }
211        self.shared.as_ref().expect("shared blit resources")
212    }
213
214    fn pipeline(
215        &mut self,
216        device: &wgpu::Device,
217        target_format: wgpu::TextureFormat,
218        sample_count: u32,
219    ) -> &wgpu::RenderPipeline {
220        let pipeline_key = (target_format, sample_count);
221        if !self.pipelines.contains_key(&pipeline_key) {
222            let shared = self.shared(device);
223            let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
224                label: Some("depth_composite_blit"),
225                layout: Some(&shared.pipeline_layout),
226                vertex: wgpu::VertexState {
227                    module: &shared.shader,
228                    entry_point: Some("vs_main"),
229                    buffers: &[],
230                    compilation_options: Default::default(),
231                },
232                fragment: Some(wgpu::FragmentState {
233                    module: &shared.shader,
234                    entry_point: Some("fs_main"),
235                    targets: &[Some(wgpu::ColorTargetState {
236                        format: target_format,
237                        // Premultiplied source over the existing UI: a
238                        // transparent scene composites instead of erasing
239                        // whatever the UI pass already drew beneath it.
240                        blend: Some(wgpu::BlendState::PREMULTIPLIED_ALPHA_BLENDING),
241                        write_mask: wgpu::ColorWrites::ALL,
242                    })],
243                    compilation_options: Default::default(),
244                }),
245                primitive: wgpu::PrimitiveState {
246                    topology: wgpu::PrimitiveTopology::TriangleList,
247                    ..Default::default()
248                },
249                // depth ops disabled, so the blit never disturbs UI depth/stencil.
250                depth_stencil: Some(wgpu::DepthStencilState {
251                    format: DEPTH_FORMAT,
252                    depth_write_enabled: Some(false),
253                    depth_compare: Some(wgpu::CompareFunction::Always),
254                    stencil: wgpu::StencilState {
255                        front: wgpu::StencilFaceState {
256                            compare: wgpu::CompareFunction::LessEqual,
257                            ..Default::default()
258                        },
259                        back: wgpu::StencilFaceState {
260                            compare: wgpu::CompareFunction::LessEqual,
261                            ..Default::default()
262                        },
263                        read_mask: 0xFF,
264                        write_mask: 0,
265                    },
266                    bias: wgpu::DepthBiasState::default(),
267                }),
268                multisample: wgpu::MultisampleState {
269                    count: sample_count,
270                    mask: !0,
271                    alpha_to_coverage_enabled: false,
272                },
273                multiview_mask: None,
274                cache: None,
275            });
276            self.pipelines.insert(pipeline_key, pipeline);
277        }
278        &self.pipelines[&pipeline_key]
279    }
280
281    /// Ensure the offscreen target for `id` at `w`x`h` (recreates on
282    /// format/sample/size change, like every other viewport-owned target).
283    /// Dimensions clamp to >= 1. Errors name the reason the target could
284    /// not be allocated; the other viewports are left untouched.
285    #[allow(clippy::too_many_arguments)] // (device, screen, id, w, h), mirrors ensure_resources conventions
286    pub fn ensure(
287        &mut self,
288        device: &wgpu::Device,
289        screen: &ScreenDescriptor,
290        id: &str,
291        w: u32,
292        h: u32,
293    ) -> anyhow::Result<()> {
294        let w = w.max(1);
295        let h = h.max(1);
296        let key = (screen.target_format, screen.sample_count, w, h);
297        self.next_tick = self.next_tick.wrapping_add(1);
298        let tick = self.next_tick;
299        if let Some(target) = self.targets.get_mut(id)
300            && target.key == key
301        {
302            target.last_used_tick = tick;
303            target.last_used_frame = self.frame_index;
304            return Ok(());
305        }
306        let scene_bytes = screen
307            .target_format
308            .theoretical_memory_footprint(wgpu::Extent3d {
309                width: w,
310                height: h,
311                depth_or_array_layers: 1,
312            });
313        let depth_bytes = DEPTH_FORMAT.theoretical_memory_footprint(wgpu::Extent3d {
314            width: w,
315            height: h,
316            depth_or_array_layers: 1,
317        });
318        let bytes = scene_bytes.saturating_add(depth_bytes);
319        let max_dimension = device.limits().max_texture_dimension_2d;
320        if w > max_dimension || h > max_dimension {
321            anyhow::bail!(
322                "depth composite {id}: {w}x{h} exceeds the device texture limit {max_dimension}"
323            );
324        }
325        if bytes > MAX_DEPTH_RESOURCE_BYTES {
326            anyhow::bail!(
327                "depth composite {id}: {w}x{h} needs {bytes} bytes, over the {MAX_DEPTH_RESOURCE_BYTES} byte limit"
328            );
329        }
330        // `id`'s own bytes are replaced, not added, so the budget only ever
331        // evicts siblings, and never leaves `id` without a target.
332        let replaced = self.targets.get(id).map_or(0, |old| old.bytes);
333        while self
334            .bytes_total
335            .saturating_sub(replaced)
336            .saturating_add(bytes)
337            > MAX_DEPTH_RESOURCE_BYTES
338        {
339            let Some(victim) = self
340                .targets
341                .iter()
342                .filter(|(key, _)| key.as_str() != id)
343                .min_by_key(|(_, target)| target.last_used_tick)
344                .map(|(key, _)| key.clone())
345            else {
346                anyhow::bail!(
347                    "depth composite {id}: no room in the {MAX_DEPTH_RESOURCE_BYTES} byte budget"
348                );
349            };
350            if let Some(old) = self.targets.remove(&victim) {
351                self.bytes_total = self.bytes_total.saturating_sub(old.bytes);
352            }
353        }
354        let shared = self.shared(device);
355        let scene = device.create_texture(&wgpu::TextureDescriptor {
356            label: Some("depth_composite_scene"),
357            size: wgpu::Extent3d {
358                width: w,
359                height: h,
360                depth_or_array_layers: 1,
361            },
362            mip_level_count: 1,
363            sample_count: 1,
364            dimension: wgpu::TextureDimension::D2,
365            format: screen.target_format,
366            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
367            view_formats: &[],
368        });
369        let scene_view = scene.create_view(&wgpu::TextureViewDescriptor::default());
370        let depth = device.create_texture(&wgpu::TextureDescriptor {
371            label: Some("depth_composite_depth"),
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: DEPTH_FORMAT,
381            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
382            view_formats: &[],
383        });
384        let depth_view = depth.create_view(&wgpu::TextureViewDescriptor::default());
385        let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {
386            label: Some("depth_composite_blit_bg"),
387            layout: &shared.layout,
388            entries: &[
389                wgpu::BindGroupEntry {
390                    binding: 0,
391                    resource: wgpu::BindingResource::TextureView(&scene_view),
392                },
393                wgpu::BindGroupEntry {
394                    binding: 1,
395                    resource: wgpu::BindingResource::Sampler(&shared.sampler),
396                },
397            ],
398        });
399        self.pipeline(device, screen.target_format, screen.sample_count);
400        let target = Target {
401            key,
402            scene,
403            scene_view,
404            depth,
405            depth_view,
406            blit_bind: bind,
407            last_used_tick: tick,
408            last_used_frame: self.frame_index,
409            bytes,
410        };
411        if let Some(old) = self.targets.insert(id.to_string(), target) {
412            self.bytes_total = self.bytes_total.saturating_sub(old.bytes);
413        }
414        self.bytes_total = self.bytes_total.saturating_add(bytes);
415        while self.targets.len() > MAX_DEPTH_TARGETS {
416            let Some(victim) = self
417                .targets
418                .iter()
419                .filter(|(key, _)| key.as_str() != id)
420                .min_by_key(|(_, target)| target.last_used_tick)
421                .map(|(key, _)| key.clone())
422            else {
423                break;
424            };
425            if let Some(old) = self.targets.remove(&victim) {
426                self.bytes_total = self.bytes_total.saturating_sub(old.bytes);
427            }
428        }
429        Ok(())
430    }
431
432    /// Begin the offscreen scene pass for `id` (clearing color to `clear`
433    /// and depth to 1.0). The caller issues its depth-tested draws inside
434    /// the returned pass, then ends it; [`blit`](Self::blit) composites in
435    /// `paint`. Returns `false` when `id` has no target (call [`ensure`](Self::ensure) first).
436    ///
437    /// The depth and stencil attachments are discarded when the pass ends:
438    /// [`blit`](Self::blit) only samples the color target, so nothing in
439    /// this crate reads them afterwards. A caller that needs the depth
440    /// buffer afterwards must keep its own copy.
441    ///
442    /// `clear` uses straight (un-premultiplied) alpha, which the blit
443    /// premultiplies before compositing.
444    pub fn begin_scene<'a>(
445        &'a self,
446        id: &str,
447        encoder: &'a mut wgpu::CommandEncoder,
448        clear: [f32; 4],
449    ) -> Option<wgpu::RenderPass<'a>> {
450        let t = self.targets.get(id)?;
451        let (w, h) = (t.key.2 as f32, t.key.3 as f32);
452        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
453            label: Some("depth_composite_scene"),
454            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
455                view: &t.scene_view,
456                depth_slice: None,
457                resolve_target: None,
458                ops: wgpu::Operations {
459                    load: wgpu::LoadOp::Clear(wgpu::Color {
460                        r: clear[0] as f64,
461                        g: clear[1] as f64,
462                        b: clear[2] as f64,
463                        a: clear[3] as f64,
464                    }),
465                    store: wgpu::StoreOp::Store,
466                },
467            })],
468            depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
469                view: &t.depth_view,
470                // Nothing reads the offscreen depth/stencil after this pass:
471                // the blit declares `depth_write_enabled: false` /
472                // `depth_compare: Always`.
473                depth_ops: Some(wgpu::Operations {
474                    load: wgpu::LoadOp::Clear(1.0),
475                    store: wgpu::StoreOp::Discard,
476                }),
477                stencil_ops: Some(wgpu::Operations {
478                    load: wgpu::LoadOp::Clear(0),
479                    store: wgpu::StoreOp::Discard,
480                }),
481            }),
482            timestamp_writes: None,
483            occlusion_query_set: None,
484            multiview_mask: None,
485        });
486        pass.set_viewport(0.0, 0.0, w, h, 0.0, 1.0);
487        Some(pass)
488    }
489
490    /// Depth format backing the scene target (for caller pipelines).
491    pub fn depth_format(&self, id: &str) -> Option<wgpu::TextureFormat> {
492        self.targets.get(id).map(|_| DEPTH_FORMAT)
493    }
494
495    /// Composite the offscreen scene for `id` into the main pass. The
496    /// renderer has already set the viewport to the callback rect, which
497    /// matches the offscreen texture 1:1 (both come from the painted frame
498    /// geometry). Composites over the UI already drawn this frame, so a
499    /// transparent scene keeps it visible. No-op when `id` has no target.
500    pub fn blit<P: BlitRenderPass>(&self, id: &str, rpass: &mut P) {
501        let Some(t) = self.targets.get(id) else {
502            return;
503        };
504        let Some(pipeline) = self.pipelines.get(&(t.key.0, t.key.1)) else {
505            log::warn!("depth composite {id}: blit pipeline missing");
506            return;
507        };
508        rpass.set_blit_pipeline(pipeline);
509        rpass.set_blit_bind_group(0, &t.blit_bind, &[]);
510        rpass.draw_blit_triangle();
511    }
512
513    pub fn blit_callback(&self, id: &str, rpass: &mut CallbackRenderPass<'_, '_>) {
514        self.blit(id, rpass);
515    }
516}