Implement GPUBuffer.getMappedRange()

1 files changed, 117 insertions, 68 deletions

diff --git a/components/script/dom/gpubuffer.rs b/components/script/dom/gpubuffer.rs
index 474c30dba0f..e238057eb9a 100644
--- a/components/script/dom/gpubuffer.rs
+++ b/components/script/dom/gpubuffer.rs
@@ -2,32 +2,35 @@
  * 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;
+use crate::dom::bindings::cell::{DomRefCell, RefCell};
 use crate::dom::bindings::codegen::Bindings::GPUBufferBinding::{GPUBufferMethods, GPUSize64};
 use crate::dom::bindings::codegen::Bindings::GPUMapModeBinding::GPUMapModeConstants;
-use crate::dom::bindings::error::Error;
+use crate::dom::bindings::error::{Error, Fallible};
 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 crate::script_runtime::JSContext;
 use dom_struct::dom_struct;
+use js::jsapi::DetachArrayBuffer;
+use js::jsapi::NewExternalArrayBuffer;
 use js::jsapi::{Heap, JSObject};
-use js::jsval::UndefinedValue;
-use js::rust::jsapi_wrapped::{DetachArrayBuffer, IsPromiseObject, RejectPromise};
-use js::typedarray::{ArrayBuffer, CreateWith};
 use std::cell::Cell;
+use std::ffi::c_void;
 use std::ops::Range;
-use std::ptr;
+use std::ptr::NonNull;
 use std::rc::Rc;
 use webgpu::{
     wgpu::device::HostMap, WebGPU, WebGPUBuffer, WebGPUDevice, WebGPURequest, WebGPUResponse,
 };
 
+const RANGE_OFFSET_ALIGN_MASK: u64 = 8;
+const RANGE_SIZE_ALIGN_MASK: u64 = 4;
+
 // https://gpuweb.github.io/gpuweb/#buffer-state
 #[derive(Clone, Copy, MallocSizeOf, PartialEq)]
 pub enum GPUBufferState {
@@ -49,8 +52,13 @@ pub struct GPUBuffer {
     device: WebGPUDevice,
     valid: Cell<bool>,
     #[ignore_malloc_size_of = "defined in mozjs"]
-    mapping: RootedTraceableBox<Heap<*mut JSObject>>,
-    mapping_range: DomRefCell<Range<u64>>,
+    mapping: Rc<RefCell<Option<Vec<u8>>>>,
+    mapping_range: DomRefCell<Option<Range<u64>>>,
+    mapped_ranges: DomRefCell<Option<Vec<Range<u64>>>>,
+    #[ignore_malloc_size_of = "defined in mozjs"]
+    js_buffers: DomRefCell<Option<Vec<Box<Heap<*mut JSObject>>>>>,
+    #[ignore_malloc_size_of = "defined in mozjs"]
+    map_promise: DomRefCell<Option<Rc<Promise>>>,
     size: GPUSize64,
     map_mode: Cell<Option<u32>>,
 }
@@ -63,7 +71,7 @@ impl GPUBuffer {
         state: GPUBufferState,
         size: GPUSize64,
         valid: bool,
-        mapping: RootedTraceableBox<Heap<*mut JSObject>>,
+        mapping: Rc<RefCell<Option<Vec<u8>>>>,
         mapping_range: Range<u64>,
     ) -> Self {
         Self {
@@ -75,8 +83,11 @@ impl GPUBuffer {
             device,
             buffer,
             mapping,
+            mapped_ranges: DomRefCell::new(None),
+            js_buffers: DomRefCell::new(None),
+            map_promise: DomRefCell::new(None),
             size,
-            mapping_range: DomRefCell::new(mapping_range),
+            mapping_range: DomRefCell::new(Some(mapping_range)),
             map_mode: Cell::new(None),
         }
     }
@@ -90,7 +101,7 @@ impl GPUBuffer {
         state: GPUBufferState,
         size: GPUSize64,
         valid: bool,
-        mapping: RootedTraceableBox<Heap<*mut JSObject>>,
+        mapping: Rc<RefCell<Option<Vec<u8>>>>,
         mapping_range: Range<u64>,
     ) -> DomRoot<Self> {
         reflect_dom_object(
@@ -142,51 +153,34 @@ impl GPUBufferMethods for GPUBuffer {
             },
             // Step 3
             GPUBufferState::Mapped | GPUBufferState::MappedAtCreation => {
-                match ArrayBuffer::from(self.mapping.get()) {
-                    Ok(array_buffer) => {
-                        // Step 3.2
-                        self.channel
-                            .0
-                            .send(WebGPURequest::UnmapBuffer {
-                                buffer_id: self.id().0,
-                                array_buffer: array_buffer.to_vec(),
-                                is_map_read: self.map_mode.get() == Some(GPUMapModeConstants::READ),
-                                offset: self.mapping_range.borrow().start,
-                                size: self.mapping_range.borrow().end -
-                                    self.mapping_range.borrow().start,
-                            })
-                            .unwrap();
-                        // Step 3.3
-                        unsafe {
-                            DetachArrayBuffer(*cx, self.mapping.handle());
-                        }
-                    },
-                    Err(_) => {
-                        warn!(
-                            "Could not find ArrayBuffer of Mapped buffer ({:?})",
-                            self.buffer.0
-                        );
-                    },
-                };
+                if let Err(e) = self.channel.0.send(WebGPURequest::UnmapBuffer {
+                    buffer_id: self.id().0,
+                    array_buffer: self.mapping.borrow().as_ref().unwrap().clone(),
+                    is_map_read: self.map_mode.get() == Some(GPUMapModeConstants::READ),
+                    offset: self.mapping_range.borrow().as_ref().unwrap().start,
+                    size: self.mapping_range.borrow().as_ref().unwrap().end -
+                        self.mapping_range.borrow().as_ref().unwrap().start,
+                }) {
+                    warn!("Failed to send Buffer unmap ({:?}) ({})", self.buffer.0, e);
+                }
+                // Step 3.3
+                let mut bufs = self.js_buffers.borrow_mut().take().unwrap();
+                bufs.drain(..).for_each(|obj| unsafe {
+                    DetachArrayBuffer(*cx, obj.handle());
+                });
+                *self.mapped_ranges.borrow_mut() = None;
+                *self.mapping.borrow_mut() = None;
             },
             // 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");
-                }
+            GPUBufferState::MappingPending => {
+                let promise = self.map_promise.borrow_mut().take().unwrap();
+                promise.reject_error(Error::Operation);
             },
         };
-        // Step 3.3
-        self.mapping.set(ptr::null_mut());
         // Step 4
         self.state.set(GPUBufferState::Unmapped);
         self.map_mode.set(None);
-        *self.mapping_range.borrow_mut() = 0..0;
+        *self.mapping_range.borrow_mut() = None;
     }
 
     /// https://gpuweb.github.io/gpuweb/#dom-gpubuffer-destroy
@@ -213,7 +207,13 @@ impl GPUBufferMethods for GPUBuffer {
 
     #[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> {
+    fn MapAsync(
+        &self,
+        mode: u32,
+        offset: GPUSize64,
+        size: GPUSize64,
+        comp: InRealm,
+    ) -> Rc<Promise> {
         let promise = Promise::new_in_current_realm(&self.global(), comp);
         let map_range = if size == 0 {
             offset..self.size
@@ -237,7 +237,6 @@ impl GPUBufferMethods for GPUBuffer {
             promise.reject_error(Error::Abort);
             return promise;
         }
-        self.mapping.set(*promise.promise_obj());
 
         let sender = response_async(&promise, self);
         if let Err(e) = self.channel.0.send(WebGPURequest::BufferMapAsync {
@@ -256,10 +255,72 @@ impl GPUBufferMethods for GPUBuffer {
 
         self.state.set(GPUBufferState::MappingPending);
         self.map_mode.set(Some(mode));
-        *self.mapping_range.borrow_mut() = map_range;
+        *self.mapping_range.borrow_mut() = Some(map_range);
+        *self.map_promise.borrow_mut() = Some(promise.clone());
         promise
     }
 
+    /// https://gpuweb.github.io/gpuweb/#dom-gpubuffer-getmappedrange
+    #[allow(unsafe_code)]
+    fn GetMappedRange(
+        &self,
+        cx: JSContext,
+        offset: GPUSize64,
+        size: GPUSize64,
+    ) -> Fallible<NonNull<JSObject>> {
+        if self.mapped_ranges.borrow().is_none() {
+            *self.mapped_ranges.borrow_mut() = Some(Vec::new());
+        }
+        if self.js_buffers.borrow().is_none() {
+            *self.js_buffers.borrow_mut() = Some(Vec::new());
+        }
+        let act_size = if size == 0 { self.size - offset } else { size };
+        let mut valid = match self.state.get() {
+            GPUBufferState::Mapped | GPUBufferState::MappedAtCreation => true,
+            _ => false,
+        };
+        valid &= offset % RANGE_OFFSET_ALIGN_MASK == 0 &&
+            act_size % RANGE_SIZE_ALIGN_MASK == 0 &&
+            offset >= self.mapping_range.borrow().as_ref().unwrap().start &&
+            offset + act_size <= self.mapping_range.borrow().as_ref().unwrap().end;
+        valid &= self
+            .mapped_ranges
+            .borrow()
+            .as_ref()
+            .unwrap()
+            .iter()
+            .all(|range| range.start > offset + act_size || range.end < offset);
+        if !valid {
+            return Err(Error::Operation);
+        }
+
+        unsafe extern "C" fn free_func(_contents: *mut c_void, free_user_data: *mut c_void) {
+            let _ = Rc::from_raw(free_user_data as _);
+        }
+
+        let array_buffer = unsafe {
+            NewExternalArrayBuffer(
+                *cx,
+                act_size as usize,
+                self.mapping.borrow_mut().as_mut().unwrap()
+                    [offset as usize..(offset + act_size) as usize]
+                    .as_mut_ptr() as _,
+                Some(free_func),
+                Rc::into_raw(self.mapping.clone()) as _,
+            )
+        };
+        self.mapped_ranges
+            .borrow_mut()
+            .as_mut()
+            .map(|v| v.push(offset..offset + act_size));
+        self.js_buffers
+            .borrow_mut()
+            .as_mut()
+            .map(|a| a.push(Heap::boxed(array_buffer)));
+
+        Ok(NonNull::new(array_buffer).unwrap())
+    }
+
     /// https://gpuweb.github.io/gpuweb/#dom-gpuobjectbase-label
     fn GetLabel(&self) -> Option<DOMString> {
         self.label.borrow().clone()
@@ -276,21 +337,8 @@ impl AsyncWGPUListener for GPUBuffer {
     fn handle_response(&self, response: WebGPUResponse, promise: &Rc<Promise>) {
         match response {
             WebGPUResponse::BufferMapAsync(bytes) => {
-                let cx = self.global().get_cx();
-                rooted!(in(*cx) let mut array_buffer = ptr::null_mut::<JSObject>());
-                match unsafe {
-                    ArrayBuffer::create(*cx, CreateWith::Slice(&bytes), array_buffer.handle_mut())
-                } {
-                    Ok(_) => promise.resolve_native(&()),
-                    Err(()) => {
-                        warn!(
-                            "Failed to create ArrayBuffer for buffer({:?})",
-                            self.buffer.0
-                        );
-                        promise.reject_error(Error::Operation);
-                    },
-                }
-                self.mapping.set(array_buffer.get());
+                *self.mapping.borrow_mut() = Some(bytes);
+                promise.resolve_native(&());
                 self.state.set(GPUBufferState::Mapped);
             },
             _ => {
@@ -298,6 +346,7 @@ impl AsyncWGPUListener for GPUBuffer {
                 promise.reject_error(Error::Operation);
             },
         }
+        *self.map_promise.borrow_mut() = None;
         if let Err(e) = self
             .channel
             .0