Skip to main content

loam_render/
triangle_raster.rs

1use bytemuck::{Pod, Zeroable};
2use glam::Mat4;
3use loam_math::{Projection, RasterizableSpace};
4use loam_shape::TriangleMesh;
5use wgpu::{
6    BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayoutEntry, BindingType,
7    BlendComponent, BlendFactor, BlendOperation, BlendState, Buffer, BufferBindingType,
8    BufferDescriptor, BufferUsages, CompareFunction, Device, Features, Limits, LoadOp, Operations,
9    PrimitiveTopology, Queue, RenderPassColorAttachment, RenderPassDepthStencilAttachment,
10    RenderPassDescriptor, RenderPipeline, ShaderStages, StoreOp, TextureFormat, VertexAttribute,
11    VertexBufferLayout, VertexFormat, VertexStepMode,
12};
13
14use crate::device::{FeatureRequest, GpuContext, MissingGpuCapability};
15use crate::material::MaterialSpec;
16
17const TRIANGLE_RASTER_WGSL: &str = include_str!("triangle_raster.wgsl");
18
19const CAMERA_BINDING: [BindGroupLayoutEntry; 1] = [BindGroupLayoutEntry {
20    binding: 0,
21    visibility: ShaderStages::VERTEX_FRAGMENT,
22    ty: BindingType::Buffer {
23        ty: BufferBindingType::Uniform,
24        has_dynamic_offset: false,
25        min_binding_size: None,
26    },
27    count: None,
28}];
29
30/// Matches the WGSL `CameraUniform`.
31#[repr(C)]
32#[derive(Copy, Clone, Debug, Pod, Zeroable)]
33pub struct TriangleRasterUniforms {
34    pub view_projection: [[f32; 4]; 4],
35}
36
37impl Default for TriangleRasterUniforms {
38    fn default() -> Self {
39        Self {
40            view_projection: Mat4::IDENTITY.to_cols_array_2d(),
41        }
42    }
43}
44
45/// Matches the vertex attribute layout: position at 0, color at 16.
46#[repr(C)]
47#[derive(Copy, Clone, Debug, Pod, Zeroable, Default)]
48pub struct TriangleVertex {
49    pub position: [f32; 3],
50    pub _pad0: f32,
51    pub color: [f32; 4],
52}
53
54#[derive(Copy, Clone, Debug, PartialEq, Eq, Default)]
55pub enum FragmentShading {
56    #[default]
57    Flat,
58    FaceNormalLambert,
59}
60
61impl FragmentShading {
62    fn entry_point(self) -> &'static str {
63        match self {
64            Self::Flat => "fs_flat",
65            Self::FaceNormalLambert => "fs_lambert",
66        }
67    }
68}
69
70pub struct TriangleRasterNode {
71    pipeline: RenderPipeline,
72    uniform_buf: Buffer,
73    bind_group: BindGroup,
74
75    vertex_buf: Buffer,
76    vertex_capacity: u32,
77
78    index_buf: Buffer,
79    index_capacity: u32,
80    index_count: u32,
81
82    has_depth: bool,
83    vertices_scratch: Vec<TriangleVertex>,
84}
85
86impl TriangleRasterNode {
87    pub fn new(
88        gpu: &GpuContext,
89        surface_format: TextureFormat,
90        depth: crate::DepthMode,
91        convention: crate::DepthConvention,
92        shading: FragmentShading,
93    ) -> Result<Self, MissingGpuCapability> {
94        let device = &gpu.device;
95        let vertex_attrs = [
96            VertexAttribute {
97                format: VertexFormat::Float32x3,
98                offset: 0,
99                shader_location: 0,
100            },
101            VertexAttribute {
102                format: VertexFormat::Float32x4,
103                offset: 16,
104                shader_location: 1,
105            },
106        ];
107        let material = MaterialSpec {
108            label: "triangle_raster",
109            shader: TRIANGLE_RASTER_WGSL,
110            vertex_entry: "vs_main",
111            fragment_entry: shading.entry_point(),
112            vertex_layouts: &[VertexBufferLayout {
113                array_stride: std::mem::size_of::<TriangleVertex>() as u64,
114                step_mode: VertexStepMode::Vertex,
115                attributes: &vertex_attrs,
116            }],
117            resources: &[&CAMERA_BINDING],
118            features: FeatureRequest {
119                required_features: Features::empty(),
120                optional_features: Features::empty(),
121                required_limits: Limits::downlevel_webgl2_defaults(),
122            },
123            topology: PrimitiveTopology::TriangleList,
124            blend: Some(BlendState {
125                color: BlendComponent {
126                    src_factor: BlendFactor::SrcAlpha,
127                    dst_factor: BlendFactor::OneMinusSrcAlpha,
128                    operation: BlendOperation::Add,
129                },
130                alpha: BlendComponent {
131                    src_factor: BlendFactor::One,
132                    dst_factor: BlendFactor::OneMinusSrcAlpha,
133                    operation: BlendOperation::Add,
134                },
135            }),
136            depth,
137            convention,
138            standard_z_compare: CompareFunction::Less,
139        };
140        let built = material.build(gpu, surface_format)?;
141
142        let uniform_buf = device.create_buffer(&BufferDescriptor {
143            label: Some("triangle_raster uniforms"),
144            size: std::mem::size_of::<TriangleRasterUniforms>() as u64,
145            usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
146            mapped_at_creation: false,
147        });
148        let bind_group = device.create_bind_group(&BindGroupDescriptor {
149            label: Some("triangle_raster bg"),
150            layout: &built.pipeline.get_bind_group_layout(0),
151            entries: &[BindGroupEntry {
152                binding: 0,
153                resource: uniform_buf.as_entire_binding(),
154            }],
155        });
156
157        let vertex_buf = device.create_buffer(&BufferDescriptor {
158            label: Some("triangle_raster vertex buffer"),
159            size: 64,
160            usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
161            mapped_at_creation: false,
162        });
163        let index_buf = device.create_buffer(&BufferDescriptor {
164            label: Some("triangle_raster index buffer"),
165            size: 64,
166            usage: BufferUsages::INDEX | BufferUsages::COPY_DST,
167            mapped_at_creation: false,
168        });
169
170        Ok(Self {
171            pipeline: built.pipeline,
172            uniform_buf,
173            bind_group,
174            vertex_buf,
175            vertex_capacity: 0,
176            index_buf,
177            index_capacity: 0,
178            index_count: 0,
179            has_depth: depth.is_active(),
180            vertices_scratch: Vec::new(),
181        })
182    }
183
184    pub fn set_camera(&self, queue: &Queue, view_projection: Mat4) {
185        let uniforms = TriangleRasterUniforms {
186            view_projection: view_projection.to_cols_array_2d(),
187        };
188        queue.write_buffer(&self.uniform_buf, 0, bytemuck::bytes_of(&uniforms));
189    }
190
191    pub fn upload<S, const N: usize>(
192        &mut self,
193        device: &Device,
194        queue: &Queue,
195        mesh: &TriangleMesh<N>,
196        projection: &Projection<N>,
197    ) where
198        S: RasterizableSpace<N>,
199    {
200        let n_vertices = mesh.vertices.len();
201        assert_eq!(
202            mesh.colors.len(),
203            n_vertices,
204            "TriangleMesh invariant: colors.len() == vertices.len()"
205        );
206        self.vertices_scratch.clear();
207        let verts = &mut self.vertices_scratch;
208        for (v, color) in mesh.vertices.iter().zip(mesh.colors.iter()) {
209            let p_native = S::array_to_point(*v);
210            let p3 = S::project_point(p_native, projection).unwrap_or(glam::Vec3::NAN);
211            verts.push(TriangleVertex {
212                position: p3.to_array(),
213                _pad0: 0.0,
214                color: *color,
215            });
216        }
217
218        let indices: &[u32] = bytemuck::cast_slice(&mesh.indices);
219
220        if verts.len() as u32 > self.vertex_capacity {
221            let new_cap = (verts.len() as u32).next_power_of_two().max(16);
222            self.vertex_buf = device.create_buffer(&BufferDescriptor {
223                label: Some("triangle_raster vertex buffer"),
224                size: (new_cap as u64) * (std::mem::size_of::<TriangleVertex>() as u64),
225                usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
226                mapped_at_creation: false,
227            });
228            self.vertex_capacity = new_cap;
229        }
230        if indices.len() as u32 > self.index_capacity {
231            let new_cap = (indices.len() as u32).next_power_of_two().max(16);
232            self.index_buf = device.create_buffer(&BufferDescriptor {
233                label: Some("triangle_raster index buffer"),
234                size: (new_cap as u64) * (std::mem::size_of::<u32>() as u64),
235                usage: BufferUsages::INDEX | BufferUsages::COPY_DST,
236                mapped_at_creation: false,
237            });
238            self.index_capacity = new_cap;
239        }
240
241        if !verts.is_empty() {
242            queue.write_buffer(&self.vertex_buf, 0, bytemuck::cast_slice(verts));
243        }
244        if !indices.is_empty() {
245            queue.write_buffer(&self.index_buf, 0, bytemuck::cast_slice(indices));
246        }
247        self.index_count = indices.len() as u32;
248    }
249
250    /// Loads both attachments; uploads between records affect every draw in the same submission.
251    pub fn record(
252        &self,
253        encoder: &mut wgpu::CommandEncoder,
254        view: &wgpu::TextureView,
255        depth_view: Option<&wgpu::TextureView>,
256        viewport: Option<&crate::Viewport>,
257    ) {
258        match (self.has_depth, depth_view.is_some()) {
259            (true, false) => {
260                panic!(
261                    "TriangleRasterNode::record: pipeline was created with a depth format but \
262                     no depth view was provided"
263                )
264            }
265            (false, true) => {
266                panic!(
267                    "TriangleRasterNode::record: pipeline was created without a depth format \
268                     but a depth view was provided"
269                )
270            }
271            _ => {}
272        }
273        if self.index_count == 0 {
274            return;
275        }
276        let depth_attachment = depth_view.map(|dv| RenderPassDepthStencilAttachment {
277            view: dv,
278            depth_ops: Some(Operations {
279                load: LoadOp::Load,
280                store: StoreOp::Store,
281            }),
282            stencil_ops: None,
283        });
284        let mut rp = encoder.begin_render_pass(&RenderPassDescriptor {
285            label: Some("triangle_raster pass"),
286            color_attachments: &[Some(RenderPassColorAttachment {
287                view,
288                depth_slice: None,
289                resolve_target: None,
290                ops: Operations {
291                    load: LoadOp::Load,
292                    store: StoreOp::Store,
293                },
294            })],
295            depth_stencil_attachment: depth_attachment,
296            timestamp_writes: None,
297            occlusion_query_set: None,
298        });
299        self.draw(&mut rp, viewport);
300    }
301
302    pub fn draw(&self, rp: &mut wgpu::RenderPass<'_>, viewport: Option<&crate::Viewport>) {
303        if self.index_count == 0 {
304            return;
305        }
306        if let Some(vp) = viewport {
307            vp.apply(rp);
308        }
309        rp.set_pipeline(&self.pipeline);
310        rp.set_bind_group(0, &self.bind_group, &[]);
311        rp.set_vertex_buffer(0, self.vertex_buf.slice(..));
312        rp.set_index_buffer(self.index_buf.slice(..), wgpu::IndexFormat::Uint32);
313        rp.draw_indexed(0..self.index_count, 0, 0..1);
314    }
315}