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

wgpu_3dgs_core/buffer/
mod.rs
