feat(ops): fast calls for Wasm (denoland#16776)

This PR introduces Wasm ops. These calls are optimized for entry from Wasm land. The `#[op(wasm)]` attribute is opt-in. Last parameter `Option<&mut [u8]>` is the memory slice of the Wasm module *when entered from a Fast API call*. Otherwise, the user is expected to implement logic to obtain the memory if `None` ```rust #[op(wasm)] pub fn op_args_get( offset: i32, buffer_offset: i32, memory: Option<&mut [u8]>, ) { // ... } ```
DavidDeSimone · Nov 27, 2022 · ca66978 · ca66978
1 parent 9ffc6ac
commit ca66978
Show file tree

Hide file tree

Showing 11 changed files with 291 additions and 8 deletions.
diff --git a/core/examples/wasm.js b/core/examples/wasm.js
@@ -0,0 +1,28 @@
+// Copyright 2018-2022 the Deno authors. All rights reserved. MIT license.
+
+// asc wasm.ts --exportStart --initialMemory 6400 -O -o wasm.wasm
+// deno-fmt-ignore
+const bytes = new Uint8Array([
+    0,  97, 115, 109,   1,   0,   0,   0,   1,  4,   1, 96,   0,   0,   2,
+   15,   1,   3, 111, 112, 115,   7, 111, 112, 95, 119, 97, 115, 109,   0,
+    0,   3,   3,   2,   0,   0,   5,   4,   1,  0, 128, 50,   7,  36,   4,
+    7, 111, 112,  95, 119,  97, 115, 109,   0,  0,   4, 99,  97, 108, 108,
+    0,   1,   6, 109, 101, 109, 111, 114, 121,  2,   0,  6,  95, 115, 116,
+   97, 114, 116,   0,   2,  10,  10,   2,   4,  0,  16,  0,  11,   3,   0,
+    1,  11
+ ]);
+
+const { ops } = Deno.core;
+
+const module = new WebAssembly.Module(bytes);
+const instance = new WebAssembly.Instance(module, { ops });
+ops.op_set_wasm_mem(instance.exports.memory);
+
+instance.exports.call();
+
+const memory = instance.exports.memory;
+const view = new Uint8Array(memory.buffer);
+
+if (view[0] !== 69) {
+  throw new Error("Expected first byte to be 69");
+}
diff --git a/core/examples/wasm.rs b/core/examples/wasm.rs
@@ -0,0 +1,67 @@
+// Copyright 2018-2022 the Deno authors. All rights reserved. MIT license.
+
+use deno_core::op;
+use deno_core::Extension;
+use deno_core::JsRuntime;
+use deno_core::RuntimeOptions;
+use std::mem::transmute;
+use std::ptr::NonNull;
+
+// This is a hack to make the `#[op]` macro work with
+// deno_core examples.
+// You can remove this:
+
+use deno_core::*;
+
+struct WasmMemory(NonNull<v8::WasmMemoryObject>);
+
+fn wasm_memory_unchecked(state: &mut OpState) -> &mut [u8] {
+  let WasmMemory(global) = state.borrow::<WasmMemory>();
+  // SAFETY: `v8::Local` is always non-null pointer; the `HandleScope` is
+  // already on the stack, but we don't have access to it.
+  let memory_object = unsafe {
+    transmute::<NonNull<v8::WasmMemoryObject>, v8::Local<v8::WasmMemoryObject>>(
+      *global,
+    )
+  };
+  let backing_store = memory_object.buffer().get_backing_store();
+  let ptr = backing_store.data().unwrap().as_ptr() as *mut u8;
+  let len = backing_store.byte_length();
+  // SAFETY: `ptr` is a valid pointer to `len` bytes.
+  unsafe { std::slice::from_raw_parts_mut(ptr, len) }
+}
+
+#[op(wasm)]
+fn op_wasm(state: &mut OpState, memory: Option<&mut [u8]>) {
+  let memory = memory.unwrap_or_else(|| wasm_memory_unchecked(state));
+  memory[0] = 69;
+}
+
+#[op(v8)]
+fn op_set_wasm_mem(
+  scope: &mut v8::HandleScope,
+  state: &mut OpState,
+  memory: serde_v8::Value,
+) {
+  let memory =
+    v8::Local::<v8::WasmMemoryObject>::try_from(memory.v8_value).unwrap();
+  let global = v8::Global::new(scope, memory);
+  state.put(WasmMemory(global.into_raw()));
+}
+
+fn main() {
+  // Build a deno_core::Extension providing custom ops
+  let ext = Extension::builder()
+    .ops(vec![op_wasm::decl(), op_set_wasm_mem::decl()])
+    .build();
+
+  // Initialize a runtime instance
+  let mut runtime = JsRuntime::new(RuntimeOptions {
+    extensions: vec![ext],
+    ..Default::default()
+  });
+
+  runtime
+    .execute_script("<usage>", include_str!("wasm.js"))
+    .unwrap();
+}
diff --git a/core/examples/wasm.ts b/core/examples/wasm.ts
@@ -0,0 +1,7 @@
+// Copyright 2018-2022 the Deno authors. All rights reserved. MIT license.
+
+export declare function op_wasm(): void;
+
+export function call(): void {
+  op_wasm();
+}
diff --git a/ops/README.md b/ops/README.md
@@ -34,12 +34,29 @@ Cases where code is optimized away:
 The macro will infer and try to auto generate V8 fast API call trait impl for
 `sync` ops with:
 
-- arguments: integers, bool, `&mut OpState`, `&[u8]`, &mut [u8]`,`&[u32]`,`&mut
-  [u32]`
+- arguments: integers, bool, `&mut OpState`, `&[u8]`, `&mut [u8]`, `&[u32]`,
+  `&mut [u32]`
 - return_type: integers, bool
 
 The `#[op(fast)]` attribute should be used to enforce fast call generation at
 compile time.
 
 Trait gen for `async` ops & a ZeroCopyBuf equivalent type is planned and will be
 added soon.
+
+### Wasm calls
+
+The `#[op(wasm)]` attribute should be used for calls expected to be called from
+Wasm. This enables the fast call generation and allows seamless `WasmMemory`
+integration for generic and fast calls.
+
+```rust
+#[op(wasm)]
+pub fn op_args_get(
+  offset: i32,
+  buffer_offset: i32,
+  memory: Option<&[u8]>, // Must be last parameter. Some(..) when entered from Wasm.
+) {
+  // ...
+}
+```
diff --git a/ops/attrs.rs b/ops/attrs.rs
@@ -11,6 +11,7 @@ pub struct Attributes {
   pub is_v8: bool,
   pub must_be_fast: bool,
   pub deferred: bool,
+  pub is_wasm: bool,
 }
 
 impl Parse for Attributes {
@@ -20,18 +21,22 @@ impl Parse for Attributes {
     let vars: Vec<_> = vars.iter().map(Ident::to_string).collect();
     let vars: Vec<_> = vars.iter().map(String::as_str).collect();
     for var in vars.iter() {
-      if !["unstable", "v8", "fast", "deferred"].contains(var) {
+      if !["unstable", "v8", "fast", "deferred", "wasm"].contains(var) {
         return Err(Error::new(
           input.span(),
-          "invalid attribute, expected one of: unstable, v8, fast, deferred",
+          "invalid attribute, expected one of: unstable, v8, fast, deferred, wasm",
         ));
       }
     }
+
+    let is_wasm = vars.contains(&"wasm");
+
     Ok(Self {
       is_unstable: vars.contains(&"unstable"),
       is_v8: vars.contains(&"v8"),
-      must_be_fast: vars.contains(&"fast"),
       deferred: vars.contains(&"deferred"),
+      must_be_fast: is_wasm || vars.contains(&"fast"),
+      is_wasm,
     })
   }
 }
diff --git a/ops/fast_call.rs b/ops/fast_call.rs
@@ -139,6 +139,7 @@ pub(crate) fn generate(
 
   // Apply *hard* optimizer hints.
   if optimizer.has_fast_callback_option
+    || optimizer.has_wasm_memory
     || optimizer.needs_opstate()
     || optimizer.is_async
     || optimizer.needs_fast_callback_option
@@ -147,7 +148,7 @@ pub(crate) fn generate(
       fast_api_callback_options: *mut #core::v8::fast_api::FastApiCallbackOptions
     };
 
-    if optimizer.has_fast_callback_option {
+    if optimizer.has_fast_callback_option || optimizer.has_wasm_memory {
       // Replace last parameter.
       assert!(fast_fn_inputs.pop().is_some());
       fast_fn_inputs.push(decl);
@@ -174,6 +175,7 @@ pub(crate) fn generate(
   if optimizer.needs_opstate()
     || optimizer.is_async
     || optimizer.has_fast_callback_option
+    || optimizer.has_wasm_memory
   {
     // Dark arts 🪄 ✨
     //

diff --git a/ops/lib.rs b/ops/lib.rs
@@ -386,7 +386,9 @@ fn codegen_arg(
   let ident = quote::format_ident!("{name}");
   let (pat, ty) = match arg {
     syn::FnArg::Typed(pat) => {
-      if is_optional_fast_callback_option(&pat.ty) {
+      if is_optional_fast_callback_option(&pat.ty)
+        || is_optional_wasm_memory(&pat.ty)
+      {
         return quote! { let #ident = None; };
       }
       (&pat.pat, &pat.ty)
@@ -663,6 +665,10 @@ fn is_optional_fast_callback_option(ty: impl ToTokens) -> bool {
   tokens(&ty).contains("Option < & mut FastApiCallbackOptions")
 }
 
+fn is_optional_wasm_memory(ty: impl ToTokens) -> bool {
+  tokens(&ty).contains("Option < & mut [u8]")
+}
+
 /// Detects if the type can be set using `rv.set_uint32` fast path
 fn is_u32_rv(ty: impl ToTokens) -> bool {
   ["u32", "u8", "u16"].iter().any(|&s| tokens(&ty) == s) || is_resource_id(&ty)
@@ -743,6 +749,10 @@ mod tests {
     if source.contains("// @test-attr:fast") {
       attrs.must_be_fast = true;
     }
+    if source.contains("// @test-attr:wasm") {
+      attrs.is_wasm = true;
+      attrs.must_be_fast = true;
+    }
 
     let item = syn::parse_str(&source).expect("Failed to parse test file");
     let op = Op::new(item, attrs);

diff --git a/ops/optimizer.rs b/ops/optimizer.rs
@@ -26,6 +26,7 @@ enum TransformKind {
   SliceU32(bool),
   SliceU8(bool),
   PtrU8,
+  WasmMemory,
 }
 
 impl Transform {
@@ -50,6 +51,13 @@ impl Transform {
     }
   }
 
+  fn wasm_memory(index: usize) -> Self {
+    Transform {
+      kind: TransformKind::WasmMemory,
+      index,
+    }
+  }
+
   fn u8_ptr(index: usize) -> Self {
     Transform {
       kind: TransformKind::PtrU8,
@@ -124,6 +132,16 @@ impl Transform {
           };
         })
       }
+      TransformKind::WasmMemory => {
+        // Note: `ty` is correctly set to __opts by the fast call tier.
+        q!(Vars { var: &ident, core }, {
+          let var = unsafe {
+            &*(__opts.wasm_memory
+              as *const core::v8::fast_api::FastApiTypedArray<u8>)
+          }
+          .get_storage_if_aligned();
+        })
+      }
       // *const u8
       TransformKind::PtrU8 => {
         *ty =
@@ -201,6 +219,8 @@ pub(crate) struct Optimizer {
   // Do we depend on FastApiCallbackOptions?
   pub(crate) needs_fast_callback_option: bool,
 
+  pub(crate) has_wasm_memory: bool,
+
   pub(crate) fast_result: Option<FastValue>,
   pub(crate) fast_parameters: Vec<FastValue>,
 
@@ -262,6 +282,9 @@ impl Optimizer {
 
     self.is_async = op.is_async;
     self.fast_compatible = true;
+    // Just assume for now. We will validate later.
+    self.has_wasm_memory = op.attrs.is_wasm;
+
     let sig = &op.item.sig;
 
     // Analyze return type
@@ -419,7 +442,32 @@ impl Optimizer {
                   TypeReference { elem, .. },
                 ))) = args.last()
                 {
-                  if let Type::Path(TypePath {
+                  if self.has_wasm_memory {
+                    // -> Option<&mut [u8]>
+                    if let Type::Slice(TypeSlice { elem, .. }) = &**elem {
+                      if let Type::Path(TypePath {
+                        path: Path { segments, .. },
+                        ..
+                      }) = &**elem
+                      {
+                        let segment = single_segment(segments)?;
+
+                        match segment {
+                          // Is `T` a u8?
+                          PathSegment { ident, .. } if ident == "u8" => {
+                            self.needs_fast_callback_option = true;
+                            assert!(self
+                              .transforms
+                              .insert(index, Transform::wasm_memory(index))
+                              .is_none());
+                          }
+                          _ => {
+                            return Err(BailoutReason::FastUnsupportedParamType)
+                          }
+                        }
+                      }
+                    }
+                  } else if let Type::Path(TypePath {
                     path: Path { segments, .. },
                     ..
                   }) = &**elem
@@ -654,6 +702,10 @@ mod tests {
       .expect("Failed to read expected file");
 
     let mut attrs = Attributes::default();
+    if source.contains("// @test-attr:wasm") {
+      attrs.must_be_fast = true;
+      attrs.is_wasm = true;
+    }
     if source.contains("// @test-attr:fast") {
       attrs.must_be_fast = true;
     }

diff --git a/ops/optimizer_tests/wasm_op.expected b/ops/optimizer_tests/wasm_op.expected
@@ -0,0 +1,11 @@
+=== Optimizer Dump ===
+returns_result: false
+has_ref_opstate: false
+has_rc_opstate: false
+has_fast_callback_option: false
+needs_fast_callback_option: true
+fast_result: Some(Void)
+fast_parameters: [V8Value]
+transforms: {0: Transform { kind: WasmMemory, index: 0 }}
+is_async: false
+fast_compatible: true