Implement mapReadAsync function of GPUBuffer

Implemented the `mapReadAsync` and fixed the `unmap` functions of `GPUBuffer`.
Added `mapped` internal slot for tracking the ArrayBuffer/Promise.
Added more states to the `GPUBufferState` enum.

components/script/dom/bindings/codegen/Bindings.conf

@@ -150,6 +150,10 @@ DOMInterfaces = {
 
 'GPUAdapter': {
     'inRealms': ['RequestDevice'],
+},
+
+'GPUBuffer': {
+    'inRealms': ['MapReadAsync'],
 }
 
 }

components/script/dom/gpubuffer.rs

@@ -6,17 +6,36 @@ use crate::dom::bindings::cell::{DomRefCell, Ref};
 use crate::dom::bindings::codegen::Bindings::GPUBufferBinding::{
     self, GPUBufferMethods, GPUBufferSize,
 };
+use crate::dom::bindings::error::Error;
+use crate::dom::bindings::reflector::DomObject;
 use crate::dom::bindings::reflector::{reflect_dom_object, Reflector};
 use crate::dom::bindings::root::DomRoot;
 use crate::dom::bindings::str::DOMString;
+use crate::dom::bindings::trace::RootedTraceableBox;
 use crate::dom::globalscope::GlobalScope;
+use crate::dom::gpu::{response_async, AsyncWGPUListener};
+use crate::dom::promise::Promise;
+use crate::realms::InRealm;
 use dom_struct::dom_struct;
+use js::jsapi::{Heap, JSObject};
+use js::jsval::UndefinedValue;
+use js::rust::jsapi_wrapped::{DetachArrayBuffer, IsPromiseObject, RejectPromise};
+use js::rust::MutableHandle;
+use js::typedarray::{ArrayBuffer, CreateWith};
 use std::cell::Cell;
-use webgpu::{WebGPU, WebGPUBuffer, WebGPUDevice, WebGPURequest};
+use std::ptr;
+use std::rc::Rc;
+use webgpu::{
+    wgpu::resource::BufferUsage, WebGPU, WebGPUBuffer, WebGPUDevice, WebGPURequest, WebGPUResponse,
+};
 
-#[derive(MallocSizeOf)]
+// https://gpuweb.github.io/gpuweb/#buffer-state
+#[derive(Clone, MallocSizeOf)]
 pub enum GPUBufferState {
-    Mapped,
+    MappedForReading,
+    MappedForWriting,
+    MappedPendingForReading,
+    MappedPendingForWriting,
     Unmapped,
     Destroyed,
 }
