Bug 1899637 - Only suspend wrapped imports that return Promises. r=rhunt

Don't use promise if the result of suspending function is know. Modify stack
to shortcut the exection without promise queue.

Differential Revision: https://phabricator.services.mozilla.com/D221330

js/src/jit-test/tests/wasm/js-promise-integration/js-promise-integration.new.js

@@ -124,8 +124,9 @@ promise_test(async () => {
     )
   )`, {m: {import: wasm_js_import}});
   let wrapped_export = WebAssembly.promising(instance.exports.test);
-  await wrapped_export();
+  const p = wrapped_export();
   assert_equals(42, instance.exports.g.value);
+  await p;
 }, "Do not suspend if the import's return value is not a Promise");
 
 test(t => {

js/src/jit/CodeGenerator.cpp

@@ -10113,6 +10113,14 @@ void CodeGenerator::visitWasmStackSwitchToMain(LWasmStackSwitchToMain* lir) {
   // DataReg is not needed anymore, using it as a scratch register.
   const Register ScratchReg2 = DataReg;
 
+  // Save future of main stack exit frame pointer.
+  masm.computeEffectiveAddress(
+      Address(masm.getStackPointer(), -int32_t(sizeof(wasm::Frame))),
+      ScratchReg2);
+  masm.storePtr(ScratchReg2,
+                Address(SuspenderDataReg,
+                        wasm::SuspenderObjectData::offsetOfMainExitFP()));
+
   // Load InstanceReg from suspendable stack exit frame.
   masm.loadPtr(Address(SuspenderDataReg,
                        wasm::SuspenderObjectData::offsetOfSuspendableExitFP()),

js/src/wasm/GenerateBuiltinModules.py

@@ -54,6 +54,7 @@ def specTypeToMIRType(specType):
         specType == "externref"
         or specType == "anyref"
         or specType == "funcref"
+        or specType == "exnref"
         or isinstance(specType, dict)
     ):
         return "MIRType::WasmAnyRef"
@@ -67,6 +68,8 @@ def specHeapTypeToTypeCode(specHeapType):
         return "Any"
     if specHeapType == "extern":
         return "Extern"
+    if specHeapType == "exn":
+        return "Exn"
     if specHeapType == "array":
         return "Array"
     if specHeapType == "struct":
@@ -81,6 +84,9 @@ def specTypeToValType(specType):
     if specType == "externref":
         return "ValType(RefType::extern_())"
 
+    if specType == "exnref":
+        return "ValType(RefType::exn())"
+
     if specType == "anyref":
         return "ValType(RefType::any())"
 

js/src/wasm/WasmBuiltinModule.yaml

@@ -458,14 +458,28 @@
 - op: AddPromiseReactions
   symbolic_address:
     name: AddPromiseReactions
-    type: Args_Int32_GeneralGeneralGeneralGeneral
+    type: Args_General4
   entry: AddPromiseReactions
   export: addPromiseReactions
   params:
     - 'externref'
     - 'externref'
     - 'funcref'
-  fail_mode: FailOnNegI32
+  result: 'anyref'
+  fail_mode: FailOnInvalidRef
+  uses_memory: false
+
+- op: ForwardExceptionToSuspended
+  symbolic_address:
+    name: ForwardExceptionToSuspended
+    type: Args_General3
+  entry: ForwardExceptionToSuspended
+  export: forwardExceptionToSuspended
+  params:
+    - 'externref'
+    - 'exnref'
+  result: 'anyref'
+  fail_mode: Infallible
   uses_memory: false
 
 - op: SetPromisingPromiseResults

js/src/wasm/WasmCompile.cpp

@@ -219,6 +219,9 @@ FeatureArgs FeatureArgs::build(JSContext* cx, const FeatureOptions& options) {
     features.tailCalls = true;
   }
 #endif
+  if (options.requireExnref) {
+    features.exnref = true;
+  }
 
   return features;
 }

js/src/wasm/WasmCompileArgs.h

@@ -102,7 +102,8 @@ struct FeatureOptions {
   FeatureOptions()
       : isBuiltinModule(false),
         jsStringBuiltins(false),
-        jsStringConstants(false)
+        jsStringConstants(false),
+        requireExnref(false)
 #ifdef ENABLE_WASM_GC
         ,
         requireGC(false)
@@ -124,6 +125,8 @@ struct FeatureOptions {
   bool jsStringConstants;
   SharedChars jsStringConstantsNamespace;
 
+  // Enable exnref support.
+  bool requireExnref;
 #ifdef ENABLE_WASM_GC
   // Enable GC support.
   bool requireGC;

js/src/wasm/WasmFrameIter.cpp

@@ -305,6 +305,7 @@ void WasmFrameIter::popFrame() {
 #endif
 
   MOZ_ASSERT(code_ == &instance()->code());
+
   lineOrBytecode_ = callsite->lineOrBytecode();
   failedUnwindSignatureMismatch_ = false;
 

js/src/wasm/WasmOpIter.h

@@ -4310,9 +4310,13 @@ inline bool OpIter<Policy>::readStackSwitch(StackSwitchKind* kind,
     ValType valType = stackType.valType();
     MOZ_ASSERT(valType.isRefType() && valType.typeDef()->isFuncType());
     const FuncType& func = valType.typeDef()->funcType();
-    MOZ_ASSERT(func.args().length() == 2 && func.results().empty() &&
+    MOZ_ASSERT(func.args().length() == 2 &&
                func.arg(0).isExternRef() &&
                ValType::isSubTypeOf(func.arg(1), RefType::any()));
+    MOZ_ASSERT_IF(*kind != StackSwitchKind::SwitchToMain,
+                  func.results().empty());
+    MOZ_ASSERT_IF(*kind == StackSwitchKind::SwitchToMain,
+                  func.results().length() == 1 && func.result(0).isAnyRef());
   }
 #  endif
   if (!popWithType(ValType(RefType::extern_()), suspender)) {

js/src/wasm/WasmPI.cpp

@@ -212,7 +212,9 @@ class SuspenderObject : public NativeObject {
  public:
   static const JSClass class_;
 
-  enum { DataSlot, PromisingPromiseSlot, SlotCount };
+  enum { DataSlot, PromisingPromiseSlot, SuspendingReturnTypeSlot, SlotCount };
+
+  enum class ReturnType : int32_t { Unknown, Promise, Value, Exception };
 
   static SuspenderObject* create(JSContext* cx) {
     for (;;) {
@@ -254,6 +256,8 @@ class SuspenderObject : public NativeObject {
 
     suspender->initReservedSlot(DataSlot, PrivateValue(data));
     suspender->initReservedSlot(PromisingPromiseSlot, NullValue());
+    suspender->initReservedSlot(SuspendingReturnTypeSlot,
+                                Int32Value(int32_t(ReturnType::Unknown)));
     return suspender;
   }
 
@@ -267,6 +271,20 @@ class SuspenderObject : public NativeObject {
     setReservedSlot(PromisingPromiseSlot, ObjectOrNullValue(promise));
   }
 
+  ReturnType suspendingReturnType() const {
+    return ReturnType(getReservedSlot(SuspendingReturnTypeSlot).toInt32());
+  }
+
+  void setSuspendingReturnType(ReturnType type) {
+    // The SuspendingReturnTypeSlot will change after result is defined,
+    // and becomes invalid after GetSuspendingPromiseResult. The assert is
+    // checking if the result was processed by GetSuspendingPromiseResult.
+    MOZ_ASSERT((type == ReturnType::Unknown) !=
+               (suspendingReturnType() == ReturnType::Unknown));
+
+    setReservedSlot(SuspendingReturnTypeSlot, Int32Value(int32_t(type)));
+  }
+
   JS::NativeStackLimit getStackMemoryLimit() {
     return JS::NativeStackLimit(data()->stackMemory()) + SuspendableRedZoneSize;
   }
@@ -289,6 +307,9 @@ class SuspenderObject : public NativeObject {
   void resume(JSContext* cx);
   void leave(JSContext* cx);
 
+  // Modifies frames to inject the suspendable stack back into the main one.
+  void forwardToSuspendable();
+
  private:
   static const JSClassOps classOps_;
 
@@ -463,6 +484,16 @@ void SuspenderObject::leave(JSContext* cx) {
   }
 }
 
+void SuspenderObject::forwardToSuspendable() {
+  // Injecting suspendable stack back into main one at the exit frame.
+  SuspenderObjectData* data = this->data();
+  uint8_t* mainExitFP = (uint8_t*)data->mainExitFP();
+  *reinterpret_cast<void**>(mainExitFP + Frame::callerFPOffset()) =
+      data->suspendableFP();
+  *reinterpret_cast<void**>(mainExitFP + Frame::returnAddressOffset()) =
+      data->suspendedReturnAddress();
+}
+
 bool CallOnMainStack(JSContext* cx, CallOnMainStackFn fn, void* data) {
   Rooted<SuspenderObject*> suspender(
       cx, cx->wasm().promiseIntegration.activeSuspender());
@@ -866,13 +897,21 @@ class SuspendingFunctionModuleFactory {
 
   // Builds function that is called on main stack:
   // (func $suspending.trampoline
-  //   (param $params (ref $suspender) (ref $param-type))
-  //   (result externref)
-  //   local.get $suspender ;; for call $add-promise-reactions
-  //   (struct.get $param-type $i (local.get $param))*
-  //   call $suspending.wrappedfn
-  //   ref.func $suspending.continue-on-suspendable
-  //   call $builtin.add-promise-reactions
+  //   (param $params (ref $suspender)) (param $param (ref $param-type))
+  //   (result anyref)
+  //   local.get $suspender ;; for $builtin.forward-exn-to-suspended below
+  //   block (result exnref)
+  //    try_table (catch_all_ref 0)
+  //     local.get $suspender ;; for call $add-promise-reactions
+  //     (struct.get $param-type $i (local.get $param))*
+  //     call $suspending.wrappedfn
+  //     ref.func $suspending.continue-on-suspendable
+  //     call $builtin.add-promise-reactions
+  //     return
+  //    end
+  //    unreachable
+  //   end
+  //   call $builtin.forward-exn-to-suspended
   // )
   // The function calls suspending import and returns into the
   // $promising.exported function because that was the top function
@@ -886,6 +925,24 @@ class SuspendingFunctionModuleFactory {
     const uint32_t SuspenderIndex = 0;
     const uint32_t ParamsIndex = 1;
 
+    if (!encoder.writeOp(Op::LocalGet) ||
+        !encoder.writeVarU32(SuspenderIndex)) {
+      return false;
+    }
+
+    if (!encoder.writeOp(Op::Block) ||
+        !encoder.writeFixedU8(uint8_t(TypeCode::ExnRef))) {
+      return false;
+    }
+
+    if (!encoder.writeOp(Op::TryTable) ||
+        !encoder.writeFixedU8(uint8_t(TypeCode::BlockVoid)) ||
+        !encoder.writeVarU32(1) ||
+        !encoder.writeFixedU8(/* catch_all_ref = */ 0x03) ||
+        !encoder.writeVarU32(0)) {
+      return false;
+    }
+
     // For AddPromiseReactions call below.
     if (!encoder.writeOp(Op::LocalGet) ||
         !encoder.writeVarU32(SuspenderIndex)) {
@@ -915,6 +972,17 @@ class SuspendingFunctionModuleFactory {
       return false;
     }
 
+    if (!encoder.writeOp(Op::Return) || !encoder.writeOp(Op::End) ||
+        !encoder.writeOp(Op::Unreachable) || !encoder.writeOp(Op::End)) {
+      return false;
+    }
+
+    if (!encoder.writeOp(MozOp::CallBuiltinModuleFunc) ||
+        !encoder.writeVarU32(
+            (uint32_t)BuiltinModuleFuncId::ForwardExceptionToSuspended)) {
+      return false;
+    }
+
     return encoder.writeOp(Op::End);
   }
 
@@ -967,7 +1035,7 @@ class SuspendingFunctionModuleFactory {
     FeatureOptions options;
     options.isBuiltinModule = true;
     options.requireGC = true;
-    options.requireTailCalls = true;
+    options.requireExnref = true;
 
     ScriptedCaller scriptedCaller;
     SharedCompileArgs compileArgs =
@@ -1048,7 +1116,8 @@ class SuspendingFunctionModuleFactory {
     ValTypeVector paramsTrampoline, resultsTrampoline;
     if (!paramsTrampoline.emplaceBack(suspenderType) ||
         !paramsTrampoline.emplaceBack(RefType::fromTypeDef(
-            &(*codeMeta->types)[ParamsTypeIndex], false))) {
+            &(*codeMeta->types)[ParamsTypeIndex], false)) ||
+        !resultsTrampoline.emplaceBack(RefType::any())) {
       ReportOutOfMemory(cx);
       return nullptr;
     }
@@ -1161,6 +1230,20 @@ static bool WasmPISuspendTaskContinue(JSContext* cx, unsigned argc, Value* vp) {
   return RejectPromiseWithPendingError(cx, promise);
 }
 
+static bool IsPromiseValue(JSContext* aCx, JS::Handle<JS::Value> aValue) {
+  if (!aValue.isObject()) {
+    return false;
+  }
+
+  // We only care about Promise here, so CheckedUnwrapStatic is fine.
+  JS::Rooted<JSObject*> obj(aCx, js::CheckedUnwrapStatic(&aValue.toObject()));
+  if (!obj) {
+    return false;
+  }
+
+  return JS::IsPromiseObject(obj);
+}
+
 // Wraps original import to catch all exceptions and convert result to a
 // promise.
 // Seen as $suspending.wrappedfn in wasm.
@@ -1175,29 +1258,14 @@ static bool WasmPIWrapSuspendingImport(JSContext* cx, unsigned argc,
   // Catching exceptions here.
   RootedValue rval(cx);
   if (Call(cx, UndefinedHandleValue, originalImportFunc, args, &rval)) {
-    // Convert the result to a resolved promise.
-    RootedObject promiseConstructor(cx, GetPromiseConstructor(cx));
-    RootedObject promiseObj(cx, PromiseResolve(cx, promiseConstructor, rval));
-    if (!promiseObj) {
-      return false;
-    }
-    args.rval().setObject(*promiseObj);
+    // Convert the result to a resolved promise later in AddPromiseReactions.
+    args.rval().set(rval);
     return true;
   }
 
-  if (cx->isThrowingOutOfMemory()) {
-    return false;
-  }
-
-  // Convert failure to a rejected promise.
-  RootedObject promiseObject(cx, NewPromiseObject(cx, nullptr));
-  if (!promiseObject) {
-    return false;
-  }
-  args.rval().setObject(*promiseObject);
-
-  Rooted<PromiseObject*> promise(cx, &promiseObject->as<PromiseObject>());
-  return RejectPromiseWithPendingError(cx, promise);
+  // Deferring pending exception to the handler in the
+  // $suspending.trampoline.
+  return false;
 }
 
 JSFunction* WasmSuspendingFunctionCreate(JSContext* cx, HandleObject func,
@@ -1419,7 +1487,6 @@ class PromisingFunctionModuleFactory {
     FeatureOptions options;
     options.isBuiltinModule = true;
     options.requireGC = true;
-    options.requireTailCalls = true;
 
     ScriptedCaller scriptedCaller;
     SharedCompileArgs compileArgs =
@@ -1672,16 +1739,35 @@ PromiseObject* CreatePromisingPromise(Instance* instance,
 // Converts promise results into actual function result, or exception/trap
 // if rejected.
 // Seen as $builtin.get-suspending-promise-result to wasm.
-JSObject* GetSuspendingPromiseResult(Instance* instance, PromiseObject* promise,
+JSObject* GetSuspendingPromiseResult(Instance* instance, void* result,
                                      SuspenderObject* suspender) {
   MOZ_ASSERT(SASigGetSuspendingPromiseResult.failureMode ==
              FailureMode::FailOnNullPtr);
   JSContext* cx = instance->cx();
   Rooted<SuspenderObject*> suspenderObject(cx, suspender);
+  RootedAnyRef resultRef(cx, AnyRef::fromCompiledCode(result));
 
-  if (promise->state() == JS::PromiseState::Rejected) {
-    RootedValue reason(cx, promise->reason());
-    // Result is also the reason of promise rejection.
+  SuspenderObject::ReturnType returnType =
+      suspenderObject->suspendingReturnType();
+  MOZ_ASSERT(returnType != SuspenderObject::ReturnType::Unknown);
+  Rooted<PromiseObject*> promise(
+      cx, returnType == SuspenderObject::ReturnType::Promise
+              ? &resultRef.toJSObject().as<PromiseObject>()
+              : nullptr);
+
+#  ifdef DEBUG
+  auto resetReturnType = mozilla::MakeScopeExit([&suspenderObject]() {
+    suspenderObject->setSuspendingReturnType(
+        SuspenderObject::ReturnType::Unknown);
+  });
+#  endif
+
+  if (promise ? promise->state() == JS::PromiseState::Rejected
+              : returnType == SuspenderObject::ReturnType::Exception) {
+    // Promise was rejected or an exception was thrown, set pending exception
+    // and fail.
+    RootedValue reason(
+        cx, promise ? promise->reason() : resultRef.get().toJSValue());
     cx->setPendingException(reason, ShouldCaptureStack::Maybe);
     return nullptr;
   }
@@ -1693,7 +1779,13 @@ JSObject* GetSuspendingPromiseResult(Instance* instance, PromiseObject* promise,
   const FieldTypeVector& fields = results->typeDef().structType().fields_;
 
   if (fields.length() > 0) {
-    RootedValue jsValue(cx, promise->value());
+    MOZ_ASSERT_IF(promise, promise->state() == JS::PromiseState::Fulfilled);
+    RootedValue jsValue(cx);
+    if (promise) {
+      jsValue.set(promise->value());
+    } else {
+      jsValue.set(resultRef.get().toJSValue());
+    }
 
     // The struct object is constructed based on returns of exported function.
     // It is the only way we can get ValType for Val::fromJSValue call.
@@ -1744,17 +1836,36 @@ JSObject* GetSuspendingPromiseResult(Instance* instance, PromiseObject* promise,
 // Collects returned suspending promising, and registers callbacks to
 // react on it using WasmPISuspendTaskContinue.
 // Seen as $builtin.add-promise-reactions to wasm.
-int32_t AddPromiseReactions(Instance* instance, SuspenderObject* suspender,
-                            PromiseObject* promise,
-                            JSFunction* continueOnSuspendable) {
-  MOZ_ASSERT(SASigSetPromisingPromiseResults.failureMode ==
-             FailureMode::FailOnNegI32);
+void* AddPromiseReactions(Instance* instance, SuspenderObject* suspender,
+                          void* result, JSFunction* continueOnSuspendable) {
+  MOZ_ASSERT(SASigAddPromiseReactions.failureMode ==
+             FailureMode::FailOnInvalidRef);
   JSContext* cx = instance->cx();
 
+  RootedAnyRef resultRef(cx, AnyRef::fromCompiledCode(result));
+  RootedValue resultValue(cx, resultRef.get().toJSValue());
   Rooted<SuspenderObject*> suspenderObject(cx, suspender);
-  Rooted<PromiseObject*> promiseObject(cx, promise);
   RootedFunction fn(cx, continueOnSuspendable);
 
+  if (!IsPromiseValue(cx, resultValue)) {
+    suspenderObject->forwardToSuspendable();
+    suspenderObject->setSuspendingReturnType(
+        SuspenderObject::ReturnType::Value);
+    return resultRef.get().forCompiledCode();
+  }
+
+  // Wrap a promise.
+  RootedObject promiseConstructor(cx, GetPromiseConstructor(cx));
+  RootedObject promiseObj(cx,
+                          PromiseResolve(cx, promiseConstructor, resultValue));
+  if (!promiseObj) {
+    return AnyRef::invalid().forCompiledCode();
+  }
+  Rooted<PromiseObject*> promiseObject(cx, &promiseObj->as<PromiseObject>());
+
+  suspenderObject->setSuspendingReturnType(
+      SuspenderObject::ReturnType::Promise);
+
   // Add promise reactions
   RootedFunction then_(
       cx, NewNativeFunction(cx, WasmPISuspendTaskContinue, 1, nullptr,
@@ -1762,7 +1873,22 @@ int32_t AddPromiseReactions(Instance* instance, SuspenderObject* suspender,
   then_->initExtendedSlot(SUSPENDER_SLOT, ObjectValue(*suspenderObject));
   then_->initExtendedSlot(CONTINUE_ON_SUSPENDABLE_SLOT, ObjectValue(*fn));
   then_->initExtendedSlot(PROMISE_SLOT, ObjectValue(*promiseObject));
-  return JS::AddPromiseReactions(cx, promiseObject, then_, then_) ? 0 : -1;
+  if (!JS::AddPromiseReactions(cx, promiseObject, then_, then_)) {
+    return AnyRef::invalid().forCompiledCode();
+  }
+  return AnyRef::fromJSObject(*promiseObject).forCompiledCode();
+}
+
+// Changes exit stack frame pointers to suspendable stack and recast exception
+// to wasm reference. Seen as $builtin.forward-exn-to-suspended to wasm.
+void* ForwardExceptionToSuspended(Instance* instance,
+                                  SuspenderObject* suspender, void* exception) {
+  MOZ_ASSERT(SASigForwardExceptionToSuspended.failureMode ==
+             FailureMode::Infallible);
+
+  suspender->forwardToSuspendable();
+  suspender->setSuspendingReturnType(SuspenderObject::ReturnType::Exception);
+  return exception;
 }
 
 // Resolves the promise using results packed by wasm.

js/src/wasm/WasmPI.h

@@ -143,6 +143,9 @@ class SuspenderObjectData
   // Stored return address for return to suspendable stack.
   void* suspendedReturnAddress_;
 
+  // Stored main stack exit/top frame pointer.
+  void* mainExitFP_;
+
   SuspenderState state_;
 
   // Identify context that is holding suspended stack, otherwise nullptr.
@@ -170,6 +173,7 @@ class SuspenderObjectData
   inline void* stackMemory() const { return stackMemory_; }
   inline void* mainFP() const { return mainFP_; }
   inline void* mainSP() const { return mainSP_; }
+  inline void* mainExitFP() const { return mainExitFP_; }
   inline void* suspendableFP() const { return suspendableFP_; }
   inline void* suspendableSP() const { return suspendableSP_; }
   inline void* suspendableExitFP() const { return suspendableExitFP_; }
@@ -209,6 +213,10 @@ class SuspenderObjectData
     return offsetof(SuspenderObjectData, suspendableExitFP_);
   }
 
+  static constexpr size_t offsetOfMainExitFP() {
+    return offsetof(SuspenderObjectData, mainExitFP_);
+  }
+
   static constexpr size_t offsetOfSuspendedReturnAddress() {
     return offsetof(SuspenderObjectData, suspendedReturnAddress_);
   }
@@ -237,12 +245,14 @@ SuspenderObject* CreateSuspender(Instance* instance, int reserved);
 PromiseObject* CreatePromisingPromise(Instance* instance,
                                       SuspenderObject* suspender);
 
-JSObject* GetSuspendingPromiseResult(Instance* instance, PromiseObject* promise,
+JSObject* GetSuspendingPromiseResult(Instance* instance, void* result,
                                      SuspenderObject* suspender);
 
-int32_t AddPromiseReactions(Instance* instance, SuspenderObject* suspender,
-                            PromiseObject* promise,
-                            JSFunction* continueOnSuspendable);
+void* AddPromiseReactions(Instance* instance, SuspenderObject* suspender,
+                          void* result, JSFunction* continueOnSuspendable);
+
+void* ForwardExceptionToSuspended(Instance* instance,
+                                  SuspenderObject* suspender, void* exception);
 
 int32_t SetPromisingPromiseResults(Instance* instance,
                                    SuspenderObject* suspender,