Skip to main content

polyscope_render/
surface_mesh_render.rs

1//! Surface mesh GPU rendering resources.
2
3use glam::{Vec3, Vec4};
4use wgpu::util::DeviceExt;
5
6/// Uniforms for surface mesh rendering.
7/// Note: Layout must match WGSL `MeshUniforms` exactly.
8/// WGSL `vec3<f32>` has 16-byte alignment, requiring extra padding.
9#[repr(C)]
10#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
11#[allow(clippy::pub_underscore_fields)]
12pub struct MeshUniforms {
13    /// Model transform matrix (must be first for alignment)
14    pub model_matrix: [[f32; 4]; 4],
15    /// Shading style: 0 = smooth, 1 = flat, 2 = tri-flat
16    pub shade_style: u32,
17    /// Show edges: 0 = off, 1 = on
18    pub show_edges: u32,
19    /// Edge width in pixels
20    pub edge_width: f32,
21    /// Surface transparency (0.0 = opaque, 1.0 = fully transparent)
22    pub transparency: f32,
23    /// Surface color (RGBA)
24    pub surface_color: [f32; 4],
25    /// Edge color (RGBA)
26    pub edge_color: [f32; 4],
27    /// Backface policy: 0 = identical, 1 = different, 2 = custom, 3 = cull
28    pub backface_policy: u32,
29    /// Slice plane clipping enable: 0 = off, 1 = on
30    pub slice_planes_enabled: u32,
31    /// Use per-vertex colors (1) or surface color (0)
32    pub use_vertex_color: u32,
33    /// Padding to align to 16 bytes
34    pub _pad1: f32,
35    /// Padding matching WGSL layout (12 bytes)
36    pub _pad2: [f32; 3],
37    /// Padding to align vec4 to 16 bytes
38    pub _pad3: f32,
39    /// Backface color (RGBA), used when `backface_policy` is custom
40    pub backface_color: [f32; 4],
41}
42
43/// Model uniforms for shadow rendering.
44/// Only contains the model matrix (64 bytes) - matches `shadow_map.wgsl` `ModelUniforms`.
45#[repr(C)]
46#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
47pub struct ShadowModelUniforms {
48    /// Model transform matrix.
49    pub model: [[f32; 4]; 4],
50}
51
52impl Default for ShadowModelUniforms {
53    fn default() -> Self {
54        Self {
55            model: [
56                [1.0, 0.0, 0.0, 0.0],
57                [0.0, 1.0, 0.0, 0.0],
58                [0.0, 0.0, 1.0, 0.0],
59                [0.0, 0.0, 0.0, 1.0],
60            ],
61        }
62    }
63}
64
65impl Default for MeshUniforms {
66    fn default() -> Self {
67        Self {
68            model_matrix: [
69                [1.0, 0.0, 0.0, 0.0],
70                [0.0, 1.0, 0.0, 0.0],
71                [0.0, 0.0, 1.0, 0.0],
72                [0.0, 0.0, 0.0, 1.0],
73            ],
74            shade_style: 0,                      // smooth shading
75            show_edges: 0,                       // edges off
76            edge_width: 1.0,                     // 1 pixel edge
77            transparency: 0.0,                   // fully opaque
78            surface_color: [0.5, 0.5, 0.5, 1.0], // gray
79            edge_color: [0.0, 0.0, 0.0, 1.0],    // black edges
80            backface_policy: 0,                  // identical to front
81            slice_planes_enabled: 1,
82            use_vertex_color: 0,
83            _pad1: 0.0,
84            _pad2: [0.0; 3],
85            _pad3: 0.0,
86            backface_color: [0.3, 0.3, 0.3, 1.0], // darker gray
87        }
88    }
89}
90
91/// GPU resources for rendering a surface mesh.
92pub struct SurfaceMeshRenderData {
93    /// Position buffer (storage buffer, vec4 for alignment).
94    pub vertex_buffer: wgpu::Buffer,
95    /// Index buffer (triangle indices).
96    pub index_buffer: wgpu::Buffer,
97    /// Normal buffer (vertex normals, vec4 for alignment).
98    pub normal_buffer: wgpu::Buffer,
99    /// Barycentric coordinate buffer for wireframe rendering.
100    /// Each triangle vertex gets `[1,0,0]`, `[0,1,0]`, `[0,0,1]`.
101    pub barycentric_buffer: wgpu::Buffer,
102    /// Color buffer (per-vertex colors for quantities, vec4).
103    pub color_buffer: wgpu::Buffer,
104    /// Edge is real buffer - marks which edges are real polygon edges vs triangulation internal.
105    pub edge_is_real_buffer: wgpu::Buffer,
106    /// Uniform buffer for mesh-specific settings.
107    pub uniform_buffer: wgpu::Buffer,
108    /// Bind group for this surface mesh.
109    pub bind_group: wgpu::BindGroup,
110    /// Number of triangles.
111    pub num_triangles: u32,
112    /// Number of indices (`num_triangles` * 3).
113    pub num_indices: u32,
114    /// Shadow pass bind group (for rendering to shadow map).
115    pub shadow_bind_group: Option<wgpu::BindGroup>,
116    /// Shadow model uniform buffer (contains only model matrix).
117    pub shadow_model_buffer: Option<wgpu::Buffer>,
118}
119
120impl SurfaceMeshRenderData {
121    /// Creates new render data from mesh geometry.
122    ///
123    /// # Arguments
124    /// * `device` - The wgpu device
125    /// * `bind_group_layout` - The bind group layout for surface meshes
126    /// * `camera_buffer` - The camera uniform buffer
127    /// * `vertices` - Vertex positions
128    /// * `triangles` - Triangle indices (each [u32; 3] is one triangle)
129    /// * `vertex_normals` - Per-vertex normals
130    /// * `edge_is_real` - Per-triangle-vertex flags marking real polygon edges vs internal triangulation edges
131    #[must_use]
132    pub fn new(
133        device: &wgpu::Device,
134        bind_group_layout: &wgpu::BindGroupLayout,
135        camera_buffer: &wgpu::Buffer,
136        vertices: &[Vec3],
137        triangles: &[[u32; 3]],
138        vertex_normals: &[Vec3],
139        edge_is_real: &[Vec3],
140    ) -> Self {
141        let num_triangles = triangles.len() as u32;
142        let num_indices = num_triangles * 3;
143
144        // Expand vertex data to be per-triangle-vertex (not per-original-vertex)
145        // This is needed because the shader accesses data by vertex_index (0, 1, 2, ...)
146        // and we use non-indexed drawing with storage buffers
147        let mut expanded_positions: Vec<f32> = Vec::with_capacity(triangles.len() * 3 * 4);
148        let mut expanded_normals: Vec<f32> = Vec::with_capacity(triangles.len() * 3 * 4);
149        let mut expanded_colors: Vec<f32> = Vec::with_capacity(triangles.len() * 3 * 4);
150        let mut barycentric_data: Vec<f32> = Vec::with_capacity(triangles.len() * 3 * 4);
151        let mut edge_is_real_data: Vec<f32> = Vec::with_capacity(triangles.len() * 3 * 4);
152
153        let mut tri_vertex_idx = 0;
154        for tri in triangles {
155            // Barycentric coordinates for wireframe
156            let bary_coords = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
157
158            for (i, &vi) in tri.iter().enumerate() {
159                let v = vertices[vi as usize];
160                expanded_positions.extend_from_slice(&[v.x, v.y, v.z, 1.0]);
161
162                let n = vertex_normals[vi as usize];
163                expanded_normals.extend_from_slice(&[n.x, n.y, n.z, 0.0]);
164
165                // Default color (will be overwritten by surface_color in shader)
166                expanded_colors.extend_from_slice(&[0.0, 0.0, 0.0, 0.0]);
167
168                barycentric_data.extend_from_slice(&[
169                    bary_coords[i][0],
170                    bary_coords[i][1],
171                    bary_coords[i][2],
172                    0.0,
173                ]);
174
175                // Edge is real flags (same for all vertices of triangle)
176                let eir = edge_is_real[tri_vertex_idx];
177                edge_is_real_data.extend_from_slice(&[eir.x, eir.y, eir.z, 0.0]);
178
179                tri_vertex_idx += 1;
180            }
181        }
182
183        // Create vertex position buffer (vec4 for alignment)
184        let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
185            label: Some("mesh vertices"),
186            contents: bytemuck::cast_slice(&expanded_positions),
187            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
188        });
189
190        // Create index buffer (sequential indices since data is expanded)
191        let index_data: Vec<u32> = (0..num_indices).collect();
192        let index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
193            label: Some("mesh indices"),
194            contents: bytemuck::cast_slice(&index_data),
195            usage: wgpu::BufferUsages::INDEX | wgpu::BufferUsages::COPY_DST,
196        });
197
198        // Create normal buffer (vec4 for alignment)
199        let normal_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
200            label: Some("mesh normals"),
201            contents: bytemuck::cast_slice(&expanded_normals),
202            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
203        });
204
205        // Create barycentric coordinate buffer
206        let barycentric_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
207            label: Some("mesh barycentrics"),
208            contents: bytemuck::cast_slice(&barycentric_data),
209            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
210        });
211
212        // Create color buffer (expanded, default to zero - shader uses surface_color when zero)
213        let color_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
214            label: Some("mesh colors"),
215            contents: bytemuck::cast_slice(&expanded_colors),
216            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
217        });
218
219        // Create edge_is_real buffer
220        let edge_is_real_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
221            label: Some("mesh edge_is_real"),
222            contents: bytemuck::cast_slice(&edge_is_real_data),
223            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
224        });
225
226        // Create uniform buffer
227        let uniforms = MeshUniforms::default();
228        let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
229            label: Some("mesh uniforms"),
230            contents: bytemuck::cast_slice(&[uniforms]),
231            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
232        });
233
234        // Create bind group
235        // Bindings:
236        // 0: camera uniforms (uniform)
237        // 1: mesh uniforms (uniform)
238        // 2: positions (storage)
239        // 3: normals (storage)
240        // 4: barycentrics (storage)
241        // 5: colors (storage)
242        // 6: edge_is_real (storage)
243        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
244            label: Some("surface mesh bind group"),
245            layout: bind_group_layout,
246            entries: &[
247                wgpu::BindGroupEntry {
248                    binding: 0,
249                    resource: camera_buffer.as_entire_binding(),
250                },
251                wgpu::BindGroupEntry {
252                    binding: 1,
253                    resource: uniform_buffer.as_entire_binding(),
254                },
255                wgpu::BindGroupEntry {
256                    binding: 2,
257                    resource: vertex_buffer.as_entire_binding(),
258                },
259                wgpu::BindGroupEntry {
260                    binding: 3,
261                    resource: normal_buffer.as_entire_binding(),
262                },
263                wgpu::BindGroupEntry {
264                    binding: 4,
265                    resource: barycentric_buffer.as_entire_binding(),
266                },
267                wgpu::BindGroupEntry {
268                    binding: 5,
269                    resource: color_buffer.as_entire_binding(),
270                },
271                wgpu::BindGroupEntry {
272                    binding: 6,
273                    resource: edge_is_real_buffer.as_entire_binding(),
274                },
275            ],
276        });
277
278        Self {
279            vertex_buffer,
280            index_buffer,
281            normal_buffer,
282            barycentric_buffer,
283            color_buffer,
284            edge_is_real_buffer,
285            uniform_buffer,
286            bind_group,
287            num_triangles,
288            num_indices,
289            shadow_bind_group: None,
290            shadow_model_buffer: None,
291        }
292    }
293
294    /// Updates the mesh uniform buffer.
295    pub fn update_uniforms(&self, queue: &wgpu::Queue, uniforms: &MeshUniforms) {
296        queue.write_buffer(&self.uniform_buffer, 0, bytemuck::cast_slice(&[*uniforms]));
297    }
298
299    /// Updates the per-vertex color buffer with per-original-vertex colors.
300    /// This expands the colors to match the per-triangle-vertex buffer layout.
301    pub fn update_colors(&self, queue: &wgpu::Queue, colors: &[Vec4], triangles: &[[u32; 3]]) {
302        // Expand per-vertex colors to per-triangle-vertex
303        let mut expanded_colors: Vec<f32> = Vec::with_capacity(triangles.len() * 3 * 4);
304        for tri in triangles {
305            for &vi in tri {
306                let c = colors[vi as usize];
307                expanded_colors.extend_from_slice(&[c.x, c.y, c.z, c.w]);
308            }
309        }
310        queue.write_buffer(
311            &self.color_buffer,
312            0,
313            bytemuck::cast_slice(&expanded_colors),
314        );
315    }
316
317    /// Clears the color buffer (sets all colors to zero, which means use `surface_color`).
318    pub fn clear_colors(&self, queue: &wgpu::Queue) {
319        let zero_colors: Vec<f32> = vec![0.0; self.num_indices as usize * 4];
320        queue.write_buffer(&self.color_buffer, 0, bytemuck::cast_slice(&zero_colors));
321    }
322
323    /// Initializes shadow rendering resources.
324    ///
325    /// Creates the bind group and model buffer needed for the shadow pass.
326    /// The shadow pass renders this mesh from the light's perspective to create
327    /// shadows on the ground plane.
328    pub fn init_shadow_resources(
329        &mut self,
330        device: &wgpu::Device,
331        shadow_bind_group_layout: &wgpu::BindGroupLayout,
332        light_buffer: &wgpu::Buffer,
333    ) {
334        // Create model uniform buffer for shadow pass
335        let shadow_model_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
336            label: Some("Shadow Model Uniform Buffer"),
337            contents: bytemuck::cast_slice(&[ShadowModelUniforms::default()]),
338            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
339        });
340
341        // Create bind group matching shadow_map.wgsl:
342        // binding 0: light uniforms (view_proj, light_dir)
343        // binding 1: model uniforms (model matrix)
344        // binding 2: vertex positions (storage buffer)
345        let shadow_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
346            label: Some("Surface Mesh Shadow Bind Group"),
347            layout: shadow_bind_group_layout,
348            entries: &[
349                wgpu::BindGroupEntry {
350                    binding: 0,
351                    resource: light_buffer.as_entire_binding(),
352                },
353                wgpu::BindGroupEntry {
354                    binding: 1,
355                    resource: shadow_model_buffer.as_entire_binding(),
356                },
357                wgpu::BindGroupEntry {
358                    binding: 2,
359                    resource: self.vertex_buffer.as_entire_binding(),
360                },
361            ],
362        });
363
364        self.shadow_model_buffer = Some(shadow_model_buffer);
365        self.shadow_bind_group = Some(shadow_bind_group);
366    }
367
368    /// Returns whether shadow resources have been initialized.
369    #[must_use]
370    pub fn has_shadow_resources(&self) -> bool {
371        self.shadow_bind_group.is_some()
372    }
373
374    /// Returns the number of vertices for shadow pass rendering.
375    ///
376    /// The shadow pass uses non-indexed drawing with storage buffers,
377    /// so this returns the number of expanded triangle vertices.
378    #[must_use]
379    pub fn num_vertices(&self) -> u32 {
380        self.num_indices // Same as expanded triangle vertex count
381    }
382
383    /// Updates the shadow model uniform buffer with the current transform.
384    ///
385    /// This should be called whenever the mesh's transform changes to ensure
386    /// shadows are rendered at the correct position.
387    pub fn update_shadow_model(&self, queue: &wgpu::Queue, model_matrix: [[f32; 4]; 4]) {
388        if let Some(buffer) = &self.shadow_model_buffer {
389            let uniforms = ShadowModelUniforms {
390                model: model_matrix,
391            };
392            queue.write_buffer(buffer, 0, bytemuck::cast_slice(&[uniforms]));
393        }
394    }
395
396    // Buffer getters for reflection rendering
397
398    /// Returns the uniform buffer.
399    #[must_use]
400    pub fn uniform_buffer(&self) -> &wgpu::Buffer {
401        &self.uniform_buffer
402    }
403
404    /// Returns the position buffer.
405    #[must_use]
406    pub fn position_buffer(&self) -> &wgpu::Buffer {
407        &self.vertex_buffer
408    }
409
410    /// Returns the normal buffer.
411    #[must_use]
412    pub fn normal_buffer(&self) -> &wgpu::Buffer {
413        &self.normal_buffer
414    }
415
416    /// Returns the barycentric buffer.
417    #[must_use]
418    pub fn barycentric_buffer(&self) -> &wgpu::Buffer {
419        &self.barycentric_buffer
420    }
421
422    /// Returns the color buffer.
423    #[must_use]
424    pub fn color_buffer(&self) -> &wgpu::Buffer {
425        &self.color_buffer
426    }
427
428    /// Returns the edge is real buffer.
429    #[must_use]
430    pub fn edge_is_real_buffer(&self) -> &wgpu::Buffer {
431        &self.edge_is_real_buffer
432    }
433
434    /// Returns the number of vertices (for non-indexed drawing).
435    #[must_use]
436    pub fn vertex_count(&self) -> u32 {
437        self.num_indices
438    }
439}
440
441#[cfg(test)]
442mod tests {
443    use super::*;
444
445    #[test]
446    fn test_mesh_uniforms_size() {
447        let size = std::mem::size_of::<MeshUniforms>();
448
449        // Verify size is 16-byte aligned for GPU uniform buffers
450        assert_eq!(
451            size % 16,
452            0,
453            "MeshUniforms size ({} bytes) must be 16-byte aligned",
454            size
455        );
456
457        // Expected size breakdown:
458        // model_matrix: 64 bytes ([[f32; 4]; 4])
459        // shade_style: 4 bytes (u32)
460        // show_edges: 4 bytes (u32)
461        // edge_width: 4 bytes (f32)
462        // transparency: 4 bytes (f32)
463        // surface_color: 16 bytes ([f32; 4])
464        // edge_color: 16 bytes ([f32; 4])
465        // backface_policy: 4 bytes (u32)
466        // slice_planes_enabled: 4 bytes (u32)
467        // use_vertex_color: 4 bytes (u32)
468        // _pad1: 4 bytes (f32)
469        // _pad2: 12 bytes ([f32; 3])
470        // _pad3: 4 bytes (f32)
471        // backface_color: 16 bytes ([f32; 4])
472        // Total: 160 bytes (matches WGSL layout with vec3 alignment)
473        assert_eq!(
474            size, 160,
475            "MeshUniforms should be 160 bytes, got {} bytes",
476            size
477        );
478    }
479}