@@ -24,7 +43,7 @@ pub enum GPUBufferState {
 #[dom_struct]
 pub struct GPUBuffer {
     reflector_: Reflector,
-    #[ignore_malloc_size_of = "channels are hard"]
+    #[ignore_malloc_size_of = "defined in webgpu"]
     channel: WebGPU,
     label: DomRefCell<Option<DOMString>>,
     size: GPUBufferSize,
@@ -33,6 +52,8 @@ pub struct GPUBuffer {
     buffer: WebGPUBuffer,
     device: WebGPUDevice,
     valid: Cell<bool>,
+    #[ignore_malloc_size_of = "defined in mozjs"]
+    mapping: RootedTraceableBox<Heap<*mut JSObject>>,
 }
 
 impl GPUBuffer {
@@ -44,6 +65,7 @@ impl GPUBuffer {
         size: GPUBufferSize,
         usage: u32,
         valid: bool,
+        mapping: RootedTraceableBox<Heap<*mut JSObject>>,
     ) -> GPUBuffer {
         Self {
             reflector_: Reflector::new(),
@@ -55,6 +77,7 @@ impl GPUBuffer {
             valid: Cell::new(valid),
             device,
             buffer,
+            mapping,
         }
     }
 
@@ -68,10 +91,11 @@ impl GPUBuffer {
         size: GPUBufferSize,
         usage: u32,
         valid: bool,
+        mapping: RootedTraceableBox<Heap<*mut JSObject>>,
     ) -> DomRoot<GPUBuffer> {
         reflect_dom_object(
             Box::new(GPUBuffer::new_inherited(
-                channel, buffer, device, state, size, usage, valid,
+                channel, buffer, device, state, size, usage, valid, mapping,
             )),
             global,
             GPUBufferBinding::Wrap,
@@ -104,19 +128,59 @@ impl Drop for GPUBuffer {
 }
 
 impl GPUBufferMethods for GPUBuffer {
+    #[allow(unsafe_code)]
     /// https://gpuweb.github.io/gpuweb/#dom-gpubuffer-unmap
     fn Unmap(&self) {
-        self.channel
+        let cx = self.global().get_cx();
-            .0
+        // Step 1
-            .send(WebGPURequest::UnmapBuffer(self.buffer))
+        match *self.state.borrow() {
-            .unwrap();
+            GPUBufferState::Unmapped | GPUBufferState::Destroyed => {
+                // TODO: Record validation error on the current scope
+                return;
+            },
+            GPUBufferState::MappedForWriting => {
+                // Step 3.1
+                match ArrayBuffer::from(self.mapping.get()) {
+                    Ok(array_buffer) => {
+                        self.channel
+                            .0
+                            .send(WebGPURequest::UnmapBuffer(
+                                self.device.0,
+                                self.id(),
+                                array_buffer.to_vec(),
+                            ))
+                            .unwrap();
+                        // Step 3.2
+                        unsafe {
+                            DetachArrayBuffer(*cx, self.mapping.handle());
+                        }
+                    },
+                    _ => {
+                        // Step 2
+                        unsafe {
+                            if IsPromiseObject(self.mapping.handle()) {
+                                let err = Error::Abort;
+                                rooted!(in(*cx) let mut undef = UndefinedValue());
+                                err.to_jsval(*cx, &self.global(), undef.handle_mut());
+                                RejectPromise(*cx, self.mapping.handle(), undef.handle());
+                            };
+                        }
+                    },
+                };
+            },
+            _ => {},
+        };
+        // Step 3.3
+        self.mapping.set(ptr::null_mut());
+        // Step 4
         *self.state.borrow_mut() = GPUBufferState::Unmapped;
     }
 
     /// https://gpuweb.github.io/gpuweb/#dom-gpubuffer-destroy
     fn Destroy(&self) {
-        match *self.state.borrow() {
+        let state = self.state.borrow().clone();
-            GPUBufferState::Mapped => {
+        match state {
+            GPUBufferState::MappedForReading | GPUBufferState::MappedForWriting => {
                 self.Unmap();
             },
             _ => {},
@@ -128,6 +192,72 @@ impl GPUBufferMethods for GPUBuffer {
         *self.state.borrow_mut() = GPUBufferState::Destroyed;
     }
 
+    #[allow(unsafe_code)]
+    /// https://gpuweb.github.io/gpuweb/#dom-gpubuffer-mapreadasync
+    fn MapReadAsync(&self, comp: InRealm) -> Rc<Promise> {
+        // Step 1 & 2
+        let promise = Promise::new_in_current_realm(&self.global(), comp);
+        match *self.state.borrow() {
+            GPUBufferState::Unmapped => {
+                match BufferUsage::from_bits(self.usage) {
+                    Some(usage) => {
+                        if !usage.contains(BufferUsage::MAP_READ) {
+                            // TODO: Record validation error on the current scope
+                            promise.reject_error(Error::Abort);
+                            return promise;
+                        };
+                    },
+                    None => {
+                        promise.reject_error(Error::Abort);
+                        return promise;
+                    },
+                }
+            },
+            _ => {
+                promise.reject_error(Error::Abort);
+                return promise;
+            },
+        }
+        // Step 3
+        self.mapping.set(*promise.promise_obj());
+        // Step 4
+        *self.state.borrow_mut() = GPUBufferState::MappedPendingForReading;
+
+        // Step 5.1
+        if unsafe {
+            ArrayBuffer::create(
+                *self.global().get_cx(),
+                CreateWith::Length(self.size as u32),
+                MutableHandle::from_raw(self.mapping.handle_mut()),
+            )
+        }
+        .is_err()
+        {
+            promise.reject_error(Error::Operation);
+            return promise;
+        }
+
+        let sender = response_async(&promise, self);
+        if self
+            .channel
+            .0
+            .send(WebGPURequest::MapReadAsync(
+                sender,
+                self.buffer.0,
+                self.device.0,
+                self.usage,
+                self.size,
+            ))
+            .is_err()
+        {
+            promise.reject_error(Error::Operation);
+            return promise;
+        }
+
+        // Step 6
+        promise
+    }
+
     /// https://gpuweb.github.io/gpuweb/#dom-gpuobjectbase-label
     fn GetLabel(&self) -> Option<DOMString> {
         self.label.borrow().clone()
@@ -138,3 +268,25 @@ impl GPUBufferMethods for GPUBuffer {
         *self.label.borrow_mut() = value;
     }
 }
+
+impl AsyncWGPUListener for GPUBuffer {
+    #[allow(unsafe_code)]
+    fn handle_response(&self, response: WebGPUResponse, promise: &Rc<Promise>) {
+        match response {
+            WebGPUResponse::MapReadAsync(bytes) => unsafe {
+                match ArrayBuffer::from(self.mapping.get()) {
+                    Ok(mut array_buffer) => {
+                        // Step 5.2
+                        array_buffer.update(&bytes);
+                        // Step 5.3
+                        *self.state.borrow_mut() = GPUBufferState::MappedForReading;
+                        // Step 5.4
+                        promise.resolve_native(&array_buffer);
+                    },
+                    _ => promise.reject_error(Error::Operation),
+                };
+            },
+            _ => promise.reject_error(Error::Operation),
+        }
+    }
+}

components/script/dom/gpudevice.rs

@@ -46,7 +46,7 @@ use webgpu::wgpu::binding_model::{
     ShaderStage,
 };
 use webgpu::wgpu::resource::{BufferDescriptor, BufferUsage};
-use webgpu::{WebGPU, WebGPUBuffer, WebGPUDevice, WebGPUQueue, WebGPURequest};
+use webgpu::{WebGPU, WebGPUDevice, WebGPUQueue, WebGPURequest};
 
 #[dom_struct]
 pub struct GPUDevice {
@@ -106,38 +106,6 @@ impl GPUDevice {
 }
 
 impl GPUDevice {
-    unsafe fn resolve_create_buffer_mapped(
-        &self,
-        cx: SafeJSContext,
-        gpu_buffer: WebGPUBuffer,
-        array_buffer: Vec<u8>,
-        descriptor: BufferDescriptor,
-        valid: bool,
-    ) -> Vec<JSVal> {
-        rooted!(in(*cx) let mut js_array_buffer = ptr::null_mut::<JSObject>());
-        let mut out = Vec::new();
-        assert!(ArrayBuffer::create(
-            *cx,
-            CreateWith::Slice(array_buffer.as_slice()),
-            js_array_buffer.handle_mut(),
-        )
-        .is_ok());
-
-        let buff = GPUBuffer::new(
-            &self.global(),
-            self.channel.clone(),
-            gpu_buffer,
-            self.device,
-            GPUBufferState::Mapped,
-            descriptor.size,
-            descriptor.usage.bits(),
-            valid,
-        );
-        out.push(ObjectValue(buff.reflector().get_jsobject().get()));
-        out.push(ObjectValue(js_array_buffer.get()));
-        out
-    }
-
     fn validate_buffer_descriptor(
         &self,
         descriptor: &GPUBufferDescriptor,
@@ -223,6 +191,7 @@ impl GPUDeviceMethods for GPUDevice {
             descriptor.size,
             descriptor.usage,
             valid,
+            RootedTraceableBox::new(Heap::default()),
         )
     }
 
@@ -245,11 +214,33 @@ impl GPUDeviceMethods for GPUDevice {
             ))
             .expect("Failed to create WebGPU buffer");
 
-        let (buffer, array_buffer) = receiver.recv().unwrap();
+        rooted!(in(*cx) let mut js_array_buffer = ptr::null_mut::<JSObject>());
-
         unsafe {
-            self.resolve_create_buffer_mapped(cx, buffer, array_buffer, wgpu_descriptor, valid)
+            assert!(ArrayBuffer::create(
+                *cx,
+                CreateWith::Length(descriptor.size as u32),
+                js_array_buffer.handle_mut(),
+            )
+            .is_ok());
         }
+
+        let buffer = receiver.recv().unwrap();
+        let buff = GPUBuffer::new(
+            &self.global(),
+            self.channel.clone(),
+            buffer,
+            self.device,
+            GPUBufferState::MappedForWriting,
+            wgpu_descriptor.size,
+            wgpu_descriptor.usage.bits(),
+            valid,
+            RootedTraceableBox::from_box(Heap::boxed(js_array_buffer.get())),
+        );
+
+        vec![
+            ObjectValue(buff.reflector().get_jsobject().get()),
+            ObjectValue(js_array_buffer.get()),
+        ]
     }
 
     /// https://gpuweb.github.io/gpuweb/#GPUDevice-createBindGroupLayout

components/script/dom/webidls/GPUBuffer.webidl

@@ -5,7 +5,7 @@
 // https://gpuweb.github.io/gpuweb/#gpubuffer
 [Exposed=(Window, DedicatedWorker), Serializable, Pref="dom.webgpu.enabled"]
 interface GPUBuffer {
-    // Promise<ArrayBuffer> mapReadAsync();
+    Promise<ArrayBuffer> mapReadAsync();
     // Promise<ArrayBuffer> mapWriteAsync();
     void unmap();
 

components/webgpu/lib.rs

@@ -7,7 +7,7 @@ extern crate log;
 #[macro_use]
 pub extern crate wgpu_core as wgpu;
 
-use ipc_channel::ipc::{self, IpcReceiver, IpcSender};
+use ipc_channel::ipc::{self, IpcReceiver, IpcSender, IpcSharedMemory};
 use malloc_size_of::{MallocSizeOf, MallocSizeOfOps};
 use servo_config::pref;
 use smallvec::SmallVec;
@@ -16,6 +16,7 @@ use smallvec::SmallVec;
 pub enum WebGPUResponse {
     RequestAdapter(String, WebGPUAdapter, WebGPU),
     RequestDevice(WebGPUDevice, WebGPUQueue, wgpu::instance::DeviceDescriptor),
+    MapReadAsync(IpcSharedMemory),
 }
 
 pub type WebGPUResponseResult = Result<WebGPUResponse, String>;
@@ -41,7 +42,7 @@ pub enum WebGPURequest {
         wgpu::resource::BufferDescriptor,
     ),
     CreateBufferMapped(
-        IpcSender<(WebGPUBuffer, Vec<u8>)>,
+        IpcSender<WebGPUBuffer>,
         WebGPUDevice,
         wgpu::id::BufferId,
         wgpu::resource::BufferDescriptor,
@@ -79,7 +80,14 @@ pub enum WebGPURequest {
         wgpu::id::ShaderModuleId,
         Vec<u32>,
     ),
-    UnmapBuffer(WebGPUBuffer),
+    MapReadAsync(
+        IpcSender<WebGPUResponseResult>,
+        wgpu::id::BufferId,
+        wgpu::id::DeviceId,
+        u32,
+        u64,
+    ),
+    UnmapBuffer(wgpu::id::DeviceId, WebGPUBuffer, Vec<u8>),
     DestroyBuffer(WebGPUBuffer),
     CreateCommandEncoder(
         IpcSender<WebGPUCommandEncoder>,
@@ -251,27 +259,26 @@ impl WGPU {
                 },
                 WebGPURequest::CreateBufferMapped(sender, device, id, descriptor) => {
                     let global = &self.global;
-                    let buffer_size = descriptor.size as usize;
+                    let (buffer_id, _arr_buff_ptr) = gfx_select!(id =>
-
-                    let (buffer_id, arr_buff_ptr) = gfx_select!(id =>
                         global.device_create_buffer_mapped(device.0, &descriptor, id));
                     let buffer = WebGPUBuffer(buffer_id);
 
-                    let mut array_buffer = Vec::with_capacity(buffer_size);
+                    if let Err(e) = sender.send(buffer) {
-                    unsafe {
-                        array_buffer.set_len(buffer_size);
-                        std::ptr::copy(arr_buff_ptr, array_buffer.as_mut_ptr(), buffer_size);
-                    };
-                    if let Err(e) = sender.send((buffer, array_buffer)) {
                         warn!(
                             "Failed to send response to WebGPURequest::CreateBufferMapped ({})",
                             e
                         )
                     }
                 },
-                WebGPURequest::UnmapBuffer(buffer) => {
+                WebGPURequest::UnmapBuffer(device_id, buffer, array_buffer) => {
                     let global = &self.global;
-                    gfx_select!(buffer.0 => global.buffer_unmap(buffer.0));
+
+                    gfx_select!(buffer.0 => global.device_set_buffer_sub_data(
+                        device_id,
+                        buffer.0,
+                        0,
+                        array_buffer.as_slice()
+                    ));
                 },
                 WebGPURequest::DestroyBuffer(buffer) => {
                     let global = &self.global;
@@ -412,6 +419,43 @@ impl WGPU {
                         )
                     }
                 },
+                WebGPURequest::MapReadAsync(sender, buffer_id, device_id, usage, size) => {
+                    let global = &self.global;
+                    let on_read = move |status: wgpu::resource::BufferMapAsyncStatus,
+                                        ptr: *const u8| {
+                        match status {
+                            wgpu::resource::BufferMapAsyncStatus::Success => {
+                                let array_buffer =
+                                    unsafe { std::slice::from_raw_parts(ptr, size as usize) };
+                                if let Err(e) = sender.send(Ok(WebGPUResponse::MapReadAsync(
+                                    IpcSharedMemory::from_bytes(array_buffer),
+                                ))) {
+                                    warn!(
+                                        "Failed to send response to WebGPURequest::MapReadAsync ({})",
+                                        e
+                                    )
+                                }
+                            },
+                            _ => {
+                                if let Err(e) = sender
+                                    .send(Err("MapReadAsync: Failed to map buffer".to_owned()))
+                                {
+                                    warn!(
+                                        "Failed to send response to WebGPURequest::MapReadAsync ({})",
+                                        e
+                                    )
+                                }
+                            },
+                        }
+                    };
+                    gfx_select!(buffer_id => global.buffer_map_async(
+                        buffer_id,
+                        wgpu::resource::BufferUsage::from_bits(usage).unwrap(),
+                        0..size,
+                        wgpu::resource::BufferMapOperation::Read(Box::new(on_read))
+                    ));
+                    gfx_select!(device_id => global.device_poll(device_id, true));
+                },
                 WebGPURequest::Submit(queue_id, command_buffer_ids) => {
                     let global = &self.global;
                     let _ = gfx_select!(queue_id => global.queue_submit(