refactor(webgpu): use op interface idiomatically (#11835)

1 files changed, 65 insertions, 60 deletions

diff --git a/ext/webgpu/bundle.rs b/ext/webgpu/bundle.rs
index 3a228ba62..df91104f5 100644
--- a/ext/webgpu/bundle.rs
+++ b/ext/webgpu/bundle.rs
@@ -1,6 +1,5 @@
 // Copyright 2018-2021 the Deno authors. All rights reserved. MIT license.
 
-use deno_core::error::null_opbuf;
 use deno_core::error::AnyError;
 use deno_core::ResourceId;
 use deno_core::ZeroCopyBuf;
@@ -8,10 +7,13 @@ use deno_core::{OpState, Resource};
 use serde::Deserialize;
 use std::borrow::Cow;
 use std::cell::RefCell;
+use std::convert::TryInto;
 use std::rc::Rc;
 
+use crate::pipeline::GpuIndexFormat;
+use crate::texture::GpuTextureFormat;
+
 use super::error::WebGpuResult;
-use super::texture::serialize_texture_format;
 
 struct WebGpuRenderBundleEncoder(
   RefCell<wgpu_core::command::RenderBundleEncoder>,
@@ -34,9 +36,9 @@ impl Resource for WebGpuRenderBundle {
 pub struct CreateRenderBundleEncoderArgs {
   device_rid: ResourceId,
   label: Option<String>,
-  color_formats: Vec<String>,
-  depth_stencil_format: Option<String>,
-  sample_count: Option<u32>,
+  color_formats: Vec<GpuTextureFormat>,
+  depth_stencil_format: Option<GpuTextureFormat>,
+  sample_count: u32,
   depth_read_only: bool,
   stencil_read_only: bool,
 }
@@ -53,24 +55,25 @@ pub fn op_webgpu_create_render_bundle_encoder(
 
   let mut color_formats = vec![];
 
-  for format in &args.color_formats {
-    color_formats.push(serialize_texture_format(format)?);
+  for format in args.color_formats {
+    color_formats.push(format.try_into()?);
   }
 
+  let depth_stencil = if let Some(format) = args.depth_stencil_format {
+    Some(wgpu_types::RenderBundleDepthStencil {
+      format: format.try_into()?,
+      depth_read_only: args.depth_read_only,
+      stencil_read_only: args.stencil_read_only,
+    })
+  } else {
+    None
+  };
+
   let descriptor = wgpu_core::command::RenderBundleEncoderDescriptor {
     label: args.label.map(Cow::from),
     color_formats: Cow::from(color_formats),
-    sample_count: args.sample_count.unwrap_or(1),
-    depth_stencil: if let Some(depth_stencil_format) = args.depth_stencil_format
-    {
-      Some(wgpu_types::RenderBundleDepthStencil {
-        format: serialize_texture_format(&depth_stencil_format)?,
-        depth_read_only: args.depth_read_only,
-        stencil_read_only: args.stencil_read_only,
-      })
-    } else {
-      None
-    },
+    sample_count: args.sample_count,
+    depth_stencil,
   };
 
   let res =
@@ -129,8 +132,8 @@ pub fn op_webgpu_render_bundle_encoder_finish(
 pub struct RenderBundleEncoderSetBindGroupArgs {
   render_bundle_encoder_rid: ResourceId,
   index: u32,
-  bind_group: u32,
-  dynamic_offsets_data: Option<Vec<u32>>,
+  bind_group: ResourceId,
+  dynamic_offsets_data: ZeroCopyBuf,
   dynamic_offsets_data_start: usize,
   dynamic_offsets_data_length: usize,
 }
@@ -138,7 +141,7 @@ pub struct RenderBundleEncoderSetBindGroupArgs {
 pub fn op_webgpu_render_bundle_encoder_set_bind_group(
   state: &mut OpState,
   args: RenderBundleEncoderSetBindGroupArgs,
-  zero_copy: Option<ZeroCopyBuf>,
+  _: (),
 ) -> Result<WebGpuResult, AnyError> {
   let bind_group_resource =
     state
@@ -149,37 +152,35 @@ pub fn op_webgpu_render_bundle_encoder_set_bind_group(
       .resource_table
       .get::<WebGpuRenderBundleEncoder>(args.render_bundle_encoder_rid)?;
 
-  // I know this might look like it can be easily deduplicated, but it can not
-  // be due to the lifetime of the args.dynamic_offsets_data slice. Because we
-  // need to use a raw pointer here the slice can be freed before the pointer
-  // is used in wgpu_render_pass_set_bind_group. See
-  // https://matrix.to/#/!XFRnMvAfptAHthwBCx:matrix.org/$HgrlhD-Me1DwsGb8UdMu2Hqubgks8s7ILwWRwigOUAg
-  match args.dynamic_offsets_data {
-    Some(data) => unsafe {
-      wgpu_core::command::bundle_ffi::wgpu_render_bundle_set_bind_group(
-        &mut render_bundle_encoder_resource.0.borrow_mut(),
-        args.index,
-        bind_group_resource.0,
-        data.as_slice().as_ptr(),
-        args.dynamic_offsets_data_length,
-      );
-    },
-    None => {
-      let zero_copy = zero_copy.ok_or_else(null_opbuf)?;
-      let (prefix, data, suffix) = unsafe { zero_copy.align_to::<u32>() };
-      assert!(prefix.is_empty());
-      assert!(suffix.is_empty());
-      unsafe {
-        wgpu_core::command::bundle_ffi::wgpu_render_bundle_set_bind_group(
-          &mut render_bundle_encoder_resource.0.borrow_mut(),
-          args.index,
-          bind_group_resource.0,
-          data[args.dynamic_offsets_data_start..].as_ptr(),
-          args.dynamic_offsets_data_length,
-        );
-      }
-    }
-  };
+  // Align the data
+  assert!(args.dynamic_offsets_data.len() % std::mem::size_of::<u32>() == 0);
+  // SAFETY: A u8 to u32 cast is safe because we asserted that the length is a
+  // multiple of 4.
+  let (prefix, dynamic_offsets_data, suffix) =
+    unsafe { args.dynamic_offsets_data.align_to::<u32>() };
+  assert!(prefix.is_empty());
+  assert!(suffix.is_empty());
+
+  let start = args.dynamic_offsets_data_start;
+  let len = args.dynamic_offsets_data_length;
+
+  // Assert that length and start are both in bounds
+  assert!(start <= dynamic_offsets_data.len());
+  assert!(len <= dynamic_offsets_data.len() - start);
+
+  let dynamic_offsets_data: &[u32] = &dynamic_offsets_data[start..start + len];
+
+  // SAFETY: the raw pointer and length are of the same slice, and that slice
+  // lives longer than the below function invocation.
+  unsafe {
+    wgpu_core::command::bundle_ffi::wgpu_render_bundle_set_bind_group(
+      &mut render_bundle_encoder_resource.0.borrow_mut(),
+      args.index,
+      bind_group_resource.0,
+      dynamic_offsets_data.as_ptr(),
+      dynamic_offsets_data.len(),
+    );
+  }
 
   Ok(WebGpuResult::empty())
 }
@@ -201,8 +202,10 @@ pub fn op_webgpu_render_bundle_encoder_push_debug_group(
       .resource_table
       .get::<WebGpuRenderBundleEncoder>(args.render_bundle_encoder_rid)?;
 
+  let label = std::ffi::CString::new(args.group_label).unwrap();
+  // SAFETY: the string the raw pointer points to lives longer than the below
+  // function invocation.
   unsafe {
-    let label = std::ffi::CString::new(args.group_label).unwrap();
     wgpu_core::command::bundle_ffi::wgpu_render_bundle_push_debug_group(
       &mut render_bundle_encoder_resource.0.borrow_mut(),
       label.as_ptr(),
@@ -252,8 +255,10 @@ pub fn op_webgpu_render_bundle_encoder_insert_debug_marker(
       .resource_table
       .get::<WebGpuRenderBundleEncoder>(args.render_bundle_encoder_rid)?;
 
+  let label = std::ffi::CString::new(args.marker_label).unwrap();
+  // SAFETY: the string the raw pointer points to lives longer than the below
+  // function invocation.
   unsafe {
-    let label = std::ffi::CString::new(args.marker_label).unwrap();
     wgpu_core::command::bundle_ffi::wgpu_render_bundle_insert_debug_marker(
       &mut render_bundle_encoder_resource.0.borrow_mut(),
       label.as_ptr(),
@@ -267,7 +272,7 @@ pub fn op_webgpu_render_bundle_encoder_insert_debug_marker(
 #[serde(rename_all = "camelCase")]
 pub struct RenderBundleEncoderSetPipelineArgs {
   render_bundle_encoder_rid: ResourceId,
-  pipeline: u32,
+  pipeline: ResourceId,
 }
 
 pub fn op_webgpu_render_bundle_encoder_set_pipeline(
@@ -296,8 +301,8 @@ pub fn op_webgpu_render_bundle_encoder_set_pipeline(
 #[serde(rename_all = "camelCase")]
 pub struct RenderBundleEncoderSetIndexBufferArgs {
   render_bundle_encoder_rid: ResourceId,
-  buffer: u32,
-  index_format: String,
+  buffer: ResourceId,
+  index_format: GpuIndexFormat,
   offset: u64,
   size: u64,
 }
@@ -320,7 +325,7 @@ pub fn op_webgpu_render_bundle_encoder_set_index_buffer(
     .borrow_mut()
     .set_index_buffer(
       buffer_resource.0,
-      super::pipeline::serialize_index_format(args.index_format),
+      args.index_format.into(),
       args.offset,
       std::num::NonZeroU64::new(args.size),
     );
@@ -333,7 +338,7 @@ pub fn op_webgpu_render_bundle_encoder_set_index_buffer(
 pub struct RenderBundleEncoderSetVertexBufferArgs {
   render_bundle_encoder_rid: ResourceId,
   slot: u32,
-  buffer: u32,
+  buffer: ResourceId,
   offset: u64,
   size: u64,
 }
@@ -430,7 +435,7 @@ pub fn op_webgpu_render_bundle_encoder_draw_indexed(
 #[serde(rename_all = "camelCase")]
 pub struct RenderBundleEncoderDrawIndirectArgs {
   render_bundle_encoder_rid: ResourceId,
-  indirect_buffer: u32,
+  indirect_buffer: ResourceId,
   indirect_offset: u64,
 }