Implement GPUBuffer.mapAsync and update wgpu-core

1 files changed, 160 insertions, 56 deletions

diff --git a/components/script/dom/gpubuffer.rs b/components/script/dom/gpubuffer.rs
index a23a30a4b11..016f5d8bd1d 100644
--- a/components/script/dom/gpubuffer.rs
+++ b/components/script/dom/gpubuffer.rs
@@ -2,8 +2,9 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 
-use crate::dom::bindings::cell::{DomRefCell, Ref};
-use crate::dom::bindings::codegen::Bindings::GPUBufferBinding::{GPUBufferMethods, GPUSize64};
+use crate::dom::bindings::cell::DomRefCell;
+use crate::dom::bindings::codegen::Bindings::GPUBufferBinding::GPUBufferMethods;
+use crate::dom::bindings::codegen::Bindings::GPUMapModeBinding::GPUMapModeConstants;
 use crate::dom::bindings::error::Error;
 use crate::dom::bindings::reflector::DomObject;
 use crate::dom::bindings::reflector::{reflect_dom_object, Reflector};
@@ -11,22 +12,28 @@ 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::typedarray::ArrayBuffer;
+use js::rust::MutableHandle;
+use js::typedarray::{ArrayBuffer, CreateWith};
 use std::cell::Cell;
 use std::ptr;
-use webgpu::{WebGPU, WebGPUBuffer, WebGPUDevice, WebGPURequest};
+use std::rc::Rc;
+use webgpu::{
+    wgpu::device::HostMap, WebGPU, WebGPUBuffer, WebGPUDevice, WebGPURequest, WebGPUResponse,
+};
 
 // https://gpuweb.github.io/gpuweb/#buffer-state
-#[derive(Clone, MallocSizeOf)]
+#[derive(Clone, Copy, MallocSizeOf, PartialEq)]
 pub enum GPUBufferState {
-    MappedForReading,
-    MappedForWriting,
-    MappedPendingForReading,
-    MappedPendingForWriting,
+    Mapped,
+    MappedAtCreation,
+    MappingPending,
     Unmapped,
     Destroyed,
 }
