wgpu_3dgs_core/buffer/
mod.rs

1mod gaussian;
2mod gaussian_transform;
3mod model_transform;
4
5pub use gaussian::*;
6pub use gaussian_transform::*;
7pub use model_transform::*;
8
9use crate::{DownloadBufferError, FixedSizeBufferWrapperError};
10
11/// A trait to to enable any wrapper to act like a [`wgpu::Buffer`].
12pub trait BufferWrapper: Into<wgpu::Buffer> {
13    /// The default usages.
14    const DEFAULT_USAGES: wgpu::BufferUsages = wgpu::BufferUsages::from_bits_retain(
15        wgpu::BufferUsages::UNIFORM.bits() | wgpu::BufferUsages::COPY_DST.bits(),
16    );
17
18    /// Returns a reference to the buffer data.
19    fn buffer(&self) -> &wgpu::Buffer;
20}
21
22impl BufferWrapper for wgpu::Buffer {
23    fn buffer(&self) -> &wgpu::Buffer {
24        self
25    }
26}
27
28/// A trait to enable any [`BufferWrapper`] to download the buffer data.
29pub trait DownloadableBufferWrapper: BufferWrapper + Send + Sync {
30    /// Download the buffer data.
31    fn download<T: bytemuck::NoUninit + bytemuck::AnyBitPattern>(
32        &self,
33        device: &wgpu::Device,
34        queue: &wgpu::Queue,
35    ) -> impl Future<Output = Result<Vec<T>, DownloadBufferError>> + Send {
36        async {
37            let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
38                label: Some("Selection Download Encoder"),
39            });
40            let download = self.prepare_download(device, &mut encoder);
41            queue.submit(Some(encoder.finish()));
42
43            Self::map_download(&download, device).await
44        }
45    }
46
47    /// Prepare for downloading the buffer data.
48    ///
49    /// Returns the download buffer (with [`wgpu::BufferUsages::COPY_DST`] and
50    /// [`wgpu::BufferUsages::MAP_READ`]) holding the selection buffer data.
51    fn prepare_download(
52        &self,
53        device: &wgpu::Device,
54        encoder: &mut wgpu::CommandEncoder,
55    ) -> wgpu::Buffer {
56        let download = device.create_buffer(&wgpu::BufferDescriptor {
57            label: Some(format!("Selection Download Buffer").as_str()),
58            size: self.buffer().size(),
59            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
60            mapped_at_creation: false,
61        });
62
63        encoder.copy_buffer_to_buffer(self.buffer(), 0, &download, 0, download.size());
64
65        download
66    }
67
68    /// Map the download buffer to read the buffer data.
69    ///
70    /// `download` should be created with [`wgpu::BufferUsages::MAP_READ`].
71    fn map_download<T: bytemuck::NoUninit + bytemuck::AnyBitPattern>(
72        download: &wgpu::Buffer,
73        device: &wgpu::Device,
74    ) -> impl Future<Output = Result<Vec<T>, DownloadBufferError>> + Send {
75        async {
76            let (tx, rx) = oneshot::channel();
77            let buffer_slice = download.slice(..);
78            buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
79                if let Err(e) = tx.send(result) {
80                    log::error!("Error occurred while sending Gaussian selection: {e:?}");
81                }
82            });
83            device.poll(wgpu::PollType::Wait)?;
84            rx.await??;
85
86            let edits = bytemuck::allocation::pod_collect_to_vec(&buffer_slice.get_mapped_range());
87            download.unmap();
88
89            Ok(edits)
90        }
91    }
92}
93
94impl<T: BufferWrapper + Send + Sync> DownloadableBufferWrapper for T {}
95
96/// A [`BufferWrapper`] with a fixed size that can be validated from a [`wgpu::Buffer`].
97pub trait FixedSizeBufferWrapper: BufferWrapper + TryFrom<wgpu::Buffer> {
98    /// The POD element type that defines the layout/size.
99    type Pod;
100
101    /// Returns the size in bytes of the POD element.
102    fn pod_size() -> wgpu::BufferAddress {
103        std::mem::size_of::<Self::Pod>() as wgpu::BufferAddress
104    }
105
106    /// Check if the given buffer matches the expected size.
107    ///
108    /// This is a helper function for implementing [`TryFrom`].
109    fn verify_buffer_size(buffer: &wgpu::Buffer) -> Result<(), FixedSizeBufferWrapperError> {
110        let expected_size = Self::pod_size();
111        let buffer_size = buffer.size();
112        if buffer_size != expected_size {
113            return Err(FixedSizeBufferWrapperError::BufferSizeMismatched {
114                buffer_size,
115                expected_size,
116            });
117        }
118        Ok(())
119    }
120}
wgpu_3dgs_core/buffer/mod.rs

wgpu_3dgs_core/buffer/
mod.rs
