wgpu_3dgs_editor/buffer/
selection.rs

1use glam::*;
2use wgpu::util::DeviceExt;
3
4use crate::core::{self, BufferWrapper, FixedSizeBufferWrapper};
5
6/// The selection storage buffer.
7///
8/// This buffer holds the selection bitmask for Gaussians, where each bit represents whether
9/// a Gaussian is selected (1) or not (0).
10#[derive(Debug, Clone)]
11pub struct SelectionBuffer(wgpu::Buffer);
12
13impl SelectionBuffer {
14    /// Create a new selection buffer.
15    pub fn new(device: &wgpu::Device, gaussian_count: u32) -> Self {
16        Self::new_with_label(device, "", gaussian_count)
17    }
18
19    /// Create a new selection buffer with additional label.
20    pub fn new_with_label(device: &wgpu::Device, label: &str, gaussian_count: u32) -> Self {
21        let size = gaussian_count.div_ceil(32) * std::mem::size_of::<u32>() as u32;
22
23        let data = device.create_buffer(&wgpu::BufferDescriptor {
24            label: Some(format!("{label} Selection Buffer").as_str()),
25            size: size as wgpu::BufferAddress,
26            usage: Self::DEFAULT_USAGES,
27            mapped_at_creation: false,
28        });
29
30        Self(data)
31    }
32}
33
34impl BufferWrapper for SelectionBuffer {
35    const DEFAULT_USAGES: wgpu::BufferUsages = wgpu::BufferUsages::from_bits_retain(
36        wgpu::BufferUsages::STORAGE.bits() | wgpu::BufferUsages::COPY_SRC.bits(),
37    );
38
39    fn buffer(&self) -> &wgpu::Buffer {
40        &self.0
41    }
42}
43
44impl From<SelectionBuffer> for wgpu::Buffer {
45    fn from(wrapper: SelectionBuffer) -> Self {
46        wrapper.0
47    }
48}
49
50/// The selection operation uniform buffer for storing selection operations.
51///
52/// This buffer holds a single u32 value representing the index of the selection operation to be
53/// performed, this includes both primitive operations and custom operations, the value is obtained
54/// from [`SelectionExpr::as_u32`](crate::SelectionExpr::as_u32), see its documentation for more
55/// details.
56#[derive(Debug, Clone)]
57pub struct SelectionOpBuffer(wgpu::Buffer);
58
59impl SelectionOpBuffer {
60    /// Create a new selection operation buffer.
61    pub fn new(device: &wgpu::Device, op: u32) -> Self {
62        let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
63            label: Some("Selection Operation Buffer"),
64            contents: bytemuck::bytes_of(&op),
65            usage: Self::DEFAULT_USAGES,
66        });
67
68        Self(buffer)
69    }
70
71    /// Update the selection operation buffer.
72    pub fn update(&self, queue: &wgpu::Queue, op: u32) {
73        queue.write_buffer(&self.0, 0, bytemuck::bytes_of(&op));
74    }
75}
76
77impl BufferWrapper for SelectionOpBuffer {
78    fn buffer(&self) -> &wgpu::Buffer {
79        &self.0
80    }
81}
82
83impl From<SelectionOpBuffer> for wgpu::Buffer {
84    fn from(wrapper: SelectionOpBuffer) -> Self {
85        wrapper.0
86    }
87}
88
89impl TryFrom<wgpu::Buffer> for SelectionOpBuffer {
90    type Error = core::FixedSizeBufferWrapperError;
91
92    fn try_from(buffer: wgpu::Buffer) -> Result<Self, Self::Error> {
93        Self::verify_buffer_size(&buffer).map(|()| Self(buffer))
94    }
95}
96
97impl FixedSizeBufferWrapper for SelectionOpBuffer {
98    type Pod = u32;
99}
100
101/// An inverse transform uniform buffer for selection operations.
102///
103/// This buffer holds the inverse of the combined scale, rotation, and translation transform,
104/// which is used for [`SelectionBundle::create_sphere_bundle`](crate::SelectionBundle::create_sphere_bundle)
105/// and [`SelectionBundle::create_box_bundle`](crate::SelectionBundle::create_box_bundle), but
106/// can also be used for your custom selection operations.
107#[derive(Debug, Clone)]
108pub struct InvTransformBuffer(wgpu::Buffer);
109
110impl InvTransformBuffer {
111    /// Create a new inverse transform buffer.
112    pub fn new(device: &wgpu::Device) -> Self {
113        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
114            label: Some("Inverse Transform Buffer"),
115            size: std::mem::size_of::<Mat4>() as wgpu::BufferAddress,
116            usage: Self::DEFAULT_USAGES,
117            mapped_at_creation: false,
118        });
119
120        Self(buffer)
121    }
122
123    /// Update the inverse transform buffer.
124    pub fn update(&self, queue: &wgpu::Queue, inv_transform: Mat4) {
125        queue.write_buffer(&self.0, 0, bytemuck::bytes_of(&inv_transform));
126    }
127
128    /// Update the sphere selection buffer with the scale, rotation, and position.
129    pub fn update_with_scale_rot_pos(
130        &self,
131        queue: &wgpu::Queue,
132        scale: Vec3,
133        rot: Quat,
134        pos: Vec3,
135    ) {
136        let inv_transform = Mat4::from_scale_rotation_translation(scale, rot, pos).inverse();
137        self.update(queue, inv_transform);
138    }
139}
140
141impl BufferWrapper for InvTransformBuffer {
142    fn buffer(&self) -> &wgpu::Buffer {
143        &self.0
144    }
145}
146
147impl From<InvTransformBuffer> for wgpu::Buffer {
148    fn from(wrapper: InvTransformBuffer) -> Self {
149        wrapper.0
150    }
151}
152
153impl TryFrom<wgpu::Buffer> for InvTransformBuffer {
154    type Error = core::FixedSizeBufferWrapperError;
155
156    fn try_from(buffer: wgpu::Buffer) -> Result<Self, Self::Error> {
157        Self::verify_buffer_size(&buffer).map(|()| Self(buffer))
158    }
159}
160
161impl FixedSizeBufferWrapper for InvTransformBuffer {
162    type Pod = Mat4;
163}