1 files changed, 53 insertions, 1 deletions

diff --git a/ops/optimizer.rs b/ops/optimizer.rs
index 8fa2ab1f5..d25857032 100644
--- 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;
     }