[rs] prototype of async/await for buffer mapping

.gitignore

@@ -12,5 +12,8 @@ Cargo.lock
 # Other
 .DS_Store
 
+# VSCode project
+.vscode
+
 # Output from capture example
 red.png

wgpu/Cargo.toml

@@ -46,3 +46,4 @@ log = "0.4"
 png = "0.15"
 winit = "0.20.0-alpha4"
 zerocopy = "0.2"
+futures = "0.3"

wgpu/examples/capture/main.rs

@@ -4,7 +4,7 @@
 use std::fs::File;
 use std::mem::size_of;
 
-fn main() {
+async fn run() {
     env_logger::init();
 
     let adapter = wgpu::Adapter::request(
@@ -86,21 +86,21 @@ fn main() {
     queue.submit(&[command_buffer]);
 
     // Write the buffer as a PNG
-    output_buffer.map_read_async(
-        0,
-        (size * size) as usize * size_of::<u32>(),
-        move |result: wgpu::BufferMapAsyncResult<&[u8]>| {
-            let mut png_encoder = png::Encoder::new(File::create("red.png").unwrap(), size, size);
-            png_encoder.set_depth(png::BitDepth::Eight);
-            png_encoder.set_color(png::ColorType::RGBA);
-            png_encoder
-                .write_header()
-                .unwrap()
-                .write_image_data(result.unwrap().data)
-                .unwrap();
-        },
-    );
+    if let Ok(mapping) = output_buffer.map_read(0u64, (size * size) as u64 * size_of::<u32>() as u64).await {
+        let mut png_encoder = png::Encoder::new(File::create("red.png").unwrap(), size, size);
+        png_encoder.set_depth(png::BitDepth::Eight);
+        png_encoder.set_color(png::ColorType::RGBA);
+        png_encoder
+            .write_header()
+            .unwrap()
+            .write_image_data(mapping.as_slice())
+            .unwrap();
+    }
 
     // The device will be polled when it is dropped but we can also poll it explicitly
     device.poll(true);
 }
+
+fn main() {
+    futures::executor::block_on(run());
+}

wgpu/examples/hello-compute/main.rs

@@ -1,7 +1,7 @@
 use std::{convert::TryInto as _, str::FromStr};
 use zerocopy::AsBytes as _;
 
-fn main() {
+async fn run() {
     env_logger::init();
 
     // For now this just panics if you didn't pass numbers. Could add proper error handling.
@@ -93,15 +93,17 @@ fn main() {
 
     queue.submit(&[encoder.finish()]);
 
-    // FIXME: Align and use `LayoutVerified`
-    staging_buffer.map_read_async(0, slice_size, |result| {
-        if let Ok(mapping) = result {
-            let times: Box<[u32]> = mapping
-                .data
-                .chunks_exact(4)
-                .map(|b| u32::from_ne_bytes(b.try_into().unwrap()))
-                .collect();
-            println!("Times: {:?}", times);
-        }
-    });
+    if let Ok(mapping) = staging_buffer.map_read(0u64, size).await {
+        let times : Box<[u32]> = mapping
+            .as_slice()
+            .chunks_exact(4)
+            .map(|b| u32::from_ne_bytes(b.try_into().unwrap()))
+            .collect();
+
+        println!("Times: {:?}", times);
+    }
+}
+
+fn main() {
+    futures::executor::block_on(run());
 }

wgpu/src/future.rs

@@ -0,0 +1,72 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::{Arc, Mutex};
+use std::task::{Context, Poll, Waker};
+
+struct GpuFutureInner<T> {
+    id: wgc::id::DeviceId,
+    result: Option<T>,
+    waker: Option<Waker>,
+}
+
+/// A Future that can poll the wgpu::Device
+pub struct GpuFuture<T> {
+    inner: Arc<Mutex<GpuFutureInner<T>>>,
+}
+
+/// A completion handle to set the result on a GpuFuture
+pub struct GpuFutureCompletion<T> {
+    inner: Arc<Mutex<GpuFutureInner<T>>>,
+}
+
+impl<T> Future for GpuFuture<T>
+{
+    type Output = T;
+
+    fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll<Self::Output> {
+        // grab a clone of the Arc
+        let arc = Arc::clone(&Pin::into_inner(self).inner);
+
+        // grab the device id and set the waker, but release the lock, so that the native callback can write to it
+        let device_id = {
+            let mut inner = arc.lock().unwrap();
+            inner.waker.replace(context.waker().clone());
+            inner.id
+        };
+
+        // polling the device should trigger the callback
+        wgn::wgpu_device_poll(device_id, true);
+
+        // now take the lock again, and check whether the future is complete
+        let mut inner = arc.lock().unwrap();
+        match inner.result.take() {
+            Some(value) => Poll::Ready(value),
+            _ => Poll::Pending,
+        }
+    }
+}
+
+impl<T> GpuFutureCompletion<T> {
+    pub fn complete(self, value: T) {
+        let mut inner = self.inner.lock().unwrap();
+        inner.result.replace(value);
+        if let Some(waker) = &inner.waker {
+            waker.wake_by_ref();
+        }
+    }
+}
+
+pub(crate) fn new_gpu_future<T>(id: wgc::id::DeviceId) -> (GpuFuture<T>, GpuFutureCompletion<T>) {
+    let inner = Arc::new(Mutex::new(GpuFutureInner {
+        id,
+        result: None,
+        waker: None,
+    }));
+
+    (
+        GpuFuture {
+            inner: inner.clone(),
+        },
+        GpuFutureCompletion { inner },
+    )
+}

wgpu/src/lib.rs

@@ -1,5 +1,9 @@
 //! A cross-platform graphics and compute library based on WebGPU.
 
+mod future;
+use future::GpuFutureCompletion;
+pub use future::GpuFuture;
+
 use arrayvec::ArrayVec;
 
 use std::ffi::CString;
@@ -100,6 +104,7 @@ pub struct Device {
 #[derive(Debug)]
 pub struct Buffer {
     id: wgc::id::BufferId,
+    device_id: wgc::id::DeviceId,
 }
 
 /// A handle to a texture on the GPU.
@@ -494,13 +499,14 @@ impl<'a> TextureCopyView<'a> {
 pub struct CreateBufferMapped<'a> {
     id: wgc::id::BufferId,
     pub data: &'a mut [u8],
+    device_id: wgc::id::DeviceId,
 }
 
 impl CreateBufferMapped<'_> {
     /// Unmaps the buffer from host memory and returns a [`Buffer`].
     pub fn finish(self) -> Buffer {
         wgn::wgpu_buffer_unmap(self.id);
-        Buffer { id: self.id }
+        Buffer { device_id: self.device_id, id: self.id }
     }
 }
 
@@ -790,6 +796,7 @@ impl Device {
     /// Creates a new buffer.
     pub fn create_buffer(&self, desc: &BufferDescriptor) -> Buffer {
         Buffer {
+            device_id: self.id,
             id: wgn::wgpu_device_create_buffer(self.id, desc),
         }
     }
@@ -811,7 +818,7 @@ impl Device {
 
         let data = unsafe { std::slice::from_raw_parts_mut(ptr as *mut u8, size) };
 
-        CreateBufferMapped { id, data }
+        CreateBufferMapped { device_id: self.id, id, data }
     }
 
     /// Creates a new buffer, maps it into host-visible memory, copies data from the given slice,
@@ -858,6 +865,52 @@ impl Drop for Device {
     }
 }
 
+pub struct BufferReadMapping {
+    data: *const u8,
+    size: usize,
+    buffer_id: wgc::id::BufferId,
+}
+//TODO: proper error type
+pub type BufferMapReadResult = Result<BufferReadMapping, ()>;
+
+impl BufferReadMapping
+{
+    pub fn as_slice(&self) -> &[u8] {
+        unsafe {
+            slice::from_raw_parts(self.data as *const u8, self.size)
+        }
+    }
+}
+
+impl Drop for BufferReadMapping {
+    fn drop(&mut self) {
+        wgn::wgpu_buffer_unmap(self.buffer_id);
+    }
+}
+
+pub struct BufferWriteMapping {
+    data: *mut u8,
+    size: usize,
+    buffer_id: wgc::id::BufferId,
+}
+//TODO: proper error type
+pub type BufferMapWriteResult = Result<BufferWriteMapping, ()>;
+
+impl BufferWriteMapping
+{
+    pub fn as_slice(&mut self) -> &mut [u8] {
+        unsafe {
+            slice::from_raw_parts_mut(self.data as *mut u8, self.size)
+        }
+    }
+}
+
+impl Drop for BufferWriteMapping {
+    fn drop(&mut self) {
+        wgn::wgpu_buffer_unmap(self.buffer_id);
+    }
+}
+
 pub struct BufferAsyncMapping<T> {
     pub data: T,
     buffer_id: wgc::id::BufferId,
@@ -871,97 +924,99 @@ impl<T> Drop for BufferAsyncMapping<T> {
     }
 }
 
-struct BufferMapReadAsyncUserData<F>
-where
-    F: FnOnce(BufferMapAsyncResult<&[u8]>),
+struct BufferMapReadFutureUserData
 {
-    size: usize,
-    callback: F,
+    size: BufferAddress,
+    completion: GpuFutureCompletion<BufferMapReadResult>,
     buffer_id: wgc::id::BufferId,
 }
 
-struct BufferMapWriteAsyncUserData<F>
-where
-    F: FnOnce(BufferMapAsyncResult<&mut [u8]>),
+struct BufferMapWriteFutureUserData
 {
-    size: usize,
-    callback: F,
+    size: BufferAddress,
+    completion: GpuFutureCompletion<BufferMapWriteResult>,
     buffer_id: wgc::id::BufferId,
 }
 
 impl Buffer {
-    pub fn map_read_async<F>(&self, start: BufferAddress, size: usize, callback: F)
-    where
-        F: FnOnce(BufferMapAsyncResult<&[u8]>),
+    /// Map the buffer for reading. The result is returned in a future.
+    pub fn map_read(&self, start: BufferAddress, size: BufferAddress) -> GpuFuture<BufferMapReadResult>
     {
-        extern "C" fn buffer_map_read_callback_wrapper<F>(
-            status: BufferMapAsyncStatus,
+        let (future, completion) = future::new_gpu_future(self.device_id);
+
+        extern "C" fn buffer_map_read_future_wrapper(
+            status: wgc::resource::BufferMapAsyncStatus,
             data: *const u8,
             user_data: *mut u8,
-        ) where
-            F: FnOnce(BufferMapAsyncResult<&[u8]>),
+        )
         {
             let user_data =
-                unsafe { Box::from_raw(user_data as *mut BufferMapReadAsyncUserData<F>) };
-            let data: &[u8] = unsafe { slice::from_raw_parts(data as *const u8, user_data.size) };
-            match status {
-                BufferMapAsyncStatus::Success => (user_data.callback)(Ok(BufferAsyncMapping {
+                unsafe { Box::from_raw(user_data as *mut BufferMapReadFutureUserData) };
+            if let wgc::resource::BufferMapAsyncStatus::Success = status {
+                user_data.completion.complete(Ok(BufferReadMapping {
                     data,
+                    size: user_data.size as usize,
                     buffer_id: user_data.buffer_id,
-                })),
-                _ => (user_data.callback)(Err(())),
+                }));
+            } else {
+                user_data.completion.complete(Err(()));
             }
         }
 
-        let user_data = Box::new(BufferMapReadAsyncUserData {
+        let user_data = Box::new(BufferMapReadFutureUserData {
             size,
-            callback,
+            completion,
             buffer_id: self.id,
         });
         wgn::wgpu_buffer_map_read_async(
             self.id,
             start,
-            size as BufferAddress,
-            buffer_map_read_callback_wrapper::<F>,
+            size,
+            buffer_map_read_future_wrapper,
             Box::into_raw(user_data) as *mut u8,
         );
+
+        future
     }
 
-    pub fn map_write_async<F>(&self, start: BufferAddress, size: usize, callback: F)
-    where
-        F: FnOnce(BufferMapAsyncResult<&mut [u8]>),
+    /// Map the buffer for writing. The result is returned in a future.
+    pub fn map_write(&self, start: BufferAddress, size: BufferAddress) -> GpuFuture<BufferMapWriteResult>
     {
-        extern "C" fn buffer_map_write_callback_wrapper<F>(
-            status: BufferMapAsyncStatus,
+        let (future, completion) = future::new_gpu_future(self.device_id);
+
+        extern "C" fn buffer_map_write_future_wrapper(
+            status: wgc::resource::BufferMapAsyncStatus,
             data: *mut u8,
             user_data: *mut u8,
-        ) where
-            F: FnOnce(BufferMapAsyncResult<&mut [u8]>),
+        )
         {
             let user_data =
-                unsafe { Box::from_raw(user_data as *mut BufferMapWriteAsyncUserData<F>) };
-            let data = unsafe { slice::from_raw_parts_mut(data as *mut u8, user_data.size) };
-            match status {
-                BufferMapAsyncStatus::Success => (user_data.callback)(Ok(BufferAsyncMapping {
+                unsafe { Box::from_raw(user_data as *mut BufferMapWriteFutureUserData) };
+            if let wgc::resource::BufferMapAsyncStatus::Success = status {
+                user_data.completion.complete(Ok(BufferWriteMapping {
                     data,
+                    size: user_data.size as usize,
                     buffer_id: user_data.buffer_id,
-                })),
-                _ => (user_data.callback)(Err(())),
+                }));
+            } else {
+                user_data.completion.complete(Err(()));
             }
         }
 
-        let user_data = Box::new(BufferMapWriteAsyncUserData {
+        let user_data = Box::new(BufferMapWriteFutureUserData {
             size,
-            callback,
+            completion,
             buffer_id: self.id,
         });
         wgn::wgpu_buffer_map_write_async(
             self.id,
             start,
-            size as BufferAddress,
-            buffer_map_write_callback_wrapper::<F>,
+            size,
+            buffer_map_write_future_wrapper,
             Box::into_raw(user_data) as *mut u8,
         );
+
+        future
     }
 
     /// Flushes any pending write operations and unmaps the buffer from host memory.