Skip to main content

polyscope_render/
reflection_pass.rs

1//! Planar reflection rendering pass.
2
3use glam::Mat4;
4use std::num::NonZeroU64;
5use wgpu::util::DeviceExt;
6
7/// GPU representation of reflection uniforms.
8#[repr(C)]
9#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
10#[allow(clippy::pub_underscore_fields)]
11pub struct ReflectionUniforms {
12    pub reflection_matrix: [[f32; 4]; 4],
13    pub intensity: f32,
14    pub ground_height: f32,
15    pub _padding: [f32; 2],
16}
17
18impl Default for ReflectionUniforms {
19    fn default() -> Self {
20        Self {
21            reflection_matrix: Mat4::IDENTITY.to_cols_array_2d(),
22            intensity: 0.25,
23            ground_height: 0.0,
24            _padding: [0.0; 2],
25        }
26    }
27}
28
29/// Reflection pass render resources.
30///
31/// This pass handles planar reflections for the ground plane by:
32/// 1. Rendering the ground plane to stencil buffer (marking reflection region)
33/// 2. Rendering reflected scene geometry (only where stencil is set)
34/// 3. Blending the reflection with the final ground plane
35pub struct ReflectionPass {
36    /// Reflection uniform buffer.
37    uniform_buffer: wgpu::Buffer,
38    /// Bind group layout for reflection uniforms.
39    pub bind_group_layout: wgpu::BindGroupLayout,
40    /// Bind group for reflection uniforms.
41    bind_group: wgpu::BindGroup,
42}
43
44impl ReflectionPass {
45    /// Creates a new reflection pass.
46    #[must_use]
47    pub fn new(device: &wgpu::Device) -> Self {
48        // Create reflection uniform buffer
49        let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
50            label: Some("Reflection Uniform Buffer"),
51            contents: bytemuck::cast_slice(&[ReflectionUniforms::default()]),
52            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
53        });
54
55        // Create bind group layout for reflection uniforms
56        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
57            label: Some("Reflection Bind Group Layout"),
58            entries: &[wgpu::BindGroupLayoutEntry {
59                binding: 0,
60                visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
61                ty: wgpu::BindingType::Buffer {
62                    ty: wgpu::BufferBindingType::Uniform,
63                    has_dynamic_offset: false,
64                    min_binding_size: NonZeroU64::new(80),
65                },
66                count: None,
67            }],
68        });
69
70        // Create bind group
71        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
72            label: Some("Reflection Bind Group"),
73            layout: &bind_group_layout,
74            entries: &[wgpu::BindGroupEntry {
75                binding: 0,
76                resource: uniform_buffer.as_entire_binding(),
77            }],
78        });
79
80        Self {
81            uniform_buffer,
82            bind_group_layout,
83            bind_group,
84        }
85    }
86
87    /// Updates the reflection uniforms.
88    pub fn update_uniforms(
89        &self,
90        queue: &wgpu::Queue,
91        reflection_matrix: Mat4,
92        intensity: f32,
93        ground_height: f32,
94    ) {
95        let uniforms = ReflectionUniforms {
96            reflection_matrix: reflection_matrix.to_cols_array_2d(),
97            intensity,
98            ground_height,
99            _padding: [0.0; 2],
100        };
101        queue.write_buffer(&self.uniform_buffer, 0, bytemuck::cast_slice(&[uniforms]));
102    }
103
104    /// Returns the reflection bind group.
105    #[must_use]
106    pub fn bind_group(&self) -> &wgpu::BindGroup {
107        &self.bind_group
108    }
109
110    /// Returns the reflection bind group layout.
111    #[must_use]
112    pub fn bind_group_layout(&self) -> &wgpu::BindGroupLayout {
113        &self.bind_group_layout
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn test_reflection_uniforms_size() {
123        // Ensure uniform is correctly aligned for GPU
124        assert_eq!(
125            std::mem::size_of::<ReflectionUniforms>(),
126            64 + 4 + 4 + 8 // 4x4 matrix + intensity + ground_height + padding
127        );
128    }
129}