[wasm][stack-switching] Fix nested suspenders

Fix some issues with nested suspenders:
- Fix scratch register conflict when returning from an inner suspender
- The outer suspender should stay in 'Active' state
- Suspenders should become 'Inactive' when they return

CC=ahaas@chromium.org

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

src/builtins/builtins.cc

@@ -523,6 +523,7 @@ bool Builtins::CodeObjectIsExecutable(Builtin builtin) {
     case Builtin::kInstantiateAsmJs:
 #if V8_ENABLE_WEBASSEMBLY
     case Builtin::kGenericJSToWasmWrapper:
+    case Builtin::kWasmReturnPromiseOnSuspend:
 #endif  // V8_ENABLE_WEBASSEMBLY
 
     // TODO(delphick): Remove this when calls to it have the trampoline inlined

src/builtins/x64/builtins-x64.cc

@@ -3092,21 +3092,25 @@ void ReloadParentContinuation(MacroAssembler* masm, Register wasm_instance,
   __ Pop(return_reg);
 }
 
-void RestoreParentSuspender(MacroAssembler* masm) {
-  Register suspender = kScratchRegister;
+void RestoreParentSuspender(MacroAssembler* masm, Register tmp1,
+                            Register tmp2) {
+  Register suspender = tmp1;
   __ LoadRoot(suspender, RootIndex::kActiveSuspender);
+  __ StoreTaggedSignedField(
+      FieldOperand(suspender, WasmSuspenderObject::kStateOffset),
+      Smi::FromInt(WasmSuspenderObject::kInactive));
   __ LoadAnyTaggedField(
       suspender, FieldOperand(suspender, WasmSuspenderObject::kParentOffset));
   __ CompareRoot(suspender, RootIndex::kUndefinedValue);
   Label undefined;
   __ j(equal, &undefined, Label::kNear);
 #ifdef DEBUG
-  // Check that the parent suspender is inactive.
+  // Check that the parent suspender is active.
   Label parent_inactive;
-  Register state = rbx;
+  Register state = tmp2;
   __ LoadTaggedSignedField(
       state, FieldOperand(suspender, WasmSuspenderObject::kStateOffset));
-  __ SmiCompare(state, Smi::FromInt(WasmSuspenderObject::kInactive));
+  __ SmiCompare(state, Smi::FromInt(WasmSuspenderObject::kActive));
   __ j(equal, &parent_inactive, Label::kNear);
   __ Trap();
   __ bind(&parent_inactive);
@@ -3754,7 +3758,7 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
   __ bind(&return_done);
   if (stack_switch) {
     ReloadParentContinuation(masm, wasm_instance, return_reg, rbx, rcx);
-    RestoreParentSuspender(masm);
+    RestoreParentSuspender(masm, rbx, rcx);
   }
   __ bind(&suspend);
   // No need to process the return value if the stack is suspended, there is a

src/runtime/runtime-wasm.cc

@@ -818,10 +818,6 @@ RUNTIME_FUNCTION(Runtime_WasmAllocateContinuation) {
   FullObjectSlot active_suspender_slot =
       isolate->roots_table().slot(RootIndex::kActiveSuspender);
   suspender->set_parent(HeapObject::cast(*active_suspender_slot));
-  if (!(*active_suspender_slot).IsUndefined()) {
-    WasmSuspenderObject::cast(*active_suspender_slot)
-        .set_state(WasmSuspenderObject::kInactive);
-  }
   suspender->set_state(WasmSuspenderObject::kActive);
   suspender->set_continuation(*target);
   active_suspender_slot.store(*suspender);

test/mjsunit/wasm/stack-switching.js

@@ -353,11 +353,12 @@ load("test/mjsunit/wasm/wasm-module-builder.js");
           kExprCallFunction, import_index,
           kExprI32Const, 0
           ]).exportFunc();
-  function js_import() {
-    return Promise.resolve(0);
-  }
-  let instance = builder.instantiate({m: {i: js_import}});
   let suspender = new WebAssembly.Suspender();
+  let i = suspender.suspendOnReturnedPromise(
+      new WebAssembly.Function(
+          {parameters: [], results: ['externref']},
+          () => Promise.resolve(0)));
+  let instance = builder.instantiate({m: {i}});
   let wrapped_export = suspender.returnPromiseOnSuspend(instance.exports.test);
   let promise1 = wrapped_export();
   // Re-enter the suspender before resolving the promise.
@@ -365,14 +366,16 @@ load("test/mjsunit/wasm/wasm-module-builder.js");
       /re-entering an active\/suspended suspender/);
 })();
 
-(function TestNestedSuspenders() {
+function TestNestedSuspenders(suspend) {
   // Nest two suspenders. The call chain looks like:
   // outer (wasm) -> outer (js) -> inner (wasm) -> inner (js)
-  // The inner JS function returns a Promise, which suspends the inner wasm
-  // function, which returns a Promise, which suspends the outer wasm function,
-  // which returns a Promise. The inner Promise resolves first, which resumes
-  // the inner continuation. Then the outer promise resolves which resumes the
-  // outer continuation.
+  // If 'suspend' is true, the inner JS function returns a Promise, which
+  // suspends the inner wasm function, which returns a Promise, which suspends
+  // the outer wasm function, which returns a Promise. The inner Promise
+  // resolves first, which resumes the inner continuation. Then the outer
+  // promise resolves which resumes the outer continuation.
+  // If 'suspend' is false, the inner JS function returns a regular value and
+  // no computation is suspended.
   print(arguments.callee.name);
   let builder = new WasmModuleBuilder();
   inner_index = builder.addImport('m', 'inner', kSig_i_v);
@@ -391,7 +394,7 @@ load("test/mjsunit/wasm/wasm-module-builder.js");
   let inner = inner_suspender.suspendOnReturnedPromise(
       new WebAssembly.Function(
           {parameters: [], results: ['externref']},
-          () => Promise.resolve(42)));
+          () => suspend ? Promise.resolve(42) : 42));
 
   let export_inner;
   let outer = outer_suspender.suspendOnReturnedPromise(
@@ -402,5 +405,39 @@ load("test/mjsunit/wasm/wasm-module-builder.js");
   let instance = builder.instantiate({m: {inner, outer}});
   export_inner = inner_suspender.returnPromiseOnSuspend(instance.exports.inner);
   let export_outer = outer_suspender.returnPromiseOnSuspend(instance.exports.outer);
+  if (suspend) {
     assertPromiseResult(export_outer(), v => assertEquals(42, v));
+  } else {
+    assertEquals(export_outer(), 42);
+  }
+}
+
+(function TestNestedSuspendersSuspend() {
+  TestNestedSuspenders(true);
+})();
+
+(function TestNestedSuspendersNoSuspend() {
+  TestNestedSuspenders(false);
+})();
+
+(function TestReenterInactiveSuspender() {
+  print(arguments.callee.name);
+  let builder = new WasmModuleBuilder();
+  let import_index = builder.addImport("m", "i", kSig_i_v);
+  builder.addFunction("test", kSig_i_v)
+      .addBody([
+          kExprCallFunction, import_index,
+          ]).exportFunc();
+  let suspender = new WebAssembly.Suspender();
+  let i = suspender.suspendOnReturnedPromise(
+      new WebAssembly.Function(
+          {parameters: [], results: ['externref']},
+          () => Promise.resolve(0)));
+  let instance = builder.instantiate({m: {i}});
+  let wrapped_export = suspender.returnPromiseOnSuspend(instance.exports.test);
+  assertPromiseResult(
+      wrapped_export(),
+      () => assertPromiseResult(
+          wrapped_export(),
+          v => assertEquals(v, 0)));
 })();