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