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 = bytemuck::allocation::pod_collect_to_vec(&buffer_slice.get_mapped_range());
98        download.unmap();
99
100        Ok(edits)
101    }
102}
103
104impl BufferWrapper for wgpu::Buffer {
105    fn buffer(&self) -> &wgpu::Buffer {
106        self
107    }
108}
109
110/// A [`BufferWrapper`] with a fixed size that can be validated from a [`wgpu::Buffer`].
111pub trait FixedSizeBufferWrapper: BufferWrapper + TryFrom<wgpu::Buffer> {
112    /// The POD element type that defines the layout/size.
113    type Pod;
114
115    /// Returns the size in bytes of the POD element.
116    fn pod_size() -> wgpu::BufferAddress {
117        std::mem::size_of::<Self::Pod>() as wgpu::BufferAddress
118    }
119
120    /// Check if the given buffer matches the expected size.
121    ///
122    /// This is a helper function for implementing [`TryFrom`].
123    fn verify_buffer_size(buffer: &wgpu::Buffer) -> Result<(), FixedSizeBufferWrapperError> {
124        let expected_size = Self::pod_size();
125        let buffer_size = buffer.size();
126        if buffer_size != expected_size {
127            return Err(FixedSizeBufferWrapperError::BufferSizeMismatched {
128                buffer_size,
129                expected_size,
130            });
131        }
132        Ok(())
133    }
134
135    /// Download a single [`FixedSizeBufferWrapper::Pod`].
136    fn download_single(
137        &self,
138        device: &wgpu::Device,
139        queue: &wgpu::Queue,
140    ) -> impl Future<Output = Result<Self::Pod, DownloadBufferError>> + Send
141    where
142        Self: Send + Sync,
143        Self::Pod: bytemuck::NoUninit + bytemuck::AnyBitPattern,
144    {
145        async move {
146            let vec = self.download::<Self::Pod>(device, queue).await?;
147            Ok(vec.into_iter().next().expect("downloaded single element"))
148        }
149    }
150}
wgpu_3dgs_core/buffer/mod.rs

wgpu_3dgs_core/buffer/
mod.rs
