Skip to main content

polyscope_render/
vector_render.rs

1//! Vector arrow GPU rendering resources.
2
3use glam::Vec3;
4use wgpu::util::DeviceExt;
5
6/// GPU resources for rendering vectors.
7pub struct VectorRenderData {
8    pub base_buffer: wgpu::Buffer,
9    pub vector_buffer: wgpu::Buffer,
10    pub uniform_buffer: wgpu::Buffer,
11    pub bind_group: wgpu::BindGroup,
12    pub num_vectors: u32,
13}
14
15#[repr(C)]
16#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
17#[allow(clippy::pub_underscore_fields)]
18pub struct VectorUniforms {
19    pub model: [f32; 16],
20    pub length_scale: f32,
21    pub radius: f32,
22    pub _padding: [f32; 2],
23    pub color: [f32; 4],
24}
25
26impl Default for VectorUniforms {
27    fn default() -> Self {
28        #[rustfmt::skip]
29        let identity: [f32; 16] = [
30            1.0, 0.0, 0.0, 0.0,
31            0.0, 1.0, 0.0, 0.0,
32            0.0, 0.0, 1.0, 0.0,
33            0.0, 0.0, 0.0, 1.0,
34        ];
35        Self {
36            model: identity,
37            length_scale: 1.0,
38            radius: 0.005,
39            _padding: [0.0; 2],
40            color: [0.8, 0.2, 0.2, 1.0], // Red
41        }
42    }
43}
44
45impl VectorRenderData {
46    #[must_use]
47    pub fn new(
48        device: &wgpu::Device,
49        bind_group_layout: &wgpu::BindGroupLayout,
50        camera_buffer: &wgpu::Buffer,
51        bases: &[Vec3],
52        vectors: &[Vec3],
53    ) -> Self {
54        let num_vectors = bases.len().min(vectors.len()) as u32;
55
56        let base_data: Vec<f32> = bases.iter().flat_map(|p| [p.x, p.y, p.z, 0.0]).collect();
57        let base_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
58            label: Some("vector bases"),
59            contents: bytemuck::cast_slice(&base_data),
60            usage: wgpu::BufferUsages::STORAGE,
61        });
62
63        let vector_data: Vec<f32> = vectors.iter().flat_map(|v| [v.x, v.y, v.z, 0.0]).collect();
64        let vector_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
65            label: Some("vectors"),
66            contents: bytemuck::cast_slice(&vector_data),
67            usage: wgpu::BufferUsages::STORAGE,
68        });
69
70        let uniforms = VectorUniforms::default();
71        let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
72            label: Some("vector uniforms"),
73            contents: bytemuck::cast_slice(&[uniforms]),
74            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
75        });
76
77        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
78            label: Some("vector bind group"),
79            layout: bind_group_layout,
80            entries: &[
81                wgpu::BindGroupEntry {
82                    binding: 0,
83                    resource: camera_buffer.as_entire_binding(),
84                },
85                wgpu::BindGroupEntry {
86                    binding: 1,
87                    resource: uniform_buffer.as_entire_binding(),
88                },
89                wgpu::BindGroupEntry {
90                    binding: 2,
91                    resource: base_buffer.as_entire_binding(),
92                },
93                wgpu::BindGroupEntry {
94                    binding: 3,
95                    resource: vector_buffer.as_entire_binding(),
96                },
97            ],
98        });
99
100        Self {
101            base_buffer,
102            vector_buffer,
103            uniform_buffer,
104            bind_group,
105            num_vectors,
106        }
107    }
108
109    /// Updates vector uniforms.
110    pub fn update_uniforms(&self, queue: &wgpu::Queue, uniforms: &VectorUniforms) {
111        queue.write_buffer(&self.uniform_buffer, 0, bytemuck::cast_slice(&[*uniforms]));
112    }
113}