Skip to main content

polyscope_render/
slice_plane_render.rs

1//! Slice plane visualization rendering.
2//!
3//! Renders slice planes as semi-transparent grids.
4
5use glam::{Mat4, Vec3, Vec4};
6use polyscope_core::slice_plane::SlicePlane;
7use std::num::NonZeroU64;
8use wgpu::util::DeviceExt;
9
10/// GPU representation of slice plane visualization uniforms.
11/// Matches the shader's `PlaneUniforms` struct.
12#[repr(C)]
13#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
14#[allow(clippy::pub_underscore_fields)]
15pub struct PlaneRenderUniforms {
16    /// Plane's object transform matrix.
17    pub transform: [[f32; 4]; 4],
18    /// Base color of the plane.
19    pub color: [f32; 4],
20    /// Color of the grid lines.
21    pub grid_color: [f32; 4],
22    /// Length scale for grid sizing.
23    pub length_scale: f32,
24    /// Size of the plane visualization (half-extent in each direction).
25    pub plane_size: f32,
26    /// Padding for alignment.
27    pub _padding: [f32; 2],
28}
29
30impl Default for PlaneRenderUniforms {
31    fn default() -> Self {
32        Self {
33            transform: Mat4::IDENTITY.to_cols_array_2d(),
34            color: [0.5, 0.5, 0.5, 1.0],
35            grid_color: [0.3, 0.3, 0.3, 1.0],
36            length_scale: 1.0,
37            plane_size: 0.05,
38            _padding: [0.0; 2],
39        }
40    }
41}
42
43/// Computes the transform matrix for a slice plane.
44///
45/// The plane lies in X=0 in local space, with Y and Z as tangent directions.
46/// The transform positions and orients the plane quad in world space.
47fn compute_plane_transform(origin: Vec3, normal: Vec3) -> Mat4 {
48    // Build orthonormal basis for the plane
49    // The normal becomes the local X axis (plane is at X=0)
50    let x_axis = normal.normalize();
51
52    // Choose an up vector that's not parallel to normal
53    let up = if x_axis.dot(Vec3::Y).abs() < 0.99 {
54        Vec3::Y
55    } else {
56        Vec3::Z
57    };
58
59    // Y axis is the first tangent direction
60    let y_axis = up.cross(x_axis).normalize();
61    // Z axis is the second tangent direction
62    let z_axis = x_axis.cross(y_axis).normalize();
63
64    // Create transform: translation + rotation
65    Mat4::from_cols(
66        Vec4::new(x_axis.x, x_axis.y, x_axis.z, 0.0),
67        Vec4::new(y_axis.x, y_axis.y, y_axis.z, 0.0),
68        Vec4::new(z_axis.x, z_axis.y, z_axis.z, 0.0),
69        Vec4::new(origin.x, origin.y, origin.z, 1.0),
70    )
71}
72
73/// Slice plane visualization render resources.
74pub struct SlicePlaneRenderData {
75    uniform_buffer: wgpu::Buffer,
76    bind_group: wgpu::BindGroup,
77}
78
79impl SlicePlaneRenderData {
80    /// Creates new slice plane render data.
81    ///
82    /// # Arguments
83    /// * `device` - The wgpu device
84    /// * `bind_group_layout` - The bind group layout
85    /// * `camera_buffer` - The camera uniform buffer
86    #[must_use]
87    pub fn new(
88        device: &wgpu::Device,
89        bind_group_layout: &wgpu::BindGroupLayout,
90        camera_buffer: &wgpu::Buffer,
91    ) -> Self {
92        let uniforms = PlaneRenderUniforms::default();
93
94        let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
95            label: Some("Slice Plane Uniform Buffer"),
96            contents: bytemuck::cast_slice(&[uniforms]),
97            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
98        });
99
100        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
101            label: Some("Slice Plane Bind Group"),
102            layout: bind_group_layout,
103            entries: &[
104                wgpu::BindGroupEntry {
105                    binding: 0,
106                    resource: camera_buffer.as_entire_binding(),
107                },
108                wgpu::BindGroupEntry {
109                    binding: 1,
110                    resource: uniform_buffer.as_entire_binding(),
111                },
112            ],
113        });
114
115        Self {
116            uniform_buffer,
117            bind_group,
118        }
119    }
120
121    /// Updates the uniforms for a specific slice plane.
122    ///
123    /// # Arguments
124    /// * `queue` - The wgpu queue
125    /// * `plane` - The slice plane to visualize
126    /// * `length_scale` - Scene length scale for grid sizing
127    pub fn update(&self, queue: &wgpu::Queue, plane: &SlicePlane, length_scale: f32) {
128        let transform = compute_plane_transform(plane.origin(), plane.normal());
129        let color = plane.color();
130
131        let uniforms = PlaneRenderUniforms {
132            transform: transform.to_cols_array_2d(),
133            color: [color.x, color.y, color.z, 1.0],
134            grid_color: [color.x * 0.6, color.y * 0.6, color.z * 0.6, 1.0],
135            length_scale,
136            plane_size: plane.plane_size(),
137            _padding: [0.0; 2],
138        };
139
140        queue.write_buffer(&self.uniform_buffer, 0, bytemuck::cast_slice(&[uniforms]));
141    }
142
143    /// Returns the bind group for rendering.
144    #[must_use]
145    pub fn bind_group(&self) -> &wgpu::BindGroup {
146        &self.bind_group
147    }
148
149    /// Draws the slice plane visualization.
150    ///
151    /// Draws 6 vertices (2 triangles) forming a quad. Both sides are rendered
152    /// because the pipeline has `cull_mode`: None.
153    pub fn draw<'a>(&'a self, render_pass: &mut wgpu::RenderPass<'a>) {
154        render_pass.set_bind_group(0, &self.bind_group, &[]);
155        render_pass.draw(0..6, 0..1);
156    }
157}
158
159/// Creates the bind group layout for slice plane rendering.
160#[must_use]
161pub fn create_slice_plane_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
162    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
163        label: Some("Slice Plane Bind Group Layout"),
164        entries: &[
165            // Camera uniforms
166            wgpu::BindGroupLayoutEntry {
167                binding: 0,
168                visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
169                ty: wgpu::BindingType::Buffer {
170                    ty: wgpu::BufferBindingType::Uniform,
171                    has_dynamic_offset: false,
172                    min_binding_size: NonZeroU64::new(272),
173                },
174                count: None,
175            },
176            // Plane uniforms
177            wgpu::BindGroupLayoutEntry {
178                binding: 1,
179                visibility: wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT,
180                ty: wgpu::BindingType::Buffer {
181                    ty: wgpu::BufferBindingType::Uniform,
182                    has_dynamic_offset: false,
183                    min_binding_size: NonZeroU64::new(112),
184                },
185                count: None,
186            },
187        ],
188    })
189}
190
191/// Creates the render pipeline for slice plane visualization.
192#[must_use]
193pub fn create_slice_plane_pipeline(
194    device: &wgpu::Device,
195    bind_group_layout: &wgpu::BindGroupLayout,
196    color_format: wgpu::TextureFormat,
197    depth_format: wgpu::TextureFormat,
198) -> wgpu::RenderPipeline {
199    let shader_source = include_str!("shaders/slice_plane.wgsl");
200    let shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
201        label: Some("Slice Plane Shader"),
202        source: wgpu::ShaderSource::Wgsl(shader_source.into()),
203    });
204
205    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
206        label: Some("Slice Plane Pipeline Layout"),
207        bind_group_layouts: &[bind_group_layout],
208        push_constant_ranges: &[],
209    });
210
211    device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
212        label: Some("Slice Plane Pipeline"),
213        layout: Some(&pipeline_layout),
214        vertex: wgpu::VertexState {
215            module: &shader_module,
216            entry_point: Some("vs_main"),
217            buffers: &[],
218            compilation_options: Default::default(),
219        },
220        fragment: Some(wgpu::FragmentState {
221            module: &shader_module,
222            entry_point: Some("fs_main"),
223            targets: &[Some(wgpu::ColorTargetState {
224                format: color_format,
225                blend: Some(wgpu::BlendState {
226                    color: wgpu::BlendComponent {
227                        src_factor: wgpu::BlendFactor::SrcAlpha,
228                        dst_factor: wgpu::BlendFactor::OneMinusSrcAlpha,
229                        operation: wgpu::BlendOperation::Add,
230                    },
231                    alpha: wgpu::BlendComponent {
232                        src_factor: wgpu::BlendFactor::One,
233                        dst_factor: wgpu::BlendFactor::OneMinusSrcAlpha,
234                        operation: wgpu::BlendOperation::Add,
235                    },
236                }),
237                write_mask: wgpu::ColorWrites::ALL,
238            })],
239            compilation_options: Default::default(),
240        }),
241        primitive: wgpu::PrimitiveState {
242            topology: wgpu::PrimitiveTopology::TriangleList,
243            strip_index_format: None,
244            front_face: wgpu::FrontFace::Ccw,
245            cull_mode: None, // Draw both sides
246            polygon_mode: wgpu::PolygonMode::Fill,
247            unclipped_depth: false,
248            conservative: false,
249        },
250        depth_stencil: Some(wgpu::DepthStencilState {
251            format: depth_format,
252            depth_write_enabled: true, // Write depth so scene geometry can occlude the plane
253            depth_compare: wgpu::CompareFunction::LessEqual, // Respect depth so nearer planes win
254            stencil: wgpu::StencilState::default(),
255            bias: wgpu::DepthBiasState::default(),
256        }),
257        multisample: wgpu::MultisampleState {
258            count: 1,
259            mask: !0,
260            alpha_to_coverage_enabled: false,
261        },
262        multiview: None,
263        cache: None,
264    })
265}