@@ -37,14 +44,15 @@ pub struct GPUBuffer {
     #[ignore_malloc_size_of = "defined in webgpu"]
     channel: WebGPU,
     label: DomRefCell<Option<DOMString>>,
-    size: GPUSize64,
-    usage: u32,
-    state: DomRefCell<GPUBufferState>,
+    state: Cell<GPUBufferState>,
     buffer: WebGPUBuffer,
     device: WebGPUDevice,
     valid: Cell<bool>,
     #[ignore_malloc_size_of = "defined in mozjs"]
     mapping: RootedTraceableBox<Heap<*mut JSObject>>,
+    mapping_range: (u64, u64),
+    size: u64,
+    map_mode: Cell<Option<u32>>,
 }
 
 impl GPUBuffer {
@@ -53,22 +61,23 @@ impl GPUBuffer {
         buffer: WebGPUBuffer,
         device: WebGPUDevice,
         state: GPUBufferState,
-        size: GPUSize64,
-        usage: u32,
+        size: u64,
         valid: bool,
         mapping: RootedTraceableBox<Heap<*mut JSObject>>,
+        mapping_range: (u64, u64),
     ) -> Self {
         Self {
             reflector_: Reflector::new(),
             channel,
             label: DomRefCell::new(None),
-            state: DomRefCell::new(state),
-            size: size,
-            usage: usage,
+            state: Cell::new(state),
             valid: Cell::new(valid),
             device,
             buffer,
             mapping,
+            size,
+            mapping_range,
+            map_mode: Cell::new(None),
         }
     }
 
@@ -79,14 +88,21 @@ impl GPUBuffer {
         buffer: WebGPUBuffer,
         device: WebGPUDevice,
         state: GPUBufferState,
-        size: GPUSize64,
-        usage: u32,
+        size: u64,
         valid: bool,
         mapping: RootedTraceableBox<Heap<*mut JSObject>>,
+        mapping_range: (u64, u64),
     ) -> DomRoot<Self> {
         reflect_dom_object(
             Box::new(GPUBuffer::new_inherited(
-                channel, buffer, device, state, size, usage, valid, mapping,
+                channel,
+                buffer,
+                device,
+                state,
+                size,
+                valid,
+                mapping,
+                mapping_range,
             )),
             global,
         )
@@ -98,16 +114,8 @@ impl GPUBuffer {
         self.buffer
     }
 
-    pub fn size(&self) -> GPUSize64 {
-        self.size
-    }
-
-    pub fn usage(&self) -> u32 {
-        self.usage
-    }
-
-    pub fn state(&self) -> Ref<GPUBufferState> {
-        self.state.borrow()
+    pub fn state(&self) -> GPUBufferState {
+        self.state.get()
     }
 
     pub fn is_valid(&self) -> bool {
@@ -127,54 +135,64 @@ impl GPUBufferMethods for GPUBuffer {
     fn Unmap(&self) {
         let cx = self.global().get_cx();
         // Step 1
-        match *self.state.borrow() {
+        match self.state.get() {
             GPUBufferState::Unmapped | GPUBufferState::Destroyed => {
                 // TODO: Record validation error on the current scope
                 return;
             },
-            GPUBufferState::MappedForWriting => {
-                // Step 3.1
+            // Step 3
+            GPUBufferState::Mapped | GPUBufferState::MappedAtCreation => {
                 match ArrayBuffer::from(self.mapping.get()) {
                     Ok(array_buffer) => {
-                        self.channel
-                            .0
-                            .send(WebGPURequest::UnmapBuffer {
-                                device_id: self.device.0,
-                                buffer_id: self.id().0,
-                                array_buffer: array_buffer.to_vec(),
-                            })
-                            .unwrap();
                         // Step 3.2
+                        if Some(GPUMapModeConstants::READ) != self.map_mode.get() {
+                            self.channel
+                                .0
+                                .send(WebGPURequest::UnmapBuffer {
+                                    device_id: self.device.0,
+                                    buffer_id: self.id().0,
+                                    array_buffer: array_buffer.to_vec(),
+                                    mapped_at_creation: self.map_mode.get() == None,
+                                })
+                                .unwrap();
+                        }
+                        // Step 3.3
                         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());
-                            };
-                        }
+                    Err(_) => {
+                        warn!(
+                            "Could not find ArrayBuffer of Mapped buffer ({:?})",
+                            self.buffer.0
+                        );
                     },
                 };
             },
-            _ => {},
+            // Step 2
+            GPUBufferState::MappingPending => unsafe {
+                if IsPromiseObject(self.mapping.handle()) {
+                    let err = Error::Operation;
+                    rooted!(in(*cx) let mut undef = UndefinedValue());
+                    err.to_jsval(*cx, &self.global(), undef.handle_mut());
+                    RejectPromise(*cx, self.mapping.handle(), undef.handle());
+                } else {
+                    warn!("No promise object for pending mapping found");
+                }
+            },
         };
         // Step 3.3
         self.mapping.set(ptr::null_mut());
         // Step 4
-        *self.state.borrow_mut() = GPUBufferState::Unmapped;
+        self.state.set(GPUBufferState::Unmapped);
+        self.map_mode.set(None);
     }
 
     /// https://gpuweb.github.io/gpuweb/#dom-gpubuffer-destroy
     fn Destroy(&self) {
-        let state = self.state.borrow().clone();
+        let state = self.state.get();
         match state {
-            GPUBufferState::MappedForReading | GPUBufferState::MappedForWriting => {
+            GPUBufferState::Mapped | GPUBufferState::MappedAtCreation => {
                 self.Unmap();
             },
             _ => {},
@@ -189,7 +207,56 @@ impl GPUBufferMethods for GPUBuffer {
                 self.buffer.0, e
             );
         };
-        *self.state.borrow_mut() = GPUBufferState::Destroyed;
+        self.state.set(GPUBufferState::Destroyed);
+    }
+
+    #[allow(unsafe_code)]
+    /// https://gpuweb.github.io/gpuweb/#dom-gpubuffer-mapasync-offset-size
+    fn MapAsync(&self, mode: u32, offset: u64, size: u64, comp: InRealm) -> Rc<Promise> {
+        let promise = Promise::new_in_current_realm(&self.global(), comp);
+        let map_range = if size == 0 {
+            offset..self.size
+        } else {
+            offset..size
+        };
+        let host_map = match mode {
+            GPUMapModeConstants::READ => {
+                self.map_mode.set(Some(mode));
+                HostMap::Read
+            },
+            GPUMapModeConstants::WRITE => {
+                self.map_mode.set(Some(mode));
+                HostMap::Write
+            },
+            _ => {
+                promise.reject_error(Error::Abort);
+                return promise;
+            },
+        };
+        if self.state.get() != GPUBufferState::Unmapped {
+            promise.reject_error(Error::Abort);
+            return promise;
+        }
+        self.mapping.set(*promise.promise_obj());
+        self.state.set(GPUBufferState::MappingPending);
+
+        let sender = response_async(&promise, self);
+
+        if let Err(e) = self.channel.0.send(WebGPURequest::BufferMapAsync {
+            sender,
+            buffer_id: self.buffer.0,
+            host_map,
+            map_range,
+        }) {
+            warn!(
+                "Failed to send BufferMapAsync ({:?}) ({})",
+                self.buffer.0, e
+            );
+            promise.reject_error(Error::Operation);
+            return promise;
+        }
+
+        promise
     }
 
     /// https://gpuweb.github.io/gpuweb/#dom-gpuobjectbase-label
@@ -202,3 +269,40 @@ 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::BufferMapAsync(bytes) => {
+                if unsafe {
+                    ArrayBuffer::create(
+                        *self.global().get_cx(),
+                        CreateWith::Slice(&bytes),
+                        MutableHandle::from_raw(self.mapping.handle_mut()),
+                    )
+                }
+                .is_err()
+                {
+                    promise.reject_error(Error::Operation);
+                }
+                self.state.set(GPUBufferState::Mapped);
+                promise.resolve_native(&());
+            },
+            _ => {
+                warn!("Wrong WebGPUResponse received");
+                promise.reject_error(Error::Operation);
+            },
+        }
+        if let Err(e) = self
+            .channel
+            .0
+            .send(WebGPURequest::BufferMapComplete(self.buffer.0))
+        {
+            warn!(
+                "Failed to send BufferMapComplete({:?}) ({})",
+                self.buffer.0, e
+            );
+        }
+    }
+}