Skip to main content

polyscope_render/
shadow_map.rs

1//! Shadow map generation and blur passes.
2
3use glam::{Mat4, Vec3};
4use std::num::NonZeroU64;
5use wgpu::util::DeviceExt;
6
7/// Shadow map resolution.
8pub const SHADOW_MAP_SIZE: u32 = 2048;
9
10/// GPU representation of light uniforms.
11#[repr(C)]
12#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
13pub struct LightUniforms {
14    pub view_proj: [[f32; 4]; 4],
15    pub light_dir: [f32; 4],
16}
17
18impl Default for LightUniforms {
19    fn default() -> Self {
20        Self {
21            view_proj: Mat4::IDENTITY.to_cols_array_2d(),
22            light_dir: [0.5, -1.0, 0.3, 0.0],
23        }
24    }
25}
26
27/// GPU representation of blur uniforms.
28#[repr(C)]
29#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
30pub struct BlurUniforms {
31    pub direction: [f32; 2],
32    pub texel_size: [f32; 2],
33}
34
35impl Default for BlurUniforms {
36    fn default() -> Self {
37        let texel_size = 1.0 / SHADOW_MAP_SIZE as f32;
38        Self {
39            direction: [1.0, 0.0],
40            texel_size: [texel_size, texel_size],
41        }
42    }
43}
44
45/// Shadow map render resources.
46pub struct ShadowMapPass {
47    /// Shadow map depth texture (kept alive for GPU resource lifetime).
48    #[allow(dead_code)]
49    depth_texture: wgpu::Texture,
50    /// Shadow map depth view.
51    depth_view: wgpu::TextureView,
52    /// Light uniform buffer.
53    light_buffer: wgpu::Buffer,
54    /// Comparison sampler for shadow sampling.
55    comparison_sampler: wgpu::Sampler,
56    /// Bind group layout for consumers (ground plane shader).
57    pub shadow_bind_group_layout: wgpu::BindGroupLayout,
58    /// Bind group for shadow sampling.
59    shadow_bind_group: wgpu::BindGroup,
60}
61
62impl ShadowMapPass {
63    /// Creates a new shadow map pass.
64    #[must_use]
65    pub fn new(device: &wgpu::Device) -> Self {
66        // Create depth texture for shadow map
67        let depth_texture = device.create_texture(&wgpu::TextureDescriptor {
68            label: Some("Shadow Map Depth"),
69            size: wgpu::Extent3d {
70                width: SHADOW_MAP_SIZE,
71                height: SHADOW_MAP_SIZE,
72                depth_or_array_layers: 1,
73            },
74            mip_level_count: 1,
75            sample_count: 1,
76            dimension: wgpu::TextureDimension::D2,
77            format: wgpu::TextureFormat::Depth32Float,
78            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
79            view_formats: &[],
80        });
81
82        let depth_view = depth_texture.create_view(&wgpu::TextureViewDescriptor::default());
83
84        // Light uniform buffer
85        let light_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
86            label: Some("Light Uniform Buffer"),
87            contents: bytemuck::cast_slice(&[LightUniforms::default()]),
88            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
89        });
90
91        // Comparison sampler for shadow mapping
92        let comparison_sampler = device.create_sampler(&wgpu::SamplerDescriptor {
93            label: Some("Shadow Comparison Sampler"),
94            address_mode_u: wgpu::AddressMode::ClampToEdge,
95            address_mode_v: wgpu::AddressMode::ClampToEdge,
96            address_mode_w: wgpu::AddressMode::ClampToEdge,
97            mag_filter: wgpu::FilterMode::Linear,
98            min_filter: wgpu::FilterMode::Linear,
99            compare: Some(wgpu::CompareFunction::LessEqual),
100            ..Default::default()
101        });
102
103        // Bind group layout for shadow sampling (used by ground plane shader)
104        let shadow_bind_group_layout =
105            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
106                label: Some("Shadow Bind Group Layout"),
107                entries: &[
108                    // Light uniforms
109                    wgpu::BindGroupLayoutEntry {
110                        binding: 0,
111                        visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
112                        ty: wgpu::BindingType::Buffer {
113                            ty: wgpu::BufferBindingType::Uniform,
114                            has_dynamic_offset: false,
115                            min_binding_size: NonZeroU64::new(80),
116                        },
117                        count: None,
118                    },
119                    // Shadow map texture
120                    wgpu::BindGroupLayoutEntry {
121                        binding: 1,
122                        visibility: wgpu::ShaderStages::FRAGMENT,
123                        ty: wgpu::BindingType::Texture {
124                            sample_type: wgpu::TextureSampleType::Depth,
125                            view_dimension: wgpu::TextureViewDimension::D2,
126                            multisampled: false,
127                        },
128                        count: None,
129                    },
130                    // Shadow comparison sampler
131                    wgpu::BindGroupLayoutEntry {
132                        binding: 2,
133                        visibility: wgpu::ShaderStages::FRAGMENT,
134                        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Comparison),
135                        count: None,
136                    },
137                ],
138            });
139
140        // Create bind group for shadow sampling
141        let shadow_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
142            label: Some("Shadow Bind Group"),
143            layout: &shadow_bind_group_layout,
144            entries: &[
145                wgpu::BindGroupEntry {
146                    binding: 0,
147                    resource: light_buffer.as_entire_binding(),
148                },
149                wgpu::BindGroupEntry {
150                    binding: 1,
151                    resource: wgpu::BindingResource::TextureView(&depth_view),
152                },
153                wgpu::BindGroupEntry {
154                    binding: 2,
155                    resource: wgpu::BindingResource::Sampler(&comparison_sampler),
156                },
157            ],
158        });
159
160        Self {
161            depth_texture,
162            depth_view,
163            light_buffer,
164            comparison_sampler,
165            shadow_bind_group_layout,
166            shadow_bind_group,
167        }
168    }
169
170    /// Computes the light view-projection matrix for shadow mapping.
171    ///
172    /// Creates an orthographic projection from the light's perspective that
173    /// encompasses the scene.
174    #[must_use]
175    pub fn compute_light_matrix(scene_center: Vec3, scene_radius: f32, light_dir: Vec3) -> Mat4 {
176        let light_dir = light_dir.normalize();
177        let light_pos = scene_center - light_dir * scene_radius * 2.0;
178
179        // Find a stable up vector that's not parallel to light direction
180        let up = if light_dir.y.abs() > 0.99 {
181            Vec3::Z
182        } else {
183            Vec3::Y
184        };
185
186        let view = Mat4::look_at_rh(light_pos, scene_center, up);
187        let proj = Mat4::orthographic_rh(
188            -scene_radius,
189            scene_radius,
190            -scene_radius,
191            scene_radius,
192            0.1,
193            scene_radius * 4.0,
194        );
195        proj * view
196    }
197
198    /// Updates the light uniforms.
199    pub fn update_light(&self, queue: &wgpu::Queue, view_proj: Mat4, light_dir: Vec3) {
200        let uniforms = LightUniforms {
201            view_proj: view_proj.to_cols_array_2d(),
202            light_dir: [light_dir.x, light_dir.y, light_dir.z, 0.0],
203        };
204        queue.write_buffer(&self.light_buffer, 0, bytemuck::cast_slice(&[uniforms]));
205    }
206
207    /// Clears the shadow map depth buffer.
208    ///
209    /// Call this before rendering scene objects to the shadow map.
210    pub fn begin_shadow_pass<'a>(
211        &'a self,
212        encoder: &'a mut wgpu::CommandEncoder,
213    ) -> wgpu::RenderPass<'a> {
214        encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
215            label: Some("Shadow Map Pass"),
216            color_attachments: &[],
217            depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
218                view: &self.depth_view,
219                depth_ops: Some(wgpu::Operations {
220                    load: wgpu::LoadOp::Clear(1.0),
221                    store: wgpu::StoreOp::Store,
222                }),
223                stencil_ops: None,
224            }),
225            ..Default::default()
226        })
227    }
228
229    /// Returns the shadow map depth view.
230    #[must_use]
231    pub fn depth_view(&self) -> &wgpu::TextureView {
232        &self.depth_view
233    }
234
235    /// Returns the light uniform buffer.
236    #[must_use]
237    pub fn light_buffer(&self) -> &wgpu::Buffer {
238        &self.light_buffer
239    }
240
241    /// Returns the comparison sampler.
242    #[must_use]
243    pub fn comparison_sampler(&self) -> &wgpu::Sampler {
244        &self.comparison_sampler
245    }
246
247    /// Returns the bind group for shadow sampling.
248    #[must_use]
249    pub fn shadow_bind_group(&self) -> &wgpu::BindGroup {
250        &self.shadow_bind_group
251    }
252
253    /// Returns the bind group layout for shadow sampling.
254    #[must_use]
255    pub fn shadow_bind_group_layout(&self) -> &wgpu::BindGroupLayout {
256        &self.shadow_bind_group_layout
257    }
258}