1use bytemuck::{Pod, Zeroable};
2use glam::{Mat4, Vec2};
3use loam_math::Rotor4;
4use loam_shape::LineMesh;
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 DepthStencilState, Device, FragmentState, LoadOp, MultisampleState, Operations,
11 PipelineLayoutDescriptor, PrimitiveState, PrimitiveTopology, Queue, RenderPassColorAttachment,
12 RenderPassDepthStencilAttachment, RenderPassDescriptor, RenderPipeline,
13 RenderPipelineDescriptor, ShaderModuleDescriptor, ShaderSource, ShaderStages, StencilState,
14 StoreOp, TextureFormat, VertexAttribute, VertexBufferLayout, VertexFormat, VertexState,
15 VertexStepMode,
16};
17
18const SHADER_WGSL: &str = include_str!("line_raster_static_r4.wgsl");
19
20#[repr(C)]
22#[derive(Copy, Clone, Debug, Pod, Zeroable)]
23pub struct LineRasterStaticR4Uniforms {
24 pub rotor_matrix: [[f32; 4]; 4],
25 pub view_projection: [[f32; 4]; 4],
26 pub viewport_size: [f32; 2],
27 pub focal_distance: f32,
28 pub _pad: f32,
29}
30
31impl Default for LineRasterStaticR4Uniforms {
32 fn default() -> Self {
33 Self {
34 rotor_matrix: Mat4::IDENTITY.to_cols_array_2d(),
35 view_projection: Mat4::IDENTITY.to_cols_array_2d(),
36 viewport_size: [1.0, 1.0],
37 focal_distance: 2.0,
38 _pad: 0.0,
39 }
40 }
41}
42
43#[repr(C)]
44#[derive(Copy, Clone, Debug, Default, Pod, Zeroable)]
45struct LineInstance4D {
46 start_pos: [f32; 4],
47 end_pos: [f32; 4],
48 start_color: [f32; 4],
49 end_color: [f32; 4],
50 width_px: f32,
51 _pad: [f32; 3],
52}
53
54pub struct LineRasterStaticR4Node {
55 pipeline: RenderPipeline,
56 uniform_buf: Buffer,
57 bind_group: BindGroup,
58
59 corner_buf: Buffer,
60 index_buf: Buffer,
61 instance_buf: Buffer,
62 instance_count: u32,
63 instance_capacity: u32,
64 has_depth: bool,
65}
66
67impl LineRasterStaticR4Node {
68 pub fn new(
69 device: &Device,
70 surface_format: TextureFormat,
71 depth: crate::DepthMode,
72 convention: crate::DepthConvention,
73 ) -> Self {
74 let module = device.create_shader_module(ShaderModuleDescriptor {
75 label: Some("line_raster_static_r4 shader"),
76 source: ShaderSource::Wgsl(SHADER_WGSL.into()),
77 });
78
79 let uniform_buf = device.create_buffer(&BufferDescriptor {
80 label: Some("line_raster_static_r4 uniforms"),
81 size: std::mem::size_of::<LineRasterStaticR4Uniforms>() as u64,
82 usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
83 mapped_at_creation: false,
84 });
85
86 let bgl = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
87 label: Some("line_raster_static_r4 bgl"),
88 entries: &[BindGroupLayoutEntry {
89 binding: 0,
90 visibility: ShaderStages::VERTEX_FRAGMENT,
91 ty: BindingType::Buffer {
92 ty: BufferBindingType::Uniform,
93 has_dynamic_offset: false,
94 min_binding_size: None,
95 },
96 count: None,
97 }],
98 });
99
100 let bind_group = device.create_bind_group(&BindGroupDescriptor {
101 label: Some("line_raster_static_r4 bg"),
102 layout: &bgl,
103 entries: &[BindGroupEntry {
104 binding: 0,
105 resource: uniform_buf.as_entire_binding(),
106 }],
107 });
108
109 let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
110 label: Some("line_raster_static_r4 pipeline layout"),
111 bind_group_layouts: &[&bgl],
112 push_constant_ranges: &[],
113 });
114
115 let corner_attrs = [VertexAttribute {
116 format: VertexFormat::Uint32,
117 offset: 0,
118 shader_location: 0,
119 }];
120 let corner_layout = VertexBufferLayout {
121 array_stride: std::mem::size_of::<u32>() as u64,
122 step_mode: VertexStepMode::Vertex,
123 attributes: &corner_attrs,
124 };
125 let instance_attrs = [
126 VertexAttribute {
127 format: VertexFormat::Float32x4,
128 offset: 0,
129 shader_location: 1,
130 },
131 VertexAttribute {
132 format: VertexFormat::Float32x4,
133 offset: 16,
134 shader_location: 2,
135 },
136 VertexAttribute {
137 format: VertexFormat::Float32x4,
138 offset: 32,
139 shader_location: 3,
140 },
141 VertexAttribute {
142 format: VertexFormat::Float32x4,
143 offset: 48,
144 shader_location: 4,
145 },
146 VertexAttribute {
147 format: VertexFormat::Float32,
148 offset: 64,
149 shader_location: 5,
150 },
151 ];
152 let instance_layout = VertexBufferLayout {
153 array_stride: std::mem::size_of::<LineInstance4D>() as u64,
154 step_mode: VertexStepMode::Instance,
155 attributes: &instance_attrs,
156 };
157
158 let pipeline = device.create_render_pipeline(&RenderPipelineDescriptor {
159 label: Some("line_raster_static_r4 pipeline"),
160 layout: Some(&pipeline_layout),
161 vertex: VertexState {
162 module: &module,
163 entry_point: Some(match convention {
164 crate::DepthConvention::StandardZ => "vs_main",
165 crate::DepthConvention::ReversedZ => "vs_reversed_z",
166 }),
167 buffers: &[corner_layout, instance_layout],
168 compilation_options: Default::default(),
169 },
170 fragment: Some(FragmentState {
171 module: &module,
172 entry_point: Some("fs_main"),
173 targets: &[Some(ColorTargetState {
174 format: surface_format,
175 blend: Some(BlendState {
176 color: BlendComponent {
177 src_factor: BlendFactor::SrcAlpha,
178 dst_factor: BlendFactor::OneMinusSrcAlpha,
179 operation: BlendOperation::Add,
180 },
181 alpha: BlendComponent {
182 src_factor: BlendFactor::One,
183 dst_factor: BlendFactor::OneMinusSrcAlpha,
184 operation: BlendOperation::Add,
185 },
186 }),
187 write_mask: ColorWrites::ALL,
188 })],
189 compilation_options: Default::default(),
190 }),
191 primitive: PrimitiveState {
192 topology: PrimitiveTopology::TriangleList,
193 ..Default::default()
194 },
195 depth_stencil: depth.format().map(|format| DepthStencilState {
196 format,
197 depth_write_enabled: depth.writes(),
198 depth_compare: convention.compare(wgpu::CompareFunction::LessEqual),
199 stencil: StencilState::default(),
200 bias: wgpu::DepthBiasState::default(),
201 }),
202 multisample: MultisampleState {
203 count: 1,
204 ..Default::default()
205 },
206 multiview: None,
207 cache: None,
208 });
209
210 let corner_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
211 label: Some("line_raster_static_r4 corner buffer"),
212 contents: bytemuck::cast_slice(&[0u32, 1, 2, 3]),
213 usage: BufferUsages::VERTEX,
214 });
215
216 let index_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
217 label: Some("line_raster_static_r4 index buffer"),
218 contents: bytemuck::cast_slice(&[0u32, 1, 2, 2, 1, 3]),
219 usage: BufferUsages::INDEX,
220 });
221
222 let instance_buf = device.create_buffer(&BufferDescriptor {
223 label: Some("line_raster_static_r4 instance buffer"),
224 size: 64,
225 usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
226 mapped_at_creation: false,
227 });
228
229 Self {
230 pipeline,
231 uniform_buf,
232 bind_group,
233 corner_buf,
234 index_buf,
235 instance_buf,
236 instance_count: 0,
237 instance_capacity: 0,
238 has_depth: depth.is_active(),
239 }
240 }
241
242 pub fn upload_mesh(&mut self, device: &Device, queue: &Queue, mesh: &LineMesh<4>) {
244 let n = mesh.segments.len();
245 debug_assert_eq!(mesh.colors.len(), n);
246 debug_assert_eq!(mesh.widths.len(), n);
247
248 let mut instances: Vec<LineInstance4D> = Vec::with_capacity(n);
249 for ((seg, (color_a, color_b)), &width) in mesh
250 .segments
251 .iter()
252 .zip(mesh.colors.iter())
253 .zip(mesh.widths.iter())
254 {
255 instances.push(LineInstance4D {
256 start_pos: seg.0,
257 end_pos: seg.1,
258 start_color: *color_a,
259 end_color: *color_b,
260 width_px: width,
261 _pad: [0.0; 3],
262 });
263 }
264
265 let needed = instances.len() as u32;
266 if needed > self.instance_capacity {
267 let new_cap = needed.next_power_of_two().max(16);
268 self.instance_buf = device.create_buffer(&BufferDescriptor {
269 label: Some("line_raster_static_r4 instance buffer"),
270 size: (new_cap as u64) * (std::mem::size_of::<LineInstance4D>() as u64),
271 usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
272 mapped_at_creation: false,
273 });
274 self.instance_capacity = new_cap;
275 }
276
277 if !instances.is_empty() {
278 queue.write_buffer(&self.instance_buf, 0, bytemuck::cast_slice(&instances));
279 }
280 self.instance_count = needed;
281 }
282
283 pub fn set_transform(
284 &self,
285 queue: &Queue,
286 rotor: Rotor4,
287 view_projection: Mat4,
288 viewport_size: Vec2,
289 focal_distance: f32,
290 ) {
291 let uniforms = LineRasterStaticR4Uniforms {
292 rotor_matrix: rotor.to_mat4(),
293 view_projection: view_projection.to_cols_array_2d(),
294 viewport_size: viewport_size.to_array(),
295 focal_distance,
296 _pad: 0.0,
297 };
298 queue.write_buffer(&self.uniform_buf, 0, bytemuck::bytes_of(&uniforms));
299 }
300
301 pub fn record(
303 &self,
304 encoder: &mut wgpu::CommandEncoder,
305 view: &wgpu::TextureView,
306 depth_view: Option<&wgpu::TextureView>,
307 viewport: Option<&crate::Viewport>,
308 ) {
309 match (self.has_depth, depth_view.is_some()) {
310 (true, false) => panic!(
311 "LineRasterStaticR4Node::record: pipeline was created with a depth format but \
312 no depth view was provided"
313 ),
314 (false, true) => panic!(
315 "LineRasterStaticR4Node::record: pipeline was created without a depth format \
316 but a depth view was provided"
317 ),
318 _ => {}
319 }
320 if self.instance_count == 0 {
321 return;
322 }
323 let depth_attachment = depth_view.map(|dv| RenderPassDepthStencilAttachment {
324 view: dv,
325 depth_ops: Some(Operations {
326 load: LoadOp::Load,
327 store: StoreOp::Store,
328 }),
329 stencil_ops: None,
330 });
331 let mut rp = encoder.begin_render_pass(&RenderPassDescriptor {
332 label: Some("line_raster_static_r4 pass"),
333 color_attachments: &[Some(RenderPassColorAttachment {
334 view,
335 depth_slice: None,
336 resolve_target: None,
337 ops: Operations {
338 load: LoadOp::Load,
339 store: StoreOp::Store,
340 },
341 })],
342 depth_stencil_attachment: depth_attachment,
343 timestamp_writes: None,
344 occlusion_query_set: None,
345 });
346 if let Some(vp) = viewport {
347 rp.set_viewport(
348 vp.x as f32,
349 vp.y as f32,
350 vp.width as f32,
351 vp.height as f32,
352 0.0,
353 1.0,
354 );
355 }
356 rp.set_pipeline(&self.pipeline);
357 rp.set_bind_group(0, &self.bind_group, &[]);
358 rp.set_vertex_buffer(0, self.corner_buf.slice(..));
359 rp.set_vertex_buffer(1, self.instance_buf.slice(..));
360 rp.set_index_buffer(self.index_buf.slice(..), wgpu::IndexFormat::Uint32);
361 rp.draw_indexed(0..6, 0, 0..self.instance_count);
362 }
363
364 pub fn instance_count(&self) -> u32 {
365 self.instance_count
366 }
367}