Reland "[wasm] Materialize suspender in JS-to-wasm wrapper"

This is a reland of commit 8cb02753

Original change's description:
> [wasm] Materialize suspender in JS-to-wasm wrapper
>
> Instead of creating the Suspender object in JS and passing it to the
> stack-switching js-to-wasm wrapper, the wrapper now automatically
> creates the Suspender object and forwards it as an extra parameter to
> the wasm function. See:
> https://github.com/WebAssembly/js-promise-integration/pull/1/files
>
> R=ahaas@chromium.org
>
> Bug: v8:12191
> Change-Id: I2badee823f4223a293632f93e7e59f24c49d0820
> Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3779688
> Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
> Reviewed-by: Andreas Haas <ahaas@chromium.org>
> Reviewed-by: Jakob Kummerow <jkummerow@chromium.org>
> Cr-Commit-Position: refs/heads/main@{#81890}

Bug: v8:12191
Change-Id: Iea233e30aa269279d2fe17f5230c87285c33e232
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3780817
Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
Reviewed-by: Andreas Haas <ahaas@chromium.org>
Reviewed-by: Jakob Kummerow <jkummerow@chromium.org>
Cr-Commit-Position: refs/heads/main@{#82009}

src/builtins/x64/builtins-x64.cc

@@ -3014,20 +3014,19 @@ void SaveState(MacroAssembler* masm, Register active_continuation, Register tmp,
   FillJumpBuffer(masm, jmpbuf, suspend);
 }
 
-// Returns the new continuation in rax.
-void AllocateContinuation(MacroAssembler* masm, Register function_data,
-                          Register wasm_instance, Register suspender) {
+// Returns the new suspender in rax.
+void AllocateSuspender(MacroAssembler* masm, Register function_data,
+                       Register wasm_instance) {
   MemOperand GCScanSlotPlace =
       MemOperand(rbp, BuiltinWasmWrapperConstants::kGCScanSlotCountOffset);
-  __ Move(GCScanSlotPlace, 3);
+  __ Move(GCScanSlotPlace, 2);
   __ Push(wasm_instance);
   __ Push(function_data);
-  __ Push(suspender);  // Argument.
   __ LoadAnyTaggedField(
       kContextRegister,
       MemOperand(wasm_instance, wasm::ObjectAccess::ToTagged(
                                     WasmInstanceObject::kNativeContextOffset)));
-  __ CallRuntime(Runtime::kWasmAllocateContinuation);
+  __ CallRuntime(Runtime::kWasmAllocateSuspender);
   __ Pop(function_data);
   __ Pop(wasm_instance);
   static_assert(kReturnRegister0 == rax);
@@ -3189,7 +3188,9 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
   // The number of parameters according to the signature.
   constexpr int kParamCountOffset =
       BuiltinWasmWrapperConstants::kParamCountOffset;
-  constexpr int kReturnCountOffset = kParamCountOffset - kSystemPointerSize;
+  constexpr int kSuspenderOffset =
+      BuiltinWasmWrapperConstants::kSuspenderOffset;
+  constexpr int kReturnCountOffset = kSuspenderOffset - kSystemPointerSize;
   constexpr int kValueTypesArrayStartOffset =
       kReturnCountOffset - kSystemPointerSize;
   // A boolean flag to check if one of the parameters is a reference. If so, we
@@ -3201,7 +3202,9 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
   // registers (so no GC scan is needed).
   constexpr int kFunctionDataOffset = kHasRefTypesOffset - kSystemPointerSize;
   constexpr int kLastSpillOffset = kFunctionDataOffset;
-  constexpr int kNumSpillSlots = 7;
+  constexpr int kNumSpillSlots =
+      (-TypedFrameConstants::kFixedFrameSizeFromFp - kLastSpillOffset) >>
+      kSystemPointerSizeLog2;
   __ subq(rsp, Immediate(kNumSpillSlots * kSystemPointerSize));
   // Put the in_parameter count on the stack, we only  need it at the very end
   // when we pop the parameters off the stack.
@@ -3236,17 +3239,21 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
 
   Label suspend;
   if (stack_switch) {
+    // Set the suspender spill slot to a sentinel value, in case a GC happens
+    // before we set the actual value.
+    __ LoadRoot(kScratchRegister, RootIndex::kUndefinedValue);
+    __ movq(MemOperand(rbp, kSuspenderOffset), kScratchRegister);
     Register active_continuation = rbx;
     __ LoadRoot(active_continuation, RootIndex::kActiveContinuation);
     SaveState(masm, active_continuation, rcx, &suspend);
-    Register suspender = rax;
-    constexpr int kReceiverOnStackSize = kSystemPointerSize;
-    constexpr int kFirstParamOffset =
-        kFPOnStackSize + kPCOnStackSize + kReceiverOnStackSize;
-    __ movq(suspender, MemOperand(rbp, kFirstParamOffset));
-    AllocateContinuation(masm, function_data, wasm_instance, suspender);
+    AllocateSuspender(masm, function_data, wasm_instance);
+    Register suspender = rax;  // Fixed.
+    __ movq(MemOperand(rbp, kSuspenderOffset), suspender);
+    Register target_continuation = rax;
+    __ LoadAnyTaggedField(
+        target_continuation,
+        FieldOperand(suspender, WasmSuspenderObject::kContinuationOffset));
     suspender = no_reg;
-    Register target_continuation = rax;  // fixed
     // Save the old stack's rbp in r9, and use it to access the parameters in
     // the parent frame.
     // We also distribute the spill slots across the two stacks as needed by
@@ -3264,9 +3271,11 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
     //      +-----------------+           +-----------------+
     //      |kGCScanSlotCount |           |kGCScanSlotCount |
     //      +-----------------+           +-----------------+
+    //      | kInParamCount   |           |      /          |
+    //      +-----------------+           +-----------------+
     //      | kParamCount     |           |      /          |
     //      +-----------------+           +-----------------+
-    //      | kInParamCount   |           |      /          |
+    //      | kSuspender      |           |      /          |
     //      +-----------------+           +-----------------+
     //      |      /          |           | kReturnCount    |
     //      +-----------------+           +-----------------+
@@ -3298,6 +3307,9 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
     // this marks the base of the stack segment for the stack frame iterator.
     __ EnterFrame(StackFrame::STACK_SWITCH);
     __ subq(rsp, Immediate(kNumSpillSlots * kSystemPointerSize));
+    // Set a sentinel value for the suspender spill slot in the new frame.
+    __ LoadRoot(kScratchRegister, RootIndex::kUndefinedValue);
+    __ movq(MemOperand(rbp, kSuspenderOffset), kScratchRegister);
   }
   Register original_fp = stack_switch ? r9 : rbp;
 
@@ -3439,6 +3451,22 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
   returns_size = no_reg;
   Register valuetype = r12;
 
+  Label numeric_params_done;
+  if (stack_switch) {
+    // Prepare for materializing the suspender parameter. We don't materialize
+    // it here but in the next loop that processes references. Here we only
+    // adjust the pointers to keep the state consistent:
+    // - Skip the first valuetype in the signature,
+    // - Adjust the param limit which is off by one because of the extra
+    // param in the signature,
+    // - Set HasRefTypes to 1 to ensure that the reference loop is entered.
+    __ addq(valuetypes_array_ptr, Immediate(kValueTypeSize));
+    __ subq(param_limit, Immediate(kSystemPointerSize));
+    __ movq(MemOperand(rbp, kHasRefTypesOffset), Immediate(1));
+    __ cmpq(current_param, param_limit);
+    __ j(equal, &numeric_params_done);
+  }
+
   // -------------------------------------------
   // Param evaluation loop.
   // -------------------------------------------
@@ -3458,7 +3486,7 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
   __ cmpq(valuetype, Immediate(wasm::kWasmI32.raw_bit_field()));
   __ j(not_equal, &convert_param);
   __ JumpIfNotSmi(param, &convert_param);
-  // Change the paramfrom Smi to int32.
+  // Change the param from Smi to int32.
   __ SmiUntag(param);
   // Zero extend.
   __ movl(param, param);
@@ -3477,6 +3505,7 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
 
   __ cmpq(current_param, param_limit);
   __ j(not_equal, &loop_through_params);
+  __ bind(&numeric_params_done);
 
   // -------------------------------------------
   // Second loop to handle references.
@@ -3505,6 +3534,17 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
   __ Move(current_param,
           kFPOnStackSize + kPCOnStackSize + kReceiverOnStackSize);
 
+  if (stack_switch) {
+    // Materialize the suspender param
+    __ movq(param, MemOperand(original_fp, kSuspenderOffset));
+    __ movq(MemOperand(current_int_param_slot, 0), param);
+    __ subq(current_int_param_slot, Immediate(kSystemPointerSize));
+    __ addq(valuetypes_array_ptr, Immediate(kValueTypeSize));
+    __ addq(ref_param_count, Immediate(1));
+    __ cmpq(current_param, param_limit);
+    __ j(equal, &ref_params_done);
+  }
+
   Label ref_loop_through_params;
   Label ref_loop_end;
   // Start of the loop.
@@ -4157,6 +4197,10 @@ void Generate_WasmResumeHelper(MacroAssembler* masm, wasm::OnResume on_resume) {
   __ subq(rsp, Immediate(3 * kSystemPointerSize));
   __ movq(MemOperand(rbp, kParamCountOffset), param_count);
   __ movq(MemOperand(rbp, kInParamCountOffset), param_count);
+  // Set a sentinel value for the spill slot visited by the GC.
+  __ LoadRoot(kScratchRegister, RootIndex::kUndefinedValue);
+  __ movq(MemOperand(rbp, BuiltinWasmWrapperConstants::kSuspenderOffset),
+          kScratchRegister);
 
   param_count = no_reg;
 

src/common/message-template.h

@@ -652,7 +652,6 @@ namespace internal {
   T(WasmTrapStringOffsetOutOfBounds, "string offset out of bounds")            \
   T(WasmTrapStringIsolatedSurrogate,                                           \
     "Failed to encode string as UTF-8: contains unpaired surrogate")           \
-  T(WasmTrapReentrantSuspender, "re-entering an active/suspended suspender")   \
   T(WasmExceptionError, "wasm exception")                                      \
   /* Asm.js validation related */                                              \
   T(AsmJsInvalid, "Invalid asm.js: %")                                         \

src/execution/frame-constants.h

@@ -216,6 +216,7 @@ class BuiltinWasmWrapperConstants : public TypedFrameConstants {
   static constexpr int kInParamCountOffset = TYPED_FRAME_PUSHED_VALUE_OFFSET(1);
   // The number of parameters according to the signature.
   static constexpr int kParamCountOffset = TYPED_FRAME_PUSHED_VALUE_OFFSET(2);
+  static constexpr int kSuspenderOffset = TYPED_FRAME_PUSHED_VALUE_OFFSET(3);
 };
 
 class ConstructFrameConstants : public TypedFrameConstants {

src/execution/frames.cc

@@ -2367,6 +2367,10 @@ void StackSwitchFrame::Iterate(RootVisitor* v) const {
       &Memory<Address>(sp() + scan_count * kSystemPointerSize));
   v->VisitRootPointers(Root::kStackRoots, nullptr, spill_slot_base,
                        spill_slot_limit);
+  // Also visit fixed spill slots that contain references.
+  FullObjectSlot suspender_slot(
+      &Memory<Address>(fp() + BuiltinWasmWrapperConstants::kSuspenderOffset));
+  v->VisitRootPointer(Root::kStackRoots, nullptr, suspender_slot);
 }
 
 // static

src/runtime/runtime-wasm.cc

@@ -796,21 +796,12 @@ void SyncStackLimit(Isolate* isolate) {
 }
 }  // namespace
 
-// Allocate a new continuation, and prepare for stack switching by updating the
+// Allocate a new suspender, and prepare for stack switching by updating the
 // active continuation, active suspender and stack limit.
-RUNTIME_FUNCTION(Runtime_WasmAllocateContinuation) {
+RUNTIME_FUNCTION(Runtime_WasmAllocateSuspender) {
   CHECK(FLAG_experimental_wasm_stack_switching);
   HandleScope scope(isolate);
-  if (!args[0].IsWasmSuspenderObject()) {
-    return ThrowWasmError(isolate, MessageTemplate::kWasmTrapJSTypeError);
-  }
-  Handle<WasmSuspenderObject> suspender =
-      handle(WasmSuspenderObject::cast(args[0]), isolate);
-
-  if (suspender->state() != WasmSuspenderObject::kInactive) {
-    return ThrowWasmError(isolate,
-                          MessageTemplate::kWasmTrapReentrantSuspender);
-  }
+  Handle<WasmSuspenderObject> suspender = WasmSuspenderObject::New(isolate);
 
   // Update the continuation state.
   auto parent = handle(WasmContinuationObject::cast(
@@ -832,7 +823,7 @@ RUNTIME_FUNCTION(Runtime_WasmAllocateContinuation) {
   active_suspender_slot.store(*suspender);
 
   SyncStackLimit(isolate);
-  return *target;
+  return *suspender;
 }
 
 // Update the stack limit after a stack switch, and preserve pending interrupts.

src/runtime/runtime.h

@@ -609,7 +609,7 @@ namespace internal {
   F(WasmDebugBreak, 0, 1)             \
   F(WasmArrayCopy, 5, 1)              \
   F(WasmArrayNewSegment, 5, 1)        \
-  F(WasmAllocateContinuation, 1, 1)   \
+  F(WasmAllocateSuspender, 0, 1)      \
   F(WasmSyncStackLimit, 0, 1)         \
   F(WasmCreateResumePromise, 2, 1)    \
   F(WasmStringNewWtf8, 5, 1)          \

src/wasm/stacks.h

@@ -104,7 +104,7 @@ class StackMemory {
   StackMemory(Isolate* isolate, byte* limit)
       : isolate_(isolate),
         limit_(limit),
-        size_(reinterpret_cast<size_t>(limit)),
+        size_(FLAG_stack_size * KB),
         owned_(false) {
     id_ = 0;
   }

src/wasm/wasm-js.cc

@@ -2026,13 +2026,15 @@ void WebAssemblyFunctionType(const v8::FunctionCallbackInfo<v8::Value>& args) {
         handle(sfi->wasm_exported_function_data(), i_isolate);
     sig = wasm_exported_function->sig();
     if (data->suspend()) {
-      // If this export is suspendable, the function returns a
+      // If this export is suspendable, the first parameter of the original
+      // function is an externref (suspender) which does not appear in the
+      // wrapper function's signature. The wrapper function also returns a
       // promise as an externref instead of the original return type.
       size_t param_count = sig->parameter_count();
       DCHECK_GE(param_count, 1);
       DCHECK_EQ(sig->GetParam(0), i::wasm::kWasmAnyRef);
-      i::wasm::FunctionSig::Builder builder(&zone, 1, param_count);
-      for (size_t i = 0; i < param_count; ++i) {
+      i::wasm::FunctionSig::Builder builder(&zone, 1, param_count - 1);
+      for (size_t i = 1; i < param_count; ++i) {
         builder.AddParam(sig->GetParam(i));
       }
       builder.AddReturn(i::wasm::kWasmAnyRef);

src/wasm/wasm-objects.cc

@@ -1808,10 +1808,8 @@ Handle<WasmContinuationObject> WasmContinuationObject::New(
 Handle<WasmSuspenderObject> WasmSuspenderObject::New(Isolate* isolate) {
   Handle<JSFunction> suspender_cons(
       isolate->native_context()->wasm_suspender_constructor(), isolate);
-  // Suspender objects should be at least as long-lived as the instances of
-  // which it will wrap the imports/exports, allocate in old space too.
   auto suspender = Handle<WasmSuspenderObject>::cast(
-      isolate->factory()->NewJSObject(suspender_cons, AllocationType::kOld));
+      isolate->factory()->NewJSObject(suspender_cons));
   suspender->set_state(kInactive);
   // Instantiate the callable object which resumes this Suspender. This will be
   // used implicitly as the onFulfilled callback of the returned JS promise.

test/mjsunit/mjsunit.status

@@ -1558,6 +1558,7 @@
   # These tests rely on allocation site tracking which only works in the young generation.
   'array-constructor-feedback': [SKIP],
   'wasm/generic-wrapper': [SKIP],
+  'wasm/stack-switching-export': [SKIP],
   'regress/regress-trap-allocation-memento': [SKIP],
   'regress/regress-crbug-1151890': [SKIP],
   'regress/regress-crbug-1163184': [SKIP],
@@ -1716,6 +1717,7 @@
 ['arch != x64', {
   # Stack switching is only supported on x64.
   'wasm/stack-switching': [SKIP],
+  'wasm/stack-switching-export': [SKIP],
 }],  # arch != x64
 
 ##############################################################################