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::Difference(_, _) => Some(2),
128            SelectionExpr::SymmetricDifference(_, _) => 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/// It is recommended to use [`ComputeBundleBuilder`] to create the custom operation bundles,
226/// and build them using [`ComputeBundleBuilder::build_without_bind_groups`].
227///
228/// ```rust
229/// // Create the selection custom operation compute bundle
230/// let my_selection_custom_op_bundle = ComputeBundleBuilder::new()
231///     .label("My Selection")
232///     .bind_group_layouts([
233///         &SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
234///         &MY_CUSTOM_BIND_GROUP_LAYOUT_DESCRIPTOR, // Put your custom bind group layout here.
235///     ])
236///     .resolver({
237///         let mut resolver = wesl::PkgResolver::new();
238///         resolver.add_package(&core::shader::PACKAGE); // Required for using core buffer structs.
239///         resolver.add_package(&shader::PACKAGE); // Optionally add this for some utility functions.
240///         resolver
241///     })
242///     .main_shader("path::to::my::wesl::module".parse().unwrap())
243///     .entry_point("main")
244///     .wesl_compile_options(wesl::CompileOptions {
245///         features: G::wesl_features(), // Required for enabling the correct features for core structs.
246///         ..Default::default()
247///     })
248///     .build_without_bind_groups(&device)
249///     .map_err(|e| log::error!("{e}"))
250///     .expect("my selection custom op bundle");
251///
252/// // Create the selection bundle
253/// let selection_bundle = SelectionBundle::<GaussianPod>::new(
254///     &device,
255///     vec![my_selection_custom_op_bundle],
256/// );
257///
258/// // Create the bind group for your custom operation
259/// let my_selection_custom_op_bind_group = selection_bundle.bundles[0]
260///     .create_bind_group(
261///         &device,
262///         1, // Index 0 is the Gaussians buffer, so remember to start from 1 for your bind groups
263///         [my_selection_custom_op_buffer.buffer().as_entire_binding()], // Your custom bind group resources
264///     )
265///     .unwrap()
266///
267/// // Create the selection expression
268/// let selection_expr = gs::SelectionExpr::selection(
269///     0, // The bundle index for your custom operation in the selection bundle
270///     vec![my_selection_custom_op_bind_group],
271/// )
272/// .union( // Combine with other selection expressions using different functions
273///     gs::SelectionExpr::Buffer(my_existing_selection_buffer) // An existing selection buffer for example
274/// );
275///
276/// // Create a selection buffer for the result
277/// let dest_selection_buffer = SelectionBuffer::new(&device, gaussians_buffer.len() as u32);
278///
279/// // Evaluate the selection expression
280/// selection_bundle.evaluate(
281///     &device,
282///     &mut encoder,
283///     &selection_expr,
284///     &dest_selection_buffer,
285///     &model_transform,
286///     &gaussian_transform,
287///     &gaussians_buffer,
288/// );
289/// ```
290///
291/// ## Shader Format
292///
293/// You may copy and paste the following shader bindings for
294/// [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`] into your custom selection operation
295/// shader to ensure that the bindings are correct, then add your own bindings after that.
296///
297/// ```wgsl
298/// import wgpu_3dgs_core::{
299///     gaussian::Gaussian,
300///     gaussian_transform::GaussianTransform,
301///     model_transform::{model_to_world, ModelTransform},
302/// };
303///
304/// @group(0) @binding(0)
305/// var<uniform> op: u32;
306///
307/// @group(0) @binding(1)
308/// var<storage, read> source: array<u32>;
309///
310/// @group(0) @binding(2)
311/// var<storage, read_write> dest: array<atomic<u32>>;
312///
313/// @group(0) @binding(3)
314/// var<uniform> model_transform: ModelTransform;
315///
316/// @group(0) @binding(4)
317/// var<uniform> gaussian_transform: GaussianTransform;
318///
319/// @group(0) @binding(5)
320/// var<storage, read> gaussians: array<Gaussian>;
321///
322/// // Your custom bindings here...
323///
324/// override workgroup_size: u32;
325///
326/// @compute @workgroup_size(workgroup_size)
327/// fn main(@builtin(global_invocation_id) id: vec3<u32>) {
328///     let index = id.x;
329///
330///     if index >= arrayLength(&gaussians) {
331///         return;
332///     }
333///
334///     let gaussian = gaussians[index];
335///
336///     let world_pos = model_to_world(model_transform, gaussian.position);
337///
338///     // Your custom selection operation code here...
339///
340///     let word_index = index / 32u;
341///     let bit_index = index % 32u;
342///     let bit_mask = 1u << bit_index;
343///     if /* Condition for selecting the Gaussian */ {
344///         atomicOr(&dest[word_index], bit_mask);
345///     } else {
346///         atomicAnd(&dest[word_index], ~bit_mask);
347///     }
348/// }
349/// ```
350#[derive(Debug)]
351pub struct SelectionBundle<G: GaussianPod> {
352    /// The compute bundle for primitive selection operations.
353    primitive_bundle: ComputeBundle<()>,
354    /// The compute bundles for selection custom operations.
355    pub bundles: Vec<ComputeBundle<()>>,
356    /// The Gaussian pod marker.
357    gaussian_pod_marker: std::marker::PhantomData<G>,
358}
359
360impl<G: GaussianPod> SelectionBundle<G> {
361    /// The Gaussians bind group layout descriptors.
362    ///
363    /// This bind group layout takes the following buffers:
364    /// - [`SelectionOpBuffer`]
365    /// - Source [`SelectionBuffer`]
366    /// - Destination [`SelectionBuffer`]
367    /// - [`ModelTransformBuffer`]
368    /// - [`GaussianTransformBuffer`]
369    /// - [`GaussiansBuffer`]
370    pub const GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
371        wgpu::BindGroupLayoutDescriptor {
372            label: Some("Selection Gaussians Bind Group Layout"),
373            entries: &[
374                // Selection operation buffer
375                wgpu::BindGroupLayoutEntry {
376                    binding: 0,
377                    visibility: wgpu::ShaderStages::COMPUTE,
378                    ty: wgpu::BindingType::Buffer {
379                        ty: wgpu::BufferBindingType::Uniform,
380                        has_dynamic_offset: false,
381                        min_binding_size: None,
382                    },
383                    count: None,
384                },
385                // Source selection buffer
386                wgpu::BindGroupLayoutEntry {
387                    binding: 1,
388                    visibility: wgpu::ShaderStages::COMPUTE,
389                    ty: wgpu::BindingType::Buffer {
390                        ty: wgpu::BufferBindingType::Storage { read_only: true },
391                        has_dynamic_offset: false,
392                        min_binding_size: None,
393                    },
394                    count: None,
395                },
396                // Destination selection buffer
397                wgpu::BindGroupLayoutEntry {
398                    binding: 2,
399                    visibility: wgpu::ShaderStages::COMPUTE,
400                    ty: wgpu::BindingType::Buffer {
401                        ty: wgpu::BufferBindingType::Storage { read_only: false },
402                        has_dynamic_offset: false,
403                        min_binding_size: None,
404                    },
405                    count: None,
406                },
407                // Model transform buffer
408                wgpu::BindGroupLayoutEntry {
409                    binding: 3,
410                    visibility: wgpu::ShaderStages::COMPUTE,
411                    ty: wgpu::BindingType::Buffer {
412                        ty: wgpu::BufferBindingType::Uniform,
413                        has_dynamic_offset: false,
414                        min_binding_size: None,
415                    },
416                    count: None,
417                },
418                // Gaussian transform buffer
419                wgpu::BindGroupLayoutEntry {
420                    binding: 4,
421                    visibility: wgpu::ShaderStages::COMPUTE,
422                    ty: wgpu::BindingType::Buffer {
423                        ty: wgpu::BufferBindingType::Uniform,
424                        has_dynamic_offset: false,
425                        min_binding_size: None,
426                    },
427                    count: None,
428                },
429                // Gaussians buffer
430                wgpu::BindGroupLayoutEntry {
431                    binding: 5,
432                    visibility: wgpu::ShaderStages::COMPUTE,
433                    ty: wgpu::BindingType::Buffer {
434                        ty: wgpu::BufferBindingType::Storage { read_only: true },
435                        has_dynamic_offset: false,
436                        min_binding_size: None,
437                    },
438                    count: None,
439                },
440            ],
441        };
442
443    /// Create a new selection bundle.
444    ///
445    /// `bundles` are used for [`SelectionExpr::Unary`], [`SelectionExpr::Binary`], or
446    /// [`SelectionExpr::Selection`] as custom operations, they must have the same bind group 0 as
447    /// the [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`], see documentation of
448    /// [`SelectionBundle`] for more details.
449    pub fn new<'a>(device: &wgpu::Device, bundles: Vec<ComputeBundle<()>>) -> Self {
450        let primitive_bundle = Self::create_primitive_bundle(device);
451
452        Self {
453            primitive_bundle,
454            bundles,
455            gaussian_pod_marker: std::marker::PhantomData,
456        }
457    }
458
459    /// Get the Gaussians bind group layout.
460    pub fn gaussians_bind_group_layout(&self) -> &wgpu::BindGroupLayout {
461        &self.primitive_bundle.bind_group_layouts()[0]
462    }
463
464    /// Evaluate and apply the selection expression.
465    pub fn evaluate(
466        &self,
467        device: &wgpu::Device,
468        encoder: &mut wgpu::CommandEncoder,
469        expr: &SelectionExpr,
470        dest: &SelectionBuffer,
471        model_transform: &ModelTransformBuffer,
472        gaussian_transform: &GaussianTransformBuffer,
473        gaussians: &GaussiansBuffer<G>,
474    ) {
475        if let SelectionExpr::Identity = expr {
476            return;
477        } else if let SelectionExpr::Buffer(buffer) = expr {
478            encoder.copy_buffer_to_buffer(
479                &buffer.buffer(),
480                0,
481                &dest.buffer(),
482                0,
483                dest.buffer().size(),
484            );
485            return;
486        }
487
488        let d = dest;
489        let m = model_transform;
490        let g = gaussian_transform;
491        let gs = gaussians;
492
493        let op = SelectionOpBuffer::new(device, expr.as_u32().expect("operation expression"));
494        let source = SelectionBuffer::new(device, gaussians.len() as u32);
495
496        match expr {
497            SelectionExpr::Union(l, r) => {
498                self.evaluate(device, encoder, l, &source, m, g, gs);
499                self.evaluate(device, encoder, r, d, m, g, gs);
500            }
501            SelectionExpr::Intersection(l, r) => {
502                self.evaluate(device, encoder, l, &source, m, g, gs);
503                self.evaluate(device, encoder, r, d, m, g, gs);
504            }
505            SelectionExpr::Difference(l, r) => {
506                self.evaluate(device, encoder, l, &source, m, g, gs);
507                self.evaluate(device, encoder, r, d, m, g, gs);
508            }
509            SelectionExpr::SymmetricDifference(l, r) => {
510                self.evaluate(device, encoder, l, &source, m, g, gs);
511                self.evaluate(device, encoder, r, d, m, g, gs);
512            }
513            SelectionExpr::Complement(e) => {
514                self.evaluate(device, encoder, e, d, m, g, gs);
515            }
516            SelectionExpr::Unary(_, e, _) => {
517                self.evaluate(device, encoder, e, d, m, g, gs);
518            }
519            SelectionExpr::Binary(l, _, r, _) => {
520                self.evaluate(device, encoder, l, &source, m, g, gs);
521                self.evaluate(device, encoder, r, d, m, g, gs);
522            }
523            SelectionExpr::Selection(_, _) => {}
524            SelectionExpr::Identity | SelectionExpr::Buffer(_) => {
525                unreachable!();
526            }
527        }
528
529        let gaussians_bind_group = self
530            .primitive_bundle
531            .create_bind_group(
532                device,
533                0,
534                [
535                    op.buffer().as_entire_binding(),
536                    source.buffer().as_entire_binding(),
537                    d.buffer().as_entire_binding(),
538                    m.buffer().as_entire_binding(),
539                    g.buffer().as_entire_binding(),
540                    gs.buffer().as_entire_binding(),
541                ],
542            )
543            .expect("gaussians bind group");
544
545        match expr.custom_op_index_and_bind_groups() {
546            None => self.primitive_bundle.dispatch(
547                encoder,
548                (gaussians.len() as u32).div_ceil(32),
549                [&gaussians_bind_group],
550            ),
551            Some((i, bind_groups)) => {
552                let bind_groups = std::iter::once(&gaussians_bind_group)
553                    .chain(bind_groups)
554                    .collect::<Vec<_>>();
555
556                let bundle = &self.bundles[i as usize];
557
558                bundle.dispatch(encoder, gaussians.len() as u32, bind_groups);
559            }
560        }
561    }
562
563    /// Create the selection primitive operation [`ComputeBundle`].
564    ///
565    /// - Bind group 0 is [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
566    ///
567    /// You usually do not need to use this method, it is used internally for creating the
568    /// primitive operation bundle for evaluation.
569    pub fn create_primitive_bundle(device: &wgpu::Device) -> ComputeBundle<()> {
570        ComputeBundleBuilder::new()
571            .label("Selection Primitive Operations")
572            .bind_group_layout(&Self::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR)
573            .resolver({
574                let mut resolver = wesl::PkgResolver::new();
575                resolver.add_package(&core::shader::PACKAGE);
576                resolver.add_package(&shader::PACKAGE);
577                resolver
578            })
579            .main_shader(
580                "wgpu_3dgs_editor::selection::primitive"
581                    .parse()
582                    .expect("selection::primitive module path"),
583            )
584            .entry_point("main")
585            .wesl_compile_options(wesl::CompileOptions {
586                features: G::wesl_features(),
587                ..Default::default()
588            })
589            .build_without_bind_groups(&device)
590            .map_err(|e| log::error!("{e}"))
591            .expect("primitive bundle")
592    }
593
594    /// The sphere selection bind group layout descriptor.
595    ///
596    /// This bind group layout takes the following buffers:
597    /// - [`InvTransformBuffer`](crate::InvTransformBuffer)
598    pub const SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
599        wgpu::BindGroupLayoutDescriptor {
600            label: Some("Sphere Selection Bind Group Layout"),
601            entries: &[
602                // Inverse transform uniform buffer
603                wgpu::BindGroupLayoutEntry {
604                    binding: 0,
605                    visibility: wgpu::ShaderStages::COMPUTE,
606                    ty: wgpu::BindingType::Buffer {
607                        ty: wgpu::BufferBindingType::Uniform,
608                        has_dynamic_offset: false,
609                        min_binding_size: None,
610                    },
611                    count: None,
612                },
613            ],
614        };
615
616    /// Create a sphere selection custom operation.
617    ///
618    /// - Bind group 0 is [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
619    /// - Bind group 1 is [`SelectionBundle::SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR`].
620    pub fn create_sphere_bundle(device: &wgpu::Device) -> ComputeBundle<()> {
621        let mut resolver = wesl::PkgResolver::new();
622        resolver.add_package(&core::shader::PACKAGE);
623        resolver.add_package(&shader::PACKAGE);
624
625        ComputeBundleBuilder::new()
626            .label("Sphere Selection")
627            .bind_group_layouts([
628                &Self::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
629                &Self::SPHERE_BIND_GROUP_LAYOUT_DESCRIPTOR,
630            ])
631            .main_shader(
632                "wgpu_3dgs_editor::selection::sphere"
633                    .parse()
634                    .expect("selection::sphere module path"),
635            )
636            .entry_point("main")
637            .wesl_compile_options(wesl::CompileOptions {
638                features: G::wesl_features(),
639                ..Default::default()
640            })
641            .resolver(resolver)
642            .build_without_bind_groups(device)
643            .map_err(|e| log::error!("{e}"))
644            .expect("sphere selection compute bundle")
645    }
646
647    /// The box selection bind group layout descriptor.
648    ///
649    /// This bind group layout takes the following buffers:
650    /// - [`InvTransformBuffer`](crate::InvTransformBuffer)
651    pub const BOX_BIND_GROUP_LAYOUT_DESCRIPTOR: wgpu::BindGroupLayoutDescriptor<'static> =
652        wgpu::BindGroupLayoutDescriptor {
653            label: Some("Box Selection Bind Group Layout"),
654            entries: &[
655                // Inverse transform uniform buffer
656                wgpu::BindGroupLayoutEntry {
657                    binding: 0,
658                    visibility: wgpu::ShaderStages::COMPUTE,
659                    ty: wgpu::BindingType::Buffer {
660                        ty: wgpu::BufferBindingType::Uniform,
661                        has_dynamic_offset: false,
662                        min_binding_size: None,
663                    },
664                    count: None,
665                },
666            ],
667        };
668
669    /// Create a box selection custom operation.
670    ///
671    /// - Bind group 0 is [`SelectionBundle::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR`].
672    /// - Bind group 1 is [`SelectionBundle::BOX_BIND_GROUP_LAYOUT_DESCRIPTOR`].
673    pub fn create_box_bundle(device: &wgpu::Device) -> ComputeBundle<()> {
674        let mut resolver = wesl::PkgResolver::new();
675        resolver.add_package(&core::shader::PACKAGE);
676        resolver.add_package(&shader::PACKAGE);
677
678        ComputeBundleBuilder::new()
679            .label("Box Selection")
680            .bind_group_layouts([
681                &Self::GAUSSIANS_BIND_GROUP_LAYOUT_DESCRIPTOR,
682                &Self::BOX_BIND_GROUP_LAYOUT_DESCRIPTOR,
683            ])
684            .main_shader(
685                "wgpu_3dgs_editor::selection::box"
686                    .parse()
687                    .expect("selection::box module path"),
688            )
689            .entry_point("main")
690            .wesl_compile_options(wesl::CompileOptions {
691                features: G::wesl_features(),
692                ..Default::default()
693            })
694            .resolver(resolver)
695            .build_without_bind_groups(device)
696            .map_err(|e| log::error!("{e}"))
697            .expect("box selection compute bundle")
698    }
699}