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    /// Download the buffer data into a [`Vec`].
22    fn download<T: bytemuck::NoUninit + bytemuck::AnyBitPattern>(
23        &self,
24        device: &wgpu::Device,
25        queue: &wgpu::Queue,
26    ) -> impl Future<Output = Result<Vec<T>, DownloadBufferError>> + Send
27    where
28        Self: Send + Sync,
29    {
30        async {
31            let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
32                label: Some("Buffer Wrapper Download Encoder"),
33            });
34            let download = self.prepare_download(device, &mut encoder);
35            queue.submit(Some(encoder.finish()));
36
37            Self::map_download(&download, device).await
38        }
39    }
40
41    /// Prepare for downloading the buffer data.
42    ///
43    /// Returns the download buffer (with [`wgpu::BufferUsages::COPY_DST`] and
44    /// [`wgpu::BufferUsages::MAP_READ`]) holding the selection buffer data.
45    fn prepare_download(
46        &self,
47        device: &wgpu::Device,
48        encoder: &mut wgpu::CommandEncoder,
49    ) -> wgpu::Buffer {
50        let download = device.create_buffer(&wgpu::BufferDescriptor {
51            label: Some("Buffer Wrapper Prepare Download Buffer"),
52            size: self.buffer().size(),
53            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
54            mapped_at_creation: false,
55        });
56
57        encoder.copy_buffer_to_buffer(self.buffer(), 0, &download, 0, download.size());
58
59        download
60    }
61
62    /// Map the download buffer to read the buffer data.
63    ///
64    /// `download` should be created with [`wgpu::BufferUsages::MAP_READ`].
65    ///
66    /// This uses [`wgpu::PollType::wait_indefinitely()`] to wait for the mapping to complete,
67    /// you can specify a custom poll type with [`BufferWrapper::map_download_with_poll_type`].
68    fn map_download<T: bytemuck::NoUninit + bytemuck::AnyBitPattern>(
69        download: &wgpu::Buffer,
70        device: &wgpu::Device,
71    ) -> impl Future<Output = Result<Vec<T>, DownloadBufferError>> + Send {
72        Self::map_download_with_poll_type(download, device, wgpu::PollType::wait_indefinitely())
73    }
74
75    /// Map the download buffer to read the buffer data with custom [`wgpu::PollType`].
76    ///
77    /// `download` should be created with [`wgpu::BufferUsages::MAP_READ`].
78    fn map_download_with_poll_type<T: bytemuck::NoUninit + bytemuck::AnyBitPattern>(
79        download: &wgpu::Buffer,
80        device: &wgpu::Device,
81        poll_type: wgpu::PollType,
82    ) -> impl Future<Output = Result<Vec<T>, DownloadBufferError>> + Send {
83        async {
84            let (tx, rx) = oneshot::channel();
85            let buffer_slice = download.slice(..);
86            buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
87                if let Err(e) = tx.send(result) {
88                    log::error!("Error occurred while sending buffer download data: {e:?}");
89                }
90            });
91            device.poll(poll_type)?;
92            rx.await??;
93
94            let edits = bytemuck::allocation::pod_collect_to_vec(&buffer_slice.get_mapped_range());
95            download.unmap();
96
97            Ok(edits)
98        }
99    }
100}
101
102impl BufferWrapper for wgpu::Buffer {
103    fn buffer(&self) -> &wgpu::Buffer {
104        self
105    }
106}
107
108/// A [`BufferWrapper`] with a fixed size that can be validated from a [`wgpu::Buffer`].
109pub trait FixedSizeBufferWrapper: BufferWrapper + TryFrom<wgpu::Buffer> {
110    /// The POD element type that defines the layout/size.
111    type Pod;
112
113    /// Returns the size in bytes of the POD element.
114    fn pod_size() -> wgpu::BufferAddress {
115        std::mem::size_of::<Self::Pod>() as wgpu::BufferAddress
116    }
117
118    /// Check if the given buffer matches the expected size.
119    ///
120    /// This is a helper function for implementing [`TryFrom`].
121    fn verify_buffer_size(buffer: &wgpu::Buffer) -> Result<(), FixedSizeBufferWrapperError> {
122        let expected_size = Self::pod_size();
123        let buffer_size = buffer.size();
124        if buffer_size != expected_size {
125            return Err(FixedSizeBufferWrapperError::BufferSizeMismatched {
126                buffer_size,
127                expected_size,
128            });
129        }
130        Ok(())
131    }
132
133    /// Download a single [`FixedSizeBufferWrapper::Pod`].
134    fn download_single(
135        &self,
136        device: &wgpu::Device,
137        queue: &wgpu::Queue,
138    ) -> impl Future<Output = Result<Self::Pod, DownloadBufferError>> + Send
139    where
140        Self: Send + Sync,
141        Self::Pod: bytemuck::NoUninit + bytemuck::AnyBitPattern,
142    {
143        async move {
144            let vec = self.download::<Self::Pod>(device, queue).await?;
145            Ok(vec.into_iter().next().expect("downloaded single element"))
146        }
147    }
148}
wgpu_3dgs_core/buffer/mod.rs

wgpu_3dgs_core/buffer/
mod.rs
