Skip to main content

loam_render/
line_raster_static_r4.rs

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/// Matches the WGSL `TransformUniform`.
21#[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    /// Allocates per call; per-frame uploads belong in [`crate::LineRasterNode`].
243    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    /// Same contract as [`crate::LineRasterNode::record`].
302    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}