wgpu_3dgs_editor/
selection.rs

1use glam::*;
2
3use crate::{
4    SelectionBuffer, SelectionOpBuffer,
5    core::{
6        self, BufferWrapper, ComputeBundle, ComputeBundleBuilder, GaussianPod,
7        GaussianTransformBuffer, GaussiansBuffer, ModelTransformBuffer,
8    },
9    shader,
10};
11
12/// A selection expression tree.
13///
14/// ## Overview
15///
16/// This can be used to carry out operations on selection buffers, these operations are evaluated
17/// by [`SelectionBundle::evaluate`] in a recursive manner (depth-first).
18///
19/// ## Custom Operations
20///
21/// [`SelectionExpr::Unary`], [`SelectionExpr::Binary`], and [`SelectionExpr::Selection`] are
22/// custom operations that can be defined with additional [`ComputeBundle`]s, so they also
23/// carry a vector of bind groups that are used in the operation when dispatched/evaluated.
24/// These vectors should correspond to the selection bundle's bind groups starting at index 1,
25/// because index 0 must be defined by [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
26#[derive(Debug, Default)]
27pub enum SelectionExpr {
28    /// Apply an identity operation.
29    #[default]
30    Identity,
31    /// Union of the two selections.
32    Union(Box<SelectionExpr>, Box<SelectionExpr>),
33    /// Interaction of the two selections.
34    Intersection(Box<SelectionExpr>, Box<SelectionExpr>),
35    /// Difference of the two selections.
36    Difference(Box<SelectionExpr>, Box<SelectionExpr>),
37    /// Symmetric difference of the two selections.
38    SymmetricDifference(Box<SelectionExpr>, Box<SelectionExpr>),
39    /// Complement of the selection.
40    Complement(Box<SelectionExpr>),
41    /// Apply a custom unary operation.
42    Unary(usize, Box<SelectionExpr>, Vec<wgpu::BindGroup>),
43    /// Apply a custom binary operation.
44    Binary(
45        Box<SelectionExpr>,
46        usize,
47        Box<SelectionExpr>,
48        Vec<wgpu::BindGroup>,
49    ),
50    /// Create a selection.
51    Selection(usize, Vec<wgpu::BindGroup>),
52    /// Directly use a selection buffer.
53    Buffer(SelectionBuffer),
54}
55
56impl SelectionExpr {
57    /// The first u32 value for a custom operation.
58    pub const CUSTOM_OP_START: u32 = 5;
59
60    /// Create a new [`SelectionExpr::Identity`].
61    pub fn identity() -> Self {
62        Self::Identity
63    }
64
65    /// Create a new [`SelectionExpr::Union`].
66    pub fn union(self, other: Self) -> Self {
67        Self::Union(Box::new(self), Box::new(other))
68    }
69
70    /// Create a new [`SelectionExpr::Intersection`].
71    pub fn intersection(self, other: Self) -> Self {
72        Self::Intersection(Box::new(self), Box::new(other))
73    }
74
75    /// Create a new [`SelectionExpr::Difference`].
76    pub fn difference(self, other: Self) -> Self {
77        Self::Difference(Box::new(self), Box::new(other))
78    }
79
80    /// Create a new [`SelectionExpr::SymmetricDifference`].
81    pub fn symmetric_difference(self, other: Self) -> Self {
82        Self::SymmetricDifference(Box::new(self), Box::new(other))
83    }
84
85    /// Create a new [`SelectionExpr::Complement`].
86    pub fn complement(self) -> Self {
87        Self::Complement(Box::new(self))
88    }
89
90    /// Create a new [`SelectionExpr::Unary`].
91    pub fn unary(self, op: usize, bind_groups: Vec<wgpu::BindGroup>) -> Self {
92        Self::Unary(op, Box::new(self), bind_groups)
93    }
94
95    /// Create a new [`SelectionExpr::Binary`].
96    pub fn binary(self, op: usize, other: Self, bind_groups: Vec<wgpu::BindGroup>) -> Self {
97        Self::Binary(Box::new(self), op, Box::new(other), bind_groups)
98    }
99
100    /// Create a new [`SelectionExpr::Selection`].
101    pub fn selection(op: usize, bind_groups: Vec<wgpu::BindGroup>) -> Self {
102        Self::Selection(op, bind_groups)
103    }
104
105    /// Create a new [`SelectionExpr::Buffer`].
106    pub fn buffer(buffer: SelectionBuffer) -> Self {
107        Self::Buffer(buffer)
108    }
109
110    /// Update the expression in place.
111    pub fn update_with(&mut self, f: impl FnOnce(Self) -> Self) {
112        *self = f(std::mem::take(self));
113    }
114
115    /// Get the u32 associated with this expression's operation.
116    ///
117    /// The value returned is not the same as that returned by [`SelectionExpr::custom_op_index`],
118    /// but rather a value that can be used to identify the operation by the compute shader, custom
119    /// operation's index are offset by [`SelectionExpr::CUSTOM_OP_START`].
120    ///
121    /// You usually do not need to use this method, it is used internally for evaluation of the
122    /// compute shader.
123    pub fn as_u32(&self) -> Option<u32> {
124        match self {
125            SelectionExpr::Union(_, _) => Some(0),
126            SelectionExpr::Intersection(_, _) => Some(1),
127            SelectionExpr::SymmetricDifference(_, _) => Some(2),
128            SelectionExpr::Difference(_, _) => Some(3),
129            SelectionExpr::Complement(_) => Some(4),
130            SelectionExpr::Unary(op, _, _) => Some(*op as u32 + Self::CUSTOM_OP_START),
131            SelectionExpr::Binary(_, op, _, _) => Some(*op as u32 + Self::CUSTOM_OP_START),
132            SelectionExpr::Selection(op, _) => Some(*op as u32 + Self::CUSTOM_OP_START),
133            SelectionExpr::Buffer(_) => None,
134            SelectionExpr::Identity => None,
135        }
136    }
137
138    /// Whether this expression is an identity operation.
139    pub fn is_identity(&self) -> bool {
140        matches!(self, SelectionExpr::Identity)
141    }
142
143    /// Whether this expression is a primitive operation.
144    pub fn is_primitive(&self) -> bool {
145        matches!(
146            self,
147            SelectionExpr::Union(..)
148                | SelectionExpr::Intersection(..)
149                | SelectionExpr::Difference(..)
150                | SelectionExpr::SymmetricDifference(..)
151                | SelectionExpr::Complement(..)
152        )
153    }
154
155    /// Whether this expression is a custom operation.
156    pub fn is_custom(&self) -> bool {
157        matches!(
158            self,
159            SelectionExpr::Unary(..) | SelectionExpr::Binary(..) | SelectionExpr::Selection(..)
160        )
161    }
162
163    /// Whether this expression is a selection operation.
164    pub fn is_operation(&self) -> bool {
165        matches!(
166            self,
167            SelectionExpr::Union(..)
168                | SelectionExpr::Intersection(..)
169                | SelectionExpr::Difference(..)
170                | SelectionExpr::SymmetricDifference(..)
171                | SelectionExpr::Complement(..)
172                | SelectionExpr::Unary(..)
173                | SelectionExpr::Binary(..)
174                | SelectionExpr::Selection(..)
175        )
176    }
177
178    /// Whether this expression is a selection buffer.
179    pub fn is_buffer(&self) -> bool {
180        matches!(self, SelectionExpr::Buffer(_))
181    }
182
183    /// Get the custom operation index.
184    ///
185    /// This is the index of the custom operation in [`SelectionBundle::bundles`] vector.
186    pub fn custom_op_index(&self) -> Option<usize> {
187        match self {
188            SelectionExpr::Unary(op, _, _)
189            | SelectionExpr::Binary(_, op, _, _)
190            | SelectionExpr::Selection(op, _) => Some(*op),
191            _ => None,
192        }
193    }
194
195    /// Get the custom operation bind groups for this expression.
196    pub fn custom_bind_groups(&self) -> Option<&Vec<wgpu::BindGroup>> {
197        match self {
198            SelectionExpr::Unary(_, _, bind_groups) => Some(bind_groups),
199            SelectionExpr::Binary(_, _, _, bind_groups) => Some(bind_groups),
200            SelectionExpr::Selection(_, bind_groups) => Some(bind_groups),
201            _ => None,
202        }
203    }
204
205    /// Get the custom operation index and bind groups for this expression.
206    pub fn custom_op_index_and_bind_groups(&self) -> Option<(usize, &Vec<wgpu::BindGroup>)> {
207        match self {
208            SelectionExpr::Unary(op, _, bind_groups)
209            | SelectionExpr::Binary(_, op, _, bind_groups)
210            | SelectionExpr::Selection(op, bind_groups) => Some((*op, bind_groups)),
211            _ => None,
212        }
213    }
214}
215
216/// A collection of specialized [`ComputeBundle`] for selection operations.
217///
218/// ## Custom Operations
219///
220/// All [`ComputeBundle`]s supplied to this bundle as a [`SelectionExpr::Unary`],
221/// [`SelectionExpr::Binary`], or [`SelectionExpr::Selection`] custom operation must have the same
222/// bind group 0 as the [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`]. They must also
223/// not have the bind group itself, as it will be supplied automatically during evaluation.
224///
225/// Note that [`SelectionExpr::Unary`] will also get the source selection buffer, but it will be
226/// empty (all zeros), you should operate on the destination selection buffer only.
227///
228/// It is recommended to use [`ComputeBundleBuilder`] to create the custom operation bundles,
229/// and build them using [`ComputeBundleBuilder::build_without_bind_groups`].
230///
231/// ```rust no_run
232/// # pollster::block_on(async {
233/// # use wgpu_3dgs_editor::{
234/// #     Editor, MODIFIER_GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR, Modifier, SelectionBuffer,
235/// #     SelectionBundle, SelectionExpr,
236/// #     core::{
237/// #         self, BufferWrapper, GaussianPod as _, GaussianTransformBuffer,
238/// #         GaussiansBuffer, ModelTransformBuffer, glam::*,
239/// #     },
240/// #     shader,
241/// # };
242/// #
243/// # type GaussianPod = core::GaussianPodWithShSingleCov3dSingleConfigs;
244/// #
245/// # let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
246/// #
247/// # let adapter = instance
248/// #     .request_adapter(&wgpu::RequestAdapterOptions::default())
249/// #     .await
250/// #     .expect("adapter");
251/// #
252/// # let (device, _queue) = adapter
253/// #     .request_device(&wgpu::DeviceDescriptor {
254/// #         label: Some("Device"),
255/// #         required_limits: adapter.limits(),
256/// #        ..Default::default()
257/// #     })
258/// #     .await
259/// #     .expect("device");
260/// #
261/// # const MY_CUSTOM_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor =
262/// #     wgpu::BindGroupLayoutDescriptor {
263/// #         label: Some("My Custom Bind Group Layout"),
264/// #         entries: &[wgpu::BindGroupLayoutEntry {
265/// #             binding: 0,
266/// #             visibility: wgpu::ShaderStages::COMPUTE,
267/// #             ty: wgpu::BindingType::Buffer {
268/// #                 ty: wgpu::BufferBindingType::Uniform,
269/// #                 has_dynamic_offset: false,
270/// #                 min_binding_size: None,
271/// #             },
272/// #             count: None,
273/// #         }],
274/// #     };
275/// #
276/// # let my_buffer = device.create_buffer(&wgpu::BufferDescriptor {
277/// #     label: Some("My Buffer"),
278/// #     size: 4,
279/// #     usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
280/// #     mapped_at_creation: false,
281/// # });
282/// #
283/// # let my_existing_selection_buffer = SelectionBuffer::new(&device, 1024);
284/// #
285/// // Create an editor that holds the buffers for the Gaussians
286/// let editor = Editor::new(
287///     &device,
288///     &vec![core::Gaussian {
289///         rot: Quat::IDENTITY,
290///         pos: Vec3::ZERO,
291///         color: U8Vec4::ZERO,
292///         sh: [Vec3::ZERO; 15],
293///         scale: Vec3::ONE,
294///     }],
295/// );
296///
297/// // Create the selection custom operation compute bundle
298/// let my_selection_custom_op_bundle = core::ComputeBundleBuilder::new()
299///     .label("My Selection")
300///     .bind_group_layouts([
301///         &SelectionBundle::<GaussianPod>::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
302///         &MY_CUSTOM_BIND_GROUP_LAYOUT_DESCRIPTOR,
303///     ])
304///     .resolver({
305///         let mut resolver =
306///             wesl::StandardResolver::new("path/to/my/folder/containing/wesl");
307///         // Required for using core buffer structs.
308///         resolver.add_package(&core::shader::PACKAGE);
309///         // Optionally add this for some utility functions.
310///         resolver.add_package(&shader::PACKAGE);
311///         resolver
312///     })
313///     .main_shader("package::my_wesl_filename".parse().unwrap())
314///     .entry_point("main")
315///     .wesl_compile_options(wesl::CompileOptions {
316///         // Required for enabling the correct features for core struct.
317///         features: GaussianPod::wesl_features(),
318///         ..Default::default()
319///     })
320///     .build_without_bind_groups(&device)
321///     .map_err(|e| log::error!("{e}"))
322///     .expect("my selection custom op bundle");
323///
324/// // Create the selection bundle
325/// let selection_bundle =
326///     SelectionBundle::<GaussianPod>::new(&device, vec![my_selection_custom_op_bundle]);
327///
328/// // Create the bind group for your custom operation
329/// let my_selection_custom_op_bind_group = selection_bundle.bundles[0]
330///     .create_bind_group(
331///         &device,
332///         1, // Index 0 is always the Gaussians buffer
333///         [my_buffer.buffer().as_entire_binding()],
334///     )
335///     .unwrap();
336///
337/// // Create the selection expression
338/// let selection_expr = SelectionExpr::selection(
339///     0, // The bundle index for your custom operation in the selection bundle
340///     vec![my_selection_custom_op_bind_group],
341/// )
342/// .union(
343///     // Combine with other selection expressions using different functions
344///     // Here is an existing selection buffer for example
345///     SelectionExpr::Buffer(my_existing_selection_buffer),
346/// );
347///
348/// // Create a selection buffer for the result
349/// let dest_selection_buffer =
350///     SelectionBuffer::new(&device, editor.gaussians_buffer.len() as u32);
351///
352/// # let mut encoder =
353/// #     device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
354///
355/// // Evaluate the selection expression
356/// selection_bundle.evaluate(
357///     &device,
358///     &mut encoder,
359///     &selection_expr,
360///     &dest_selection_buffer,
361///     &editor.model_transform_buffer,
362///     &editor.gaussian_transform_buffer,
363///     &editor.gaussians_buffer,
364/// );
365/// # });
366/// ```
367///
368/// ## Shader Format
369///
370/// You may copy and paste the following shader bindings for
371/// [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`] into your custom selection operation
372/// shader to ensure that the bindings are correct, then add your own bindings after that.
373///
374/// ```wgsl
375/// import wgpu_3dgs_core::{
376///     gaussian::Gaussian,
377///     gaussian_transform::GaussianTransform,
378///     model_transform::{model_to_world, ModelTransform},
379/// };
380///
381/// @group(0) @binding(0)
382/// var<uniform> op: u32;
383///
384/// @group(0) @binding(1)
385/// var<storage, read> source: array<u32>;
386///
387/// @group(0) @binding(2)
388/// var<storage, read_write> dest: array<atomic<u32>>;
389///
390/// @group(0) @binding(3)
391/// var<uniform> model_transform: ModelTransform;
392///
393/// @group(0) @binding(4)
394/// var<uniform> gaussian_transform: GaussianTransform;
395///
396/// @group(0) @binding(5)
397/// var<storage, read> gaussians: array<Gaussian>;
398///
399/// // Your custom bindings here...
400///
401/// override workgroup_size: u32;
402///
403/// @compute @workgroup_size(workgroup_size)
404/// fn main(@builtin(global_invocation_id) id: vec3<u32>) {
405///     let index = id.x;
406///
407///     if index >= arrayLength(&gaussians) {
408///         return;
409///     }
410///
411///     let gaussian = gaussians[index];
412///
413///     let world_pos = model_to_world(model_transform, gaussian.position);
414///
415///     // Your custom selection operation code here...
416///
417///     let word_index = index / 32u;
418///     let bit_index = index % 32u;
419///     let bit_mask = 1u << bit_index;
420///     if /* Condition for selecting the Gaussian */ {
421///         atomicOr(&dest[word_index], bit_mask);
422///     } else {
423///         atomicAnd(&dest[word_index], ~bit_mask);
424///     }
425/// }
426/// ```
427#[derive(Debug)]
428pub struct SelectionBundle<G: GaussianPod> {
429    /// The compute bundle for primitive selection operations.
430    primitive_bundle: ComputeBundle<()>,
431    /// The compute bundles for selection custom operations.
432    pub bundles: Vec<ComputeBundle<()>>,
433    /// The Gaussian pod marker.
434    gaussian_pod_marker: std::marker::PhantomData<G>,
435}
436
437impl<G: GaussianPod> SelectionBundle<G> {
438    /// The Gaussians bind group layout descriptors.
439    ///
440    /// This bind group layout takes the following buffers:
441    /// - [`SelectionOpBuffer`]
442    /// - Source [`SelectionBuffer`]
443    /// - Destination [`SelectionBuffer`]
444    /// - [`ModelTransformBuffer`]
445    /// - [`GaussianTransformBuffer`]
446    /// - [`GaussiansBuffer`]
447    pub const GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
448        wgpu::BindGroupLayoutDescriptor {
449            label: Some("Selection Gaussians Bind Group Layout"),
450            entries: &[
451                // Selection operation buffer
452                wgpu::BindGroupLayoutEntry {
453                    binding: 0,
454                    visibility: wgpu::ShaderStages::COMPUTE,
455                    ty: wgpu::BindingType::Buffer {
456                        ty: wgpu::BufferBindingType::Uniform,
457                        has_dynamic_offset: false,
458                        min_binding_size: None,
459                    },
460                    count: None,
461                },
462                // Source selection buffer
463                wgpu::BindGroupLayoutEntry {
464                    binding: 1,
465                    visibility: wgpu::ShaderStages::COMPUTE,
466                    ty: wgpu::BindingType::Buffer {
467                        ty: wgpu::BufferBindingType::Storage { read_only: true },
468                        has_dynamic_offset: false,
469                        min_binding_size: None,
470                    },
471                    count: None,
472                },
473                // Destination selection buffer
474                wgpu::BindGroupLayoutEntry {
475                    binding: 2,
476                    visibility: wgpu::ShaderStages::COMPUTE,
477                    ty: wgpu::BindingType::Buffer {
478                        ty: wgpu::BufferBindingType::Storage { read_only: false },
479                        has_dynamic_offset: false,
480                        min_binding_size: None,
481                    },
482                    count: None,
483                },
484                // Model transform buffer
485                wgpu::BindGroupLayoutEntry {
486                    binding: 3,
487                    visibility: wgpu::ShaderStages::COMPUTE,
488                    ty: wgpu::BindingType::Buffer {
489                        ty: wgpu::BufferBindingType::Uniform,
490                        has_dynamic_offset: false,
491                        min_binding_size: None,
492                    },
493                    count: None,
494                },
495                // Gaussian transform buffer
496                wgpu::BindGroupLayoutEntry {
497                    binding: 4,
498                    visibility: wgpu::ShaderStages::COMPUTE,
499                    ty: wgpu::BindingType::Buffer {
500                        ty: wgpu::BufferBindingType::Uniform,
501                        has_dynamic_offset: false,
502                        min_binding_size: None,
503                    },
504                    count: None,
505                },
506                // Gaussians buffer
507                wgpu::BindGroupLayoutEntry {
508                    binding: 5,
509                    visibility: wgpu::ShaderStages::COMPUTE,
510                    ty: wgpu::BindingType::Buffer {
511                        ty: wgpu::BufferBindingType::Storage { read_only: true },
512                        has_dynamic_offset: false,
513                        min_binding_size: None,
514                    },
515                    count: None,
516                },
517            ],
518        };
519
520    /// Create a new selection bundle.
521    ///
522    /// `bundles` are used for [`SelectionExpr::Unary`], [`SelectionExpr::Binary`], or
523    /// [`SelectionExpr::Selection`] as custom operations, they must have the same bind group 0 as
524    /// the [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`], see documentation of
525    /// [`SelectionBundle`] for more details.
526    pub fn new(device: &wgpu::Device, bundles: Vec<ComputeBundle<()>>) -> Self {
527        let primitive_bundle = Self::create_primitive_bundle(device);
528
529        Self {
530            primitive_bundle,
531            bundles,
532            gaussian_pod_marker: std::marker::PhantomData,
533        }
534    }
535
536    /// Get the Gaussians bind group layout.
537    pub fn gaussians_bind_group_layout(&self) -> &wgpu::BindGroupLayout {
538        &self.primitive_bundle.bind_group_layouts()[0]
539    }
540
541    /// Evaluate and apply the selection expression.
542    #[allow(clippy::too_many_arguments)]
543    pub fn evaluate(
544        &self,
545        device: &wgpu::Device,
546        encoder: &mut wgpu::CommandEncoder,
547        expr: &SelectionExpr,
548        dest: &SelectionBuffer,
549        model_transform: &ModelTransformBuffer,
550        gaussian_transform: &GaussianTransformBuffer,
551        gaussians: &GaussiansBuffer<G>,
552    ) {
553        if let SelectionExpr::Identity = expr {
554            return;
555        } else if let SelectionExpr::Buffer(buffer) = expr {
556            encoder.copy_buffer_to_buffer(
557                buffer.buffer(),
558                0,
559                dest.buffer(),
560                0,
561                dest.buffer().size(),
562            );
563            return;
564        }
565
566        let d = dest;
567        let m = model_transform;
568        let g = gaussian_transform;
569        let gs = gaussians;
570
571        let op = SelectionOpBuffer::new(device, expr.as_u32().expect("operation expression"));
572        let source = SelectionBuffer::new(device, gaussians.len() as u32);
573
574        match expr {
575            SelectionExpr::Union(l, r) => {
576                self.evaluate(device, encoder, l, &source, m, g, gs);
577                self.evaluate(device, encoder, r, d, m, g, gs);
578            }
579            SelectionExpr::Intersection(l, r) => {
580                self.evaluate(device, encoder, l, &source, m, g, gs);
581                self.evaluate(device, encoder, r, d, m, g, gs);
582            }
583            SelectionExpr::Difference(l, r) => {
584                self.evaluate(device, encoder, l, &source, m, g, gs);
585                self.evaluate(device, encoder, r, d, m, g, gs);
586            }
587            SelectionExpr::SymmetricDifference(l, r) => {
588                self.evaluate(device, encoder, l, &source, m, g, gs);
589                self.evaluate(device, encoder, r, d, m, g, gs);
590            }
591            SelectionExpr::Complement(e) => {
592                self.evaluate(device, encoder, e, d, m, g, gs);
593            }
594            SelectionExpr::Unary(_, e, _) => {
595                self.evaluate(device, encoder, e, d, m, g, gs);
596            }
597            SelectionExpr::Binary(l, _, r, _) => {
598                self.evaluate(device, encoder, l, &source, m, g, gs);
599                self.evaluate(device, encoder, r, d, m, g, gs);
600            }
601            SelectionExpr::Selection(_, _) => {}
602            SelectionExpr::Identity | SelectionExpr::Buffer(_) => {
603                unreachable!();
604            }
605        }
606
607        let gaussians_bind_group = self
608            .primitive_bundle
609            .create_bind_group(
610                device,
611                0,
612                [
613                    op.buffer().as_entire_binding(),
614                    source.buffer().as_entire_binding(),
615                    d.buffer().as_entire_binding(),
616                    m.buffer().as_entire_binding(),
617                    g.buffer().as_entire_binding(),
618                    gs.buffer().as_entire_binding(),
619                ],
620            )
621            .expect("gaussians bind group");
622
623        match expr.custom_op_index_and_bind_groups() {
624            None => self.primitive_bundle.dispatch(
625                encoder,
626                (gaussians.len() as u32).div_ceil(32),
627                [&gaussians_bind_group],
628            ),
629            Some((i, bind_groups)) => {
630                let bind_groups = std::iter::once(&gaussians_bind_group)
631                    .chain(bind_groups)
632                    .collect::<Vec<_>>();
633
634                let bundle = &self.bundles[i];
635
636                bundle.dispatch(encoder, gaussians.len() as u32, bind_groups);
637            }
638        }
639    }
640
641    /// Create the selection primitive operation [`ComputeBundle`].
642    ///
643    /// - Bind group 0 is [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
644    ///
645    /// You usually do not need to use this method, it is used internally for creating the
646    /// primitive operation bundle for evaluation.
647    pub fn create_primitive_bundle(device: &wgpu::Device) -> ComputeBundle<()> {
648        ComputeBundleBuilder::new()
649            .label("Selection Primitive Operations")
650            .bind_group_layout(&Self::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR)
651            .resolver({
652                let mut resolver = wesl::PkgResolver::new();
653                resolver.add_package(&core::shader::PACKAGE);
654                resolver.add_package(&shader::PACKAGE);
655                resolver
656            })
657            .main_shader(
658                "wgpu_3dgs_editor::selection::primitive"
659                    .parse()
660                    .expect("selection::primitive module path"),
661            )
662            .entry_point("main")
663            .wesl_compile_options(wesl::CompileOptions {
664                features: G::wesl_features(),
665                ..Default::default()
666            })
667            .build_without_bind_groups(device)
668            .map_err(|e| log::error!("{e}"))
669            .expect("primitive bundle")
670    }
671
672    /// The sphere selection bind group layout descriptor.
673    ///
674    /// This bind group layout takes the following buffers:
675    /// - [`InvTransformBuffer`](crate::InvTransformBuffer)
676    pub const SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
677        wgpu::BindGroupLayoutDescriptor {
678            label: Some("Sphere Selection Bind Group Layout"),
679            entries: &[
680                // Inverse transform uniform buffer
681                wgpu::BindGroupLayoutEntry {
682                    binding: 0,
683                    visibility: wgpu::ShaderStages::COMPUTE,
684                    ty: wgpu::BindingType::Buffer {
685                        ty: wgpu::BufferBindingType::Uniform,
686                        has_dynamic_offset: false,
687                        min_binding_size: None,
688                    },
689                    count: None,
690                },
691            ],
692        };
693
694    /// Create a sphere selection custom operation.
695    ///
696    /// - Bind group 0 is [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
697    /// - Bind group 1 is [`SelectionBundle::SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR`].
698    pub fn create_sphere_bundle(device: &wgpu::Device) -> ComputeBundle<()> {
699        let mut resolver = wesl::PkgResolver::new();
700        resolver.add_package(&core::shader::PACKAGE);
701        resolver.add_package(&shader::PACKAGE);
702
703        ComputeBundleBuilder::new()
704            .label("Sphere Selection")
705            .bind_group_layouts([
706                &Self::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
707                &Self::SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR,
708            ])
709            .main_shader(
710                "wgpu_3dgs_editor::selection::sphere"
711                    .parse()
712                    .expect("selection::sphere module path"),
713            )
714            .entry_point("main")
715            .wesl_compile_options(wesl::CompileOptions {
716                features: G::wesl_features(),
717                ..Default::default()
718            })
719            .resolver(resolver)
720            .build_without_bind_groups(device)
721            .map_err(|e| log::error!("{e}"))
722            .expect("sphere selection compute bundle")
723    }
724
725    /// The box selection bind group layout descriptor.
726    ///
727    /// This bind group layout takes the following buffers:
728    /// - [`InvTransformBuffer`](crate::InvTransformBuffer)
729    pub const BOX_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
730        wgpu::BindGroupLayoutDescriptor {
731            label: Some("Box Selection Bind Group Layout"),
732            entries: &[
733                // Inverse transform uniform buffer
734                wgpu::BindGroupLayoutEntry {
735                    binding: 0,
736                    visibility: wgpu::ShaderStages::COMPUTE,
737                    ty: wgpu::BindingType::Buffer {
738                        ty: wgpu::BufferBindingType::Uniform,
739                        has_dynamic_offset: false,
740                        min_binding_size: None,
741                    },
742                    count: None,
743                },
744            ],
745        };
746
747    /// Create a box selection custom operation.
748    ///
749    /// - Bind group 0 is [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
750    /// - Bind group 1 is [`SelectionBundle::BOX_BIND_GROUP_LAYOUT_DESCRIPTOR`].
751    pub fn create_box_bundle(device: &wgpu::Device) -> ComputeBundle<()> {
752        let mut resolver = wesl::PkgResolver::new();
753        resolver.add_package(&core::shader::PACKAGE);
754        resolver.add_package(&shader::PACKAGE);
755
756        ComputeBundleBuilder::new()
757            .label("Box Selection")
758            .bind_group_layouts([
759                &Self::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
760                &Self::BOX_BIND_GROUP_LAYOUT_DESCRIPTOR,
761            ])
762            .main_shader(
763                "wgpu_3dgs_editor::selection::box"
764                    .parse()
765                    .expect("selection::box module path"),
766            )
767            .entry_point("main")
768            .wesl_compile_options(wesl::CompileOptions {
769                features: G::wesl_features(),
770                ..Default::default()
771            })
772            .resolver(resolver)
773            .build_without_bind_groups(device)
774            .map_err(|e| log::error!("{e}"))
775            .expect("box selection compute bundle")
776    }
777}