pebble/wgpu/
buffer.rs

1use crate::threading::SpawnableFuture;
2use crate::wgpu::gpu_context::GpuContext;
3
4/// A GPU buffer that can act on itself — built via
5/// [`BufferBuilder`](super::buffers::BufferBuilder), never constructed
6/// directly. Carries its own device/queue access internally, so writing to
7/// it doesn't need a `&wgpu::Queue` threaded in from the caller.
8///
9/// Opaque by design: there's no way to reach the underlying `wgpu::Buffer`
10/// from outside this crate. Binding one into a bind group goes through
11/// [`BindGroupBuilder`](super::buffers::BindGroupBuilder), which accepts
12/// `&Buffer` directly.
13pub struct Buffer {
14    pub(crate) raw: wgpu::Buffer,
15    pub(crate) ctx: GpuContext,
16}
17
18impl Buffer {
19    pub(crate) fn new(raw: wgpu::Buffer, ctx: GpuContext) -> Self {
20        Self { raw, ctx }
21    }
22
23    /// Overwrites this buffer's contents with `data`, starting at offset 0.
24    pub fn write(&self, data: &[u8]) {
25        self.ctx.queue().write_buffer(&self.raw, 0, data);
26    }
27
28    /// Writes `data` into this buffer at a byte offset — for updating one
29    /// element of a [`DynamicBuffer`] without touching the others. Prefer
30    /// [`DynamicBuffer::write_element`], which computes the offset for you
31    /// from the buffer's own stride.
32    pub fn write_at(&self, offset: u64, data: &[u8]) {
33        self.ctx.queue().write_buffer(&self.raw, offset, data);
34    }
35
36    /// Size in bytes.
37    pub fn size(&self) -> u64 {
38        self.raw.size()
39    }
40
41    /// Copies this buffer's current contents back to the CPU. The copy
42    /// itself is submitted eagerly, right away — do not call mid-frame;
43    /// call after presenting or outside of frame encoding. Only the *wait
44    /// for the GPU to finish mapping it* is deferred into the returned
45    /// future.
46    ///
47    /// This doesn't run itself — drive it with
48    /// [`AsyncEventWriter::spawn`](crate::prelude::AsyncEventWriter::spawn) to
49    /// get the result delivered as an event, or
50    /// [`BackgroundTasks::spawn_async`](crate::threading::BackgroundTasks::spawn_async)
51    /// directly if you'd rather hold onto a
52    /// [`TaskHandle`](crate::threading::TaskHandle) and poll it yourself.
53    pub fn read(&self) -> impl SpawnableFuture<Vec<u8>> {
54        readback(self.ctx.device(), self.ctx.queue(), &self.raw)
55    }
56
57    /// Same as [`read`](Self::read) but the resolved bytes are cast to `T`.
58    pub fn read_as<T: bytemuck::Pod + Send + 'static>(&self) -> impl SpawnableFuture<Vec<T>> {
59        let bytes = self.read();
60        async move {
61            let bytes = bytes.await;
62            bytemuck::cast_slice(&bytes).to_vec()
63        }
64    }
65
66    pub(crate) fn raw(&self) -> &wgpu::Buffer {
67        &self.raw
68    }
69}
70
71/// Shared by [`Buffer::read`] and (in the future) anything else that needs a
72/// GPU→CPU readback — split out so the async staging-buffer dance lives in
73/// exactly one place.
74pub(crate) fn readback(
75    device: &wgpu::Device,
76    queue: &wgpu::Queue,
77    src: &wgpu::Buffer,
78) -> impl SpawnableFuture<Vec<u8>> {
79    let size = src.size();
80    let staging = crate::wgpu::buffers::BufferBuilder::new()
81        .usage(wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ)
82        .size(size)
83        .build_raw(device);
84
85    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
86    encoder.copy_buffer_to_buffer(src, 0, &staging, 0, size);
87    let idx = queue.submit(std::iter::once(encoder.finish()));
88
89    #[cfg(not(target_arch = "wasm32"))]
90    let device = device.clone();
91
92    async move {
93        #[cfg(not(target_arch = "wasm32"))]
94        {
95            let (tx, rx) = std::sync::mpsc::channel();
96            staging.slice(..).map_async(wgpu::MapMode::Read, move |r| {
97                let _ = tx.send(r);
98            });
99            // Native backends need an explicit poll for a queued
100            // map_async callback to ever fire — nothing else drives
101            // that here, so this blocks whichever thread is driving the
102            // future until the mapping lands. Fine: this is meant to
103            // run via `BackgroundTasks::spawn_async`, which already
104            // dedicates a worker thread to exactly this kind of wait.
105            let _ = device.poll(wgpu::PollType::Wait {
106                submission_index: Some(idx),
107                timeout: None,
108            });
109            rx.recv().unwrap().unwrap();
110            let data = staging.slice(..).get_mapped_range().to_vec();
111            staging.unmap();
112            data
113        }
114
115        #[cfg(target_arch = "wasm32")]
116        {
117            let _ = idx;
118            let mapped: std::sync::Arc<std::sync::Mutex<Option<Result<(), wgpu::BufferAsyncError>>>> =
119                std::sync::Arc::new(std::sync::Mutex::new(None));
120            let waker: std::sync::Arc<std::sync::Mutex<Option<std::task::Waker>>> =
121                std::sync::Arc::new(std::sync::Mutex::new(None));
122
123            let mapped_cb = mapped.clone();
124            let waker_cb = waker.clone();
125            staging.slice(..).map_async(wgpu::MapMode::Read, move |r| {
126                *mapped_cb.lock().unwrap() = Some(r);
127                if let Some(w) = waker_cb.lock().unwrap().take() {
128                    w.wake();
129                }
130            });
131
132            std::future::poll_fn(move |cx| {
133                if let Some(result) = mapped.lock().unwrap().take() {
134                    result.unwrap();
135                    let data = staging.slice(..).get_mapped_range().to_vec();
136                    staging.unmap();
137                    std::task::Poll::Ready(data)
138                } else {
139                    *waker.lock().unwrap() = Some(cx.waker().clone());
140                    std::task::Poll::Pending
141                }
142            })
143            .await
144        }
145    }
146}
147
148/// A buffer sized to hold many dynamically-offset elements — built via
149/// [`DynamicBufferBuilder`](super::buffers::DynamicBufferBuilder). Bundles
150/// the per-element stride (for [`write_element`](Self::write_element)) and
151/// the true (unpadded) element size (for
152/// [`BindGroupBuilder::dynamic_buffer`](super::buffers::BindGroupBuilder::dynamic_buffer))
153/// alongside the buffer itself, so neither can drift out of sync with what
154/// the buffer was actually built with.
155pub struct DynamicBuffer {
156    pub(crate) buffer: Buffer,
157    pub(crate) stride: u64,
158    pub(crate) element_size: u64,
159}
160
161impl DynamicBuffer {
162    pub(crate) fn new(buffer: Buffer, stride: u64, element_size: u64) -> Self {
163        Self { buffer, stride, element_size }
164    }
165
166    /// Writes `data` (expected to be [`element_size`](Self::element_size)
167    /// bytes) into the slot for element `index`, computing its byte offset
168    /// from this buffer's own stride.
169    pub fn write_element(&self, index: u64, data: &[u8]) {
170        self.buffer.write_at(index * self.stride, data);
171    }
172
173    /// The byte size of one element, as originally given to
174    /// [`DynamicBufferBuilder::uniform`](super::buffers::DynamicBufferBuilder::uniform)/[`storage`](super::buffers::DynamicBufferBuilder::storage).
175    pub fn element_size(&self) -> u64 {
176        self.element_size
177    }
178
179    /// The aligned per-element stride — pass `index as u32 * stride as u32`
180    /// as the dynamic offset to `set_bind_group` at draw/dispatch time.
181    pub fn stride(&self) -> u64 {
182        self.stride
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189    use crate::wgpu::buffers::BufferBuilder;
190    use crate::wgpu::test_util::with_device;
191
192    fn ctx(device: &wgpu::Device, queue: &wgpu::Queue) -> GpuContext {
193        GpuContext::new(device.clone(), queue.clone())
194    }
195
196    #[test]
197    fn write_and_write_at_do_not_panic() {
198        with_device!(device, queue, {
199            let buffer = Buffer::new(
200                BufferBuilder::new().usage(wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST).size(16).build_raw(&device),
201                ctx(&device, &queue),
202            );
203            buffer.write(&[1u8, 2, 3, 4]);
204            buffer.write_at(8, &[5u8, 6, 7, 8]);
205            assert_eq!(buffer.size(), 16);
206        });
207    }
208
209    #[test]
210    fn dynamic_buffer_write_element_does_not_panic_and_reports_its_own_sizing() {
211        with_device!(device, queue, {
212            let element_size = 16u64;
213            let count = 4u64;
214            let (usage, stride) = (
215                wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
216                crate::wgpu::buffers::dynamic_uniform_offset_stride(&device, element_size),
217            );
218            let raw = BufferBuilder::new().usage(usage).size(stride * count).build_raw(&device);
219            let dynamic = DynamicBuffer::new(Buffer::new(raw, ctx(&device, &queue)), stride, element_size);
220
221            assert_eq!(dynamic.element_size(), element_size);
222            assert_eq!(dynamic.stride(), stride);
223            assert!(dynamic.stride() >= dynamic.element_size(), "stride is alignment-padded, never smaller than the element");
224
225            // Writing the last element must not overrun the buffer — this
226            // is exactly the case a wrong stride/size calculation would
227            // panic on inside wgpu's validation.
228            dynamic.write_element(count - 1, &vec![0u8; element_size as usize]);
229        });
230    }
231
232}
233
pebble/wgpu/buffer.rs

pebble/wgpu/
buffer.rs
