1use bytemuck::{Pod, Zeroable};
2use glam::{Mat4, Vec2};
3use loam_math::{Projection, RasterizableSpace};
4use loam_shape::PointMesh;
5use wgpu::util::DeviceExt;
6use wgpu::{
7 BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor,
8 BindGroupLayoutEntry, BindingType, BlendComponent, BlendFactor, BlendOperation, BlendState,
9 Buffer, BufferBindingType, BufferDescriptor, BufferUsages, ColorTargetState, ColorWrites,
10 CompareFunction, DepthStencilState, Device, FragmentState, LoadOp, MultisampleState,
11 Operations, PipelineLayoutDescriptor, PrimitiveState, PrimitiveTopology, Queue,
12 RenderPassColorAttachment, RenderPassDepthStencilAttachment, RenderPassDescriptor,
13 RenderPipeline, RenderPipelineDescriptor, ShaderModuleDescriptor, ShaderSource, ShaderStages,
14 StencilState, StoreOp, TextureFormat, VertexAttribute, VertexBufferLayout, VertexFormat,
15 VertexState, VertexStepMode,
16};
17
18const POINT_RASTER_WGSL: &str = include_str!("point_raster.wgsl");
19
20#[repr(C)]
21#[derive(Copy, Clone, Debug, Pod, Zeroable)]
22pub struct PointRasterUniforms {
23 pub view_projection: [[f32; 4]; 4],
24 pub viewport_size: [f32; 2],
25 pub _pad: [f32; 2],
26}
27
28impl Default for PointRasterUniforms {
29 fn default() -> Self {
30 Self {
31 view_projection: Mat4::IDENTITY.to_cols_array_2d(),
32 viewport_size: [1.0, 1.0],
33 _pad: [0.0; 2],
34 }
35 }
36}
37
38#[repr(C)]
39#[derive(Copy, Clone, Debug, Default, Pod, Zeroable)]
40pub struct PointInstance {
41 pub pos: [f32; 3],
42 pub radius_px: f32,
43 pub color: [f32; 4],
44}
45
46pub struct PointRasterNode {
47 pipeline: RenderPipeline,
48 uniform_buf: Buffer,
49 bind_group: BindGroup,
50
51 corner_buf: Buffer,
52 index_buf: Buffer,
53 instance_buf: Buffer,
54 instance_count: u32,
55 instance_capacity: u32,
56 has_depth: bool,
57 instances_scratch: Vec<PointInstance>,
58}
59
60impl PointRasterNode {
61 pub fn new(
62 device: &Device,
63 surface_format: TextureFormat,
64 depth: crate::DepthMode,
65 convention: crate::DepthConvention,
66 ) -> Self {
67 let module = device.create_shader_module(ShaderModuleDescriptor {
68 label: Some("point_raster shader"),
69 source: ShaderSource::Wgsl(POINT_RASTER_WGSL.into()),
70 });
71
72 let uniform_buf = device.create_buffer(&BufferDescriptor {
73 label: Some("point_raster uniforms"),
74 size: std::mem::size_of::<PointRasterUniforms>() as u64,
75 usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
76 mapped_at_creation: false,
77 });
78
79 let bgl = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
80 label: Some("point_raster bgl"),
81 entries: &[BindGroupLayoutEntry {
82 binding: 0,
83 visibility: ShaderStages::VERTEX_FRAGMENT,
84 ty: BindingType::Buffer {
85 ty: BufferBindingType::Uniform,
86 has_dynamic_offset: false,
87 min_binding_size: None,
88 },
89 count: None,
90 }],
91 });
92
93 let bind_group = device.create_bind_group(&BindGroupDescriptor {
94 label: Some("point_raster bg"),
95 layout: &bgl,
96 entries: &[BindGroupEntry {
97 binding: 0,
98 resource: uniform_buf.as_entire_binding(),
99 }],
100 });
101
102 let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
103 label: Some("point_raster pipeline layout"),
104 bind_group_layouts: &[&bgl],
105 push_constant_ranges: &[],
106 });
107
108 let corner_attrs = [VertexAttribute {
109 format: VertexFormat::Uint32,
110 offset: 0,
111 shader_location: 0,
112 }];
113 let corner_layout = VertexBufferLayout {
114 array_stride: std::mem::size_of::<u32>() as u64,
115 step_mode: VertexStepMode::Vertex,
116 attributes: &corner_attrs,
117 };
118 let instance_attrs = [
119 VertexAttribute {
120 format: VertexFormat::Float32x3,
121 offset: 0,
122 shader_location: 1,
123 },
124 VertexAttribute {
125 format: VertexFormat::Float32x4,
126 offset: 16,
127 shader_location: 2,
128 },
129 VertexAttribute {
130 format: VertexFormat::Float32,
131 offset: 12,
132 shader_location: 3,
133 },
134 ];
135 let instance_layout = VertexBufferLayout {
136 array_stride: std::mem::size_of::<PointInstance>() as u64,
137 step_mode: VertexStepMode::Instance,
138 attributes: &instance_attrs,
139 };
140
141 let pipeline = device.create_render_pipeline(&RenderPipelineDescriptor {
142 label: Some("point_raster pipeline"),
143 layout: Some(&pipeline_layout),
144 vertex: VertexState {
145 module: &module,
146 entry_point: Some("vs_main"),
147 buffers: &[corner_layout, instance_layout],
148 compilation_options: Default::default(),
149 },
150 fragment: Some(FragmentState {
151 module: &module,
152 entry_point: Some("fs_main"),
153 targets: &[Some(ColorTargetState {
154 format: surface_format,
155 blend: Some(BlendState {
156 color: BlendComponent {
157 src_factor: BlendFactor::SrcAlpha,
158 dst_factor: BlendFactor::OneMinusSrcAlpha,
159 operation: BlendOperation::Add,
160 },
161 alpha: BlendComponent {
162 src_factor: BlendFactor::One,
163 dst_factor: BlendFactor::OneMinusSrcAlpha,
164 operation: BlendOperation::Add,
165 },
166 }),
167 write_mask: ColorWrites::ALL,
168 })],
169 compilation_options: Default::default(),
170 }),
171 primitive: PrimitiveState {
172 topology: PrimitiveTopology::TriangleList,
173 ..Default::default()
174 },
175 depth_stencil: depth_state(depth, convention),
176 multisample: MultisampleState {
177 count: 1,
178 ..Default::default()
179 },
180 multiview: None,
181 cache: None,
182 });
183
184 let corner_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
185 label: Some("point_raster corner buffer"),
186 contents: bytemuck::cast_slice(&[0u32, 1, 2, 3]),
187 usage: BufferUsages::VERTEX,
188 });
189
190 let index_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
191 label: Some("point_raster index buffer"),
192 contents: bytemuck::cast_slice(&[0u32, 1, 2, 2, 1, 3]),
193 usage: BufferUsages::INDEX,
194 });
195
196 let instance_buf = device.create_buffer(&BufferDescriptor {
197 label: Some("point_raster instance buffer"),
198 size: 64,
199 usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
200 mapped_at_creation: false,
201 });
202
203 Self {
204 pipeline,
205 uniform_buf,
206 bind_group,
207 corner_buf,
208 index_buf,
209 instance_buf,
210 instance_count: 0,
211 instance_capacity: 0,
212 has_depth: depth.is_active(),
213 instances_scratch: Vec::new(),
214 }
215 }
216
217 pub fn set_camera(&self, queue: &Queue, view_projection: Mat4, viewport_size: Vec2) {
218 let uniforms = PointRasterUniforms {
219 view_projection: view_projection.to_cols_array_2d(),
220 viewport_size: viewport_size.to_array(),
221 _pad: [0.0; 2],
222 };
223 queue.write_buffer(&self.uniform_buf, 0, bytemuck::bytes_of(&uniforms));
224 }
225
226 pub fn upload<S, const N: usize>(
227 &mut self,
228 device: &Device,
229 queue: &Queue,
230 mesh: &PointMesh<N>,
231 projection: &Projection<N>,
232 ) where
233 S: RasterizableSpace<N>,
234 {
235 let n_points = mesh.positions.len();
236 assert_eq!(
237 mesh.colors.len(),
238 n_points,
239 "PointMesh invariant: colors.len() == positions.len()"
240 );
241 assert_eq!(
242 mesh.sizes.len(),
243 n_points,
244 "PointMesh invariant: sizes.len() == positions.len()"
245 );
246
247 self.instances_scratch.clear();
248 let instances = &mut self.instances_scratch;
249 for ((p, color), size) in mesh
250 .positions
251 .iter()
252 .zip(mesh.colors.iter())
253 .zip(mesh.sizes.iter())
254 {
255 let p_native = S::array_to_point(*p);
256 let Some(p3) = S::project_point(p_native, projection) else {
257 continue;
258 };
259 if !p3.is_finite() {
260 continue;
261 }
262 instances.push(PointInstance {
263 pos: p3.to_array(),
264 radius_px: *size,
265 color: *color,
266 });
267 }
268
269 if instances.len() as u32 > self.instance_capacity {
270 let new_cap = (instances.len() as u32).next_power_of_two().max(16);
271 self.instance_buf = device.create_buffer(&BufferDescriptor {
272 label: Some("point_raster instance buffer"),
273 size: (new_cap as u64) * (std::mem::size_of::<PointInstance>() as u64),
274 usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
275 mapped_at_creation: false,
276 });
277 self.instance_capacity = new_cap;
278 }
279
280 if !instances.is_empty() {
281 queue.write_buffer(&self.instance_buf, 0, bytemuck::cast_slice(instances));
282 }
283 self.instance_count = instances.len() as u32;
284 }
285
286 pub fn draw_buffer(&mut self, instances: Buffer, count: u32) {
288 self.instance_buf = instances;
289 self.instance_capacity = 0;
290 self.instance_count = count;
291 }
292
293 pub fn record(
295 &self,
296 encoder: &mut wgpu::CommandEncoder,
297 view: &wgpu::TextureView,
298 depth_view: Option<&wgpu::TextureView>,
299 viewport: Option<&crate::Viewport>,
300 ) {
301 match (self.has_depth, depth_view.is_some()) {
302 (true, false) => panic!(
303 "PointRasterNode::record: pipeline was created with a depth format but no \
304 depth view was provided"
305 ),
306 (false, true) => panic!(
307 "PointRasterNode::record: pipeline was created without a depth format but a \
308 depth view was provided"
309 ),
310 _ => {}
311 }
312 if self.instance_count == 0 {
313 return;
314 }
315 let depth_attachment = depth_view.map(|dv| RenderPassDepthStencilAttachment {
316 view: dv,
317 depth_ops: Some(Operations {
318 load: LoadOp::Load,
319 store: StoreOp::Store,
320 }),
321 stencil_ops: None,
322 });
323 let mut rp = encoder.begin_render_pass(&RenderPassDescriptor {
324 label: Some("point_raster pass"),
325 color_attachments: &[Some(RenderPassColorAttachment {
326 view,
327 depth_slice: None,
328 resolve_target: None,
329 ops: Operations {
330 load: LoadOp::Load,
331 store: StoreOp::Store,
332 },
333 })],
334 depth_stencil_attachment: depth_attachment,
335 timestamp_writes: None,
336 occlusion_query_set: None,
337 });
338 if let Some(vp) = viewport {
339 vp.apply(&mut rp);
340 }
341 rp.set_pipeline(&self.pipeline);
342 rp.set_bind_group(0, &self.bind_group, &[]);
343 rp.set_vertex_buffer(0, self.corner_buf.slice(..));
344 rp.set_vertex_buffer(1, self.instance_buf.slice(..));
345 rp.set_index_buffer(self.index_buf.slice(..), wgpu::IndexFormat::Uint32);
346 rp.draw_indexed(0..6, 0, 0..self.instance_count);
347 }
348}
349
350fn depth_state(
351 depth: crate::DepthMode,
352 convention: crate::DepthConvention,
353) -> Option<DepthStencilState> {
354 depth.format().map(|format| DepthStencilState {
355 format,
356 depth_write_enabled: depth.writes(),
357 depth_compare: convention.compare(CompareFunction::LessEqual),
358 stencil: StencilState::default(),
359 bias: wgpu::DepthBiasState::default(),
360 })
361}
362
363#[cfg(test)]
364mod tests {
365 use super::*;
366
367 use glam::Vec3;
368 use loam_math::Iso3;
369 use loam_runtime::Eye;
370
371 use crate::depth_passes;
372 use crate::view::{eye_relative, root_view_projection, DEPTH_FORMAT};
373 use crate::{DepthConvention, DepthMode};
374
375 #[test]
376 fn point_raster_hides_the_nearer_point_under_the_convention_its_pipeline_declares() {
377 let eye = Eye {
378 position: [0.0, 0.0, 3.0],
379 ..Eye::default()
380 };
381 let standing = Iso3 {
382 rotation: glam::Quat::IDENTITY,
383 translation: Vec3::from(eye.position),
384 };
385 let near_point = Vec3::new(0.2, -0.1, 0.0);
386 let far_point = Vec3::new(0.2, -0.1, -4.0);
387
388 for (convention, projection) in [
389 (DepthConvention::ReversedZ, root_view_projection(&eye)),
390 (
391 DepthConvention::StandardZ,
392 Mat4::perspective_rh(eye.fov_y, eye.aspect, eye.near, eye.far)
393 * eye_relative(standing),
394 ),
395 ] {
396 let depth = |point: Vec3| {
397 let clip = projection * point.extend(1.0);
398 clip.z / clip.w
399 };
400 let state = depth_state(
401 DepthMode::ReadWrite {
402 format: DEPTH_FORMAT,
403 },
404 convention,
405 )
406 .unwrap();
407 assert!(
408 depth_passes(state.depth_compare, depth(near_point), depth(far_point)),
409 "{convention:?} with {:?} hides the nearer point",
410 state.depth_compare
411 );
412 assert!(!depth_passes(
413 state.depth_compare,
414 depth(far_point),
415 depth(near_point)
416 ));
417 }
418 }
419}