Skip to main content

loam_render/
line_raster.rs

1//! The caller owns the depth texture and clears it before this node's draw,
2//! which loads both color and depth.
3
4use bytemuck::{Pod, Zeroable};
5use glam::{Mat4, Vec2};
6use loam_math::{Projection, RasterizableSpace};
7use loam_runtime::{Eye, SegmentRecord};
8use loam_shape::LineMesh;
9use wgpu::util::DeviceExt;
10use wgpu::{
11    BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor,
12    BindGroupLayoutEntry, BindingType, BlendComponent, BlendFactor, BlendOperation, BlendState,
13    Buffer, BufferBindingType, BufferDescriptor, BufferUsages, ColorTargetState, ColorWrites,
14    CompareFunction, DepthStencilState, Device, FragmentState, LoadOp, MultisampleState,
15    Operations, PipelineLayoutDescriptor, PrimitiveState, PrimitiveTopology, Queue,
16    RenderPassColorAttachment, RenderPassDepthStencilAttachment, RenderPassDescriptor,
17    RenderPipeline, RenderPipelineDescriptor, ShaderModuleDescriptor, ShaderSource, ShaderStages,
18    StencilState, StoreOp, TextureFormat, VertexAttribute, VertexBufferLayout, VertexFormat,
19    VertexState, VertexStepMode,
20};
21
22const LINE_RASTER_WGSL: &str = include_str!("line_raster.wgsl");
23
24/// Matches the WGSL `CameraUniform`.
25#[repr(C)]
26#[derive(Copy, Clone, Debug, Pod, Zeroable)]
27pub struct LineRasterUniforms {
28    /// Must be bit-identical to the raymarcher's for the same frame.
29    pub view_projection: [[f32; 4]; 4],
30    pub viewport_size: [f32; 2],
31    pub _pad: [f32; 2],
32}
33
34impl Default for LineRasterUniforms {
35    fn default() -> Self {
36        Self {
37            view_projection: Mat4::IDENTITY.to_cols_array_2d(),
38            viewport_size: [1.0, 1.0],
39            _pad: [0.0; 2],
40        }
41    }
42}
43
44pub struct LineRasterNode {
45    pipeline: RenderPipeline,
46    uniform_buf: Buffer,
47    bind_group: BindGroup,
48
49    corner_buf: Buffer,
50    index_buf: Buffer,
51    instance_buf: Buffer,
52    instance_count: u32,
53    instance_capacity: u32,
54    has_depth: bool,
55    instances_scratch: Vec<SegmentRecord>,
56}
57
58impl LineRasterNode {
59    pub fn new(
60        device: &Device,
61        surface_format: TextureFormat,
62        depth: crate::DepthMode,
63        convention: crate::DepthConvention,
64    ) -> Self {
65        let module = device.create_shader_module(ShaderModuleDescriptor {
66            label: Some("line_raster shader"),
67            source: ShaderSource::Wgsl(LINE_RASTER_WGSL.into()),
68        });
69
70        let uniform_buf = device.create_buffer(&BufferDescriptor {
71            label: Some("line_raster uniforms"),
72            size: std::mem::size_of::<LineRasterUniforms>() as u64,
73            usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
74            mapped_at_creation: false,
75        });
76
77        let bgl = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
78            label: Some("line_raster bgl"),
79            entries: &[BindGroupLayoutEntry {
80                binding: 0,
81                visibility: ShaderStages::VERTEX_FRAGMENT,
82                ty: BindingType::Buffer {
83                    ty: BufferBindingType::Uniform,
84                    has_dynamic_offset: false,
85                    min_binding_size: None,
86                },
87                count: None,
88            }],
89        });
90
91        let bind_group = device.create_bind_group(&BindGroupDescriptor {
92            label: Some("line_raster bg"),
93            layout: &bgl,
94            entries: &[BindGroupEntry {
95                binding: 0,
96                resource: uniform_buf.as_entire_binding(),
97            }],
98        });
99
100        let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
101            label: Some("line_raster pipeline layout"),
102            bind_group_layouts: &[&bgl],
103            push_constant_ranges: &[],
104        });
105
106        let corner_attrs = [VertexAttribute {
107            format: VertexFormat::Uint32,
108            offset: 0,
109            shader_location: 0,
110        }];
111        let corner_layout = VertexBufferLayout {
112            array_stride: std::mem::size_of::<u32>() as u64,
113            step_mode: VertexStepMode::Vertex,
114            attributes: &corner_attrs,
115        };
116        let instance_attrs = [
117            VertexAttribute {
118                format: VertexFormat::Float32x3,
119                offset: 0,
120                shader_location: 1,
121            },
122            VertexAttribute {
123                format: VertexFormat::Float32x3,
124                offset: 16,
125                shader_location: 2,
126            },
127            VertexAttribute {
128                format: VertexFormat::Float32x4,
129                offset: 32,
130                shader_location: 3,
131            },
132            VertexAttribute {
133                format: VertexFormat::Float32x4,
134                offset: 48,
135                shader_location: 4,
136            },
137            VertexAttribute {
138                format: VertexFormat::Float32,
139                offset: 64,
140                shader_location: 5,
141            },
142        ];
143        let instance_layout = VertexBufferLayout {
144            array_stride: std::mem::size_of::<SegmentRecord>() as u64,
145            step_mode: VertexStepMode::Instance,
146            attributes: &instance_attrs,
147        };
148
149        let pipeline = device.create_render_pipeline(&RenderPipelineDescriptor {
150            label: Some("line_raster pipeline"),
151            layout: Some(&pipeline_layout),
152            vertex: VertexState {
153                module: &module,
154                entry_point: Some("vs_main"),
155                buffers: &[corner_layout, instance_layout],
156                compilation_options: Default::default(),
157            },
158            fragment: Some(FragmentState {
159                module: &module,
160                entry_point: Some("fs_main"),
161                targets: &[Some(ColorTargetState {
162                    format: surface_format,
163                    blend: Some(BlendState {
164                        color: BlendComponent {
165                            src_factor: BlendFactor::SrcAlpha,
166                            dst_factor: BlendFactor::OneMinusSrcAlpha,
167                            operation: BlendOperation::Add,
168                        },
169                        alpha: BlendComponent {
170                            src_factor: BlendFactor::One,
171                            dst_factor: BlendFactor::OneMinusSrcAlpha,
172                            operation: BlendOperation::Add,
173                        },
174                    }),
175                    write_mask: ColorWrites::ALL,
176                })],
177                compilation_options: Default::default(),
178            }),
179            primitive: PrimitiveState {
180                topology: PrimitiveTopology::TriangleList,
181                ..Default::default()
182            },
183            depth_stencil: depth_state(depth, convention),
184            multisample: MultisampleState {
185                count: 1,
186                ..Default::default()
187            },
188            multiview: None,
189            cache: None,
190        });
191
192        let corner_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
193            label: Some("line_raster corner buffer"),
194            contents: bytemuck::cast_slice(&[0u32, 1, 2, 3]),
195            usage: BufferUsages::VERTEX,
196        });
197
198        let index_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
199            label: Some("line_raster index buffer"),
200            contents: bytemuck::cast_slice(&[0u32, 1, 2, 2, 1, 3]),
201            usage: BufferUsages::INDEX,
202        });
203
204        let instance_capacity = 0u32;
205        let instance_buf = device.create_buffer(&BufferDescriptor {
206            label: Some("line_raster instance buffer"),
207            size: 64,
208            usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
209            mapped_at_creation: false,
210        });
211
212        Self {
213            pipeline,
214            uniform_buf,
215            bind_group,
216            corner_buf,
217            index_buf,
218            instance_buf,
219            instance_count: 0,
220            instance_capacity,
221            has_depth: depth.is_active(),
222            instances_scratch: Vec::new(),
223        }
224    }
225
226    pub fn set_root_camera(&self, queue: &Queue, eye: &Eye, viewport_size: Vec2) {
227        self.set_camera(queue, crate::view::root_view_projection(eye), viewport_size);
228    }
229
230    pub fn set_camera(&self, queue: &Queue, view_projection: Mat4, viewport_size: Vec2) {
231        let uniforms = LineRasterUniforms {
232            view_projection: view_projection.to_cols_array_2d(),
233            viewport_size: viewport_size.to_array(),
234            _pad: [0.0; 2],
235        };
236        queue.write_buffer(&self.uniform_buf, 0, bytemuck::bytes_of(&uniforms));
237    }
238
239    pub fn upload<S, const N: usize>(
240        &mut self,
241        device: &Device,
242        queue: &Queue,
243        mesh: &LineMesh<N>,
244        projection: &Projection<N>,
245        samples_per_segment: usize,
246    ) where
247        S: RasterizableSpace<N>,
248    {
249        let samples = samples_per_segment.max(1);
250        build_line_instances::<S, N>(&mut self.instances_scratch, mesh, projection, samples);
251
252        let scratch = std::mem::take(&mut self.instances_scratch);
253        self.upload_segments(device, queue, &scratch);
254        self.instances_scratch = scratch;
255    }
256
257    pub fn segment_count(&self) -> u32 {
258        self.instance_count
259    }
260
261    pub fn upload_segments(&mut self, device: &Device, queue: &Queue, segments: &[SegmentRecord]) {
262        let needed_capacity = segments.len() as u32;
263        if needed_capacity > self.instance_capacity {
264            let new_cap = needed_capacity.next_power_of_two().max(16);
265            self.instance_buf = device.create_buffer(&BufferDescriptor {
266                label: Some("line_raster instance buffer"),
267                size: (new_cap as u64) * (std::mem::size_of::<SegmentRecord>() as u64),
268                usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
269                mapped_at_creation: false,
270            });
271            self.instance_capacity = new_cap;
272        }
273        if !segments.is_empty() {
274            queue.write_buffer(&self.instance_buf, 0, bytemuck::cast_slice(segments));
275        }
276        self.instance_count = needed_capacity;
277    }
278
279    /// Records into the caller's encoder with `LoadOp::Load` on both attachments; `depth_view` is `Some` iff the pipeline has depth.
280    pub fn record(
281        &self,
282        encoder: &mut wgpu::CommandEncoder,
283        view: &wgpu::TextureView,
284        depth_view: Option<&wgpu::TextureView>,
285        viewport: Option<&crate::Viewport>,
286    ) {
287        match (self.has_depth, depth_view.is_some()) {
288            (true, false) => {
289                panic!(
290                    "LineRasterNode::record: pipeline was created with a depth format but \
291                     no depth view was provided"
292                )
293            }
294            (false, true) => {
295                panic!(
296                    "LineRasterNode::record: pipeline was created without a depth format but \
297                     a depth view was provided"
298                )
299            }
300            _ => {}
301        }
302        if self.instance_count == 0 {
303            return;
304        }
305        let depth_attachment = depth_view.map(|dv| RenderPassDepthStencilAttachment {
306            view: dv,
307            depth_ops: Some(Operations {
308                load: LoadOp::Load,
309                store: StoreOp::Store,
310            }),
311            stencil_ops: None,
312        });
313        let mut rp = encoder.begin_render_pass(&RenderPassDescriptor {
314            label: Some("line_raster pass"),
315            color_attachments: &[Some(RenderPassColorAttachment {
316                view,
317                depth_slice: None,
318                resolve_target: None,
319                ops: Operations {
320                    load: LoadOp::Load,
321                    store: StoreOp::Store,
322                },
323            })],
324            depth_stencil_attachment: depth_attachment,
325            timestamp_writes: None,
326            occlusion_query_set: None,
327        });
328        self.draw(&mut rp, viewport);
329    }
330
331    pub fn draw(&self, rp: &mut wgpu::RenderPass<'_>, viewport: Option<&crate::Viewport>) {
332        if self.instance_count == 0 {
333            return;
334        }
335        if let Some(vp) = viewport {
336            vp.apply(rp);
337        }
338        rp.set_pipeline(&self.pipeline);
339        rp.set_bind_group(0, &self.bind_group, &[]);
340        rp.set_vertex_buffer(0, self.corner_buf.slice(..));
341        rp.set_vertex_buffer(1, self.instance_buf.slice(..));
342        rp.set_index_buffer(self.index_buf.slice(..), wgpu::IndexFormat::Uint32);
343        rp.draw_indexed(0..6, 0, 0..self.instance_count);
344    }
345}
346
347fn depth_state(
348    depth: crate::DepthMode,
349    convention: crate::DepthConvention,
350) -> Option<DepthStencilState> {
351    depth.format().map(|format| DepthStencilState {
352        format,
353        depth_write_enabled: depth.writes(),
354        // Equal depth passes so coplanar outlines survive the filled surface.
355        depth_compare: convention.compare(CompareFunction::LessEqual),
356        stencil: StencilState::default(),
357        bias: wgpu::DepthBiasState::default(),
358    })
359}
360
361fn build_line_instances<S, const N: usize>(
362    out: &mut Vec<SegmentRecord>,
363    mesh: &LineMesh<N>,
364    projection: &Projection<N>,
365    samples: usize,
366) where
367    S: RasterizableSpace<N>,
368{
369    out.clear();
370    out.reserve(mesh.segments.len() * samples);
371
372    for ((seg, (color_a, color_b)), &width) in mesh
373        .segments
374        .iter()
375        .zip(mesh.colors.iter())
376        .zip(mesh.widths.iter())
377    {
378        let p0 = S::array_to_point(seg.0);
379        let p1 = S::array_to_point(seg.1);
380        let mut previous = None;
381        let mut index = 0;
382        S::tessellate_segment(p0, p1, samples, |point| {
383            let q1 = S::project_point(point, projection).unwrap_or(glam::Vec3::NAN);
384            if let Some(q0) = previous {
385                let t0 = index as f32 / samples as f32;
386                let t1 = (index + 1) as f32 / samples as f32;
387                if glam::Vec3::is_finite(q0) && q1.is_finite() {
388                    out.push(SegmentRecord {
389                        start: q0.to_array(),
390                        _pad0: 0.0,
391                        end: q1.to_array(),
392                        _pad1: 0.0,
393                        start_color: lerp_color(*color_a, *color_b, t0),
394                        end_color: lerp_color(*color_a, *color_b, t1),
395                        width_px: width,
396                        _pad2: [0.0; 3],
397                    });
398                }
399                index += 1;
400            }
401            previous = Some(q1);
402        });
403    }
404}
405
406fn lerp_color(a: [f32; 4], b: [f32; 4], t: f32) -> [f32; 4] {
407    [
408        a[0] + (b[0] - a[0]) * t,
409        a[1] + (b[1] - a[1]) * t,
410        a[2] + (b[2] - a[2]) * t,
411        a[3] + (b[3] - a[3]) * t,
412    ]
413}
414
415#[cfg(test)]
416mod tests {
417    use super::*;
418
419    use glam::Vec3;
420    use loam_math::{EuclideanR3, Iso3};
421    use loam_shape::LineMesh;
422
423    use crate::depth_passes as passes;
424    use crate::view::{eye_relative, projective_depth, root_view_projection, DEPTH_FORMAT};
425    use crate::{DepthConvention, DepthMode};
426
427    #[test]
428    fn line_raster_hides_the_nearer_line_under_the_convention_its_pipeline_declares() {
429        let eye = Eye {
430            position: [0.0, 0.0, 3.0],
431            ..Eye::default()
432        };
433        let standing = Iso3 {
434            rotation: glam::Quat::IDENTITY,
435            translation: Vec3::from(eye.position),
436        };
437        let near_point = Vec3::new(0.2, -0.1, 0.0);
438        let far_point = Vec3::new(0.2, -0.1, -4.0);
439
440        for (convention, projection) in [
441            (DepthConvention::ReversedZ, root_view_projection(&eye)),
442            (
443                DepthConvention::StandardZ,
444                Mat4::perspective_rh(eye.fov_y, eye.aspect, eye.near, eye.far)
445                    * eye_relative(standing),
446            ),
447        ] {
448            let depth = |point: Vec3| {
449                let clip = projection * point.extend(1.0);
450                clip.z / clip.w
451            };
452            let state = depth_state(
453                DepthMode::ReadWrite {
454                    format: DEPTH_FORMAT,
455                },
456                convention,
457            )
458            .unwrap();
459            assert!(
460                passes(state.depth_compare, depth(near_point), depth(far_point)),
461                "{convention:?} with {:?} hides the nearer line",
462                state.depth_compare
463            );
464            assert!(!passes(
465                state.depth_compare,
466                depth(far_point),
467                depth(near_point)
468            ));
469        }
470
471        let reversed = root_view_projection(&eye);
472        for point in [near_point, far_point] {
473            let clip = reversed * point.extend(1.0);
474            let image = projective_depth(point - Vec3::from(eye.position), eye.near);
475            assert!(
476                (clip.z / clip.w - image).abs() <= 1e-6,
477                "{point} rasterizes at depth {}, not the image-space {image}",
478                clip.z / clip.w
479            );
480        }
481    }
482
483    fn one_segment_mesh(a: [f32; 3], b: [f32; 3]) -> LineMesh<3> {
484        LineMesh {
485            segments: vec![(a, b)],
486            colors: vec![([1.0; 4], [1.0; 4])],
487            widths: vec![1.0],
488        }
489    }
490
491    #[test]
492    fn upload_drops_non_finite_segments() {
493        let mut scratch: Vec<SegmentRecord> = Vec::new();
494
495        let finite = one_segment_mesh([0.0, 0.0, 0.0], [1.0, 1.0, 1.0]);
496        build_line_instances::<EuclideanR3, 3>(&mut scratch, &finite, &Projection::Identity, 1);
497        assert_eq!(scratch.len(), 1, "finite segment must produce one instance");
498
499        let nan = one_segment_mesh([0.0, 0.0, 0.0], [f32::NAN, 1.0, 1.0]);
500        build_line_instances::<EuclideanR3, 3>(&mut scratch, &nan, &Projection::Identity, 1);
501        assert_eq!(scratch.len(), 0, "NaN-endpoint segment must be dropped");
502    }
503
504    #[test]
505    fn upload_drops_non_finite_without_reallocating() {
506        let mut scratch: Vec<SegmentRecord> = Vec::new();
507
508        let mut many = LineMesh::<3>::default();
509        for k in 0..8 {
510            let f = k as f32;
511            many.segments.push(([f, 0.0, 0.0], [f, 1.0, 0.0]));
512            many.colors.push(([1.0; 4], [1.0; 4]));
513            many.widths.push(1.0);
514        }
515        build_line_instances::<EuclideanR3, 3>(&mut scratch, &many, &Projection::Identity, 1);
516        assert_eq!(scratch.len(), 8);
517        let cap_after_growth = scratch.capacity();
518
519        let nan = one_segment_mesh([f32::NAN, 0.0, 0.0], [0.0, 0.0, 0.0]);
520        build_line_instances::<EuclideanR3, 3>(&mut scratch, &nan, &Projection::Identity, 1);
521        assert_eq!(scratch.len(), 0);
522        assert_eq!(
523            scratch.capacity(),
524            cap_after_growth,
525            "dropping a segment must not change the scratch heap capacity"
526        );
527    }
528
529    #[test]
530    fn upload_drop_predicate_catches_infinity_not_just_nan() {
531        let mut scratch: Vec<SegmentRecord> = Vec::new();
532
533        let pos_inf = one_segment_mesh([0.0, 0.0, 0.0], [f32::INFINITY, 1.0, 1.0]);
534        build_line_instances::<EuclideanR3, 3>(&mut scratch, &pos_inf, &Projection::Identity, 1);
535        assert_eq!(scratch.len(), 0, "+inf endpoint must be dropped");
536
537        let neg_inf = one_segment_mesh([f32::NEG_INFINITY, 0.0, 0.0], [1.0, 1.0, 1.0]);
538        build_line_instances::<EuclideanR3, 3>(&mut scratch, &neg_inf, &Projection::Identity, 1);
539        assert_eq!(scratch.len(), 0, "-inf endpoint must be dropped");
540    }
541}