wgpu_3dgs_editor/
non_destructive_modifier.rs

1use crate::{
2    Modifier, NonDestructiveModifierCreateError,
3    core::{
4        BufferWrapper, GaussianPod, GaussianTransformBuffer, GaussiansBuffer,
5        GaussiansBufferUpdateError, ModelTransformBuffer,
6    },
7};
8
9/// A wrapper for any [`Modifier`] to not modify the original Gaussians.
10///
11/// This modifier stores a copy of the original Gaussians, whenever it is applied, it first
12/// copies original Gaussians to the destination Gaussians buffer, then applies the inner modifier.
13pub struct NonDestructiveModifier<G: GaussianPod, M: Modifier<G>> {
14    pub modifier: M,
15    pub original_gaussians: GaussiansBuffer<G>,
16}
17
18impl<G: GaussianPod, M: Modifier<G>> NonDestructiveModifier<G, M> {
19    /// Create a new non-destructive modifier.
20    ///
21    /// `gaussians` must have [`wgpu::BufferUsages::COPY_SRC`] in its usage.
22    pub fn new(
23        device: &wgpu::Device,
24        queue: &wgpu::Queue,
25        modifier: M,
26        gaussians: &GaussiansBuffer<G>,
27    ) -> Result<Self, NonDestructiveModifierCreateError> {
28        if !gaussians
29            .buffer()
30            .usage()
31            .contains(wgpu::BufferUsages::COPY_SRC)
32        {
33            return Err(NonDestructiveModifierCreateError::MissingCopySrcBufferUsage);
34        }
35
36        let original_gaussians = GaussiansBuffer::new_empty_with_usage(
37            device,
38            gaussians.len(),
39            GaussiansBuffer::<G>::DEFAULT_USAGES | wgpu::BufferUsages::COPY_SRC,
40        );
41
42        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
43            label: Some("Command Encoder"),
44        });
45
46        encoder.copy_buffer_to_buffer(
47            gaussians.buffer(),
48            0,
49            original_gaussians.buffer(),
50            0,
51            gaussians.buffer().size(),
52        );
53
54        queue.submit(Some(encoder.finish()));
55        device.poll(wgpu::PollType::Wait).inspect_err(|e| {
56            log::error!("Failed to poll device: {e:?}");
57        })?;
58
59        Ok(Self {
60            modifier,
61            original_gaussians,
62        })
63    }
64
65    /// Try apply the modifier.
66    ///
67    /// Returns [`Err`] if the original Gaussians count does not match the target Gaussians count.
68    pub fn try_apply(
69        &self,
70        device: &wgpu::Device,
71        encoder: &mut wgpu::CommandEncoder,
72        gaussians: &GaussiansBuffer<G>,
73        model_transform: &ModelTransformBuffer,
74        gaussian_transform: &GaussianTransformBuffer,
75    ) -> Result<(), GaussiansBufferUpdateError> {
76        self.try_apply_with(encoder, gaussians, |encoder, modifier, gaussians| {
77            modifier.apply(
78                device,
79                encoder,
80                gaussians,
81                model_transform,
82                gaussian_transform,
83            );
84        })
85    }
86
87    /// Apply the modifier with a function.
88    pub fn try_apply_with(
89        &self,
90        encoder: &mut wgpu::CommandEncoder,
91        gaussians: &GaussiansBuffer<G>,
92        f: impl FnOnce(&mut wgpu::CommandEncoder, &M, &GaussiansBuffer<G>),
93    ) -> Result<(), GaussiansBufferUpdateError> {
94        if self.original_gaussians.len() != gaussians.len() {
95            return Err(GaussiansBufferUpdateError::CountMismatch {
96                count: gaussians.len(),
97                expected_count: self.original_gaussians.len(),
98            });
99        }
100
101        encoder.copy_buffer_to_buffer(
102            self.original_gaussians.buffer(),
103            0,
104            gaussians.buffer(),
105            0,
106            self.original_gaussians.buffer().size(),
107        );
108
109        f(encoder, &self.modifier, gaussians);
110
111        Ok(())
112    }
113}
114
115impl<G: GaussianPod, M: Modifier<G>> Modifier<G> for NonDestructiveModifier<G, M> {
116    fn apply(
117        &self,
118        device: &wgpu::Device,
119        encoder: &mut wgpu::CommandEncoder,
120        gaussians: &GaussiansBuffer<G>,
121        model_transform: &ModelTransformBuffer,
122        gaussian_transform: &GaussianTransformBuffer,
123    ) {
124        self.try_apply(
125            device,
126            encoder,
127            gaussians,
128            model_transform,
129            gaussian_transform,
130        )
131        .expect("apply non-destructive modifier")
132    }
133}