1use 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#[repr(C)]
26#[derive(Copy, Clone, Debug, Pod, Zeroable)]
27pub struct LineRasterUniforms {
28 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 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 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}