Bug 1690965 - part 4: add unwind wasm instruction r=rhunt

Adds support for the `unwind` instruction from the wasm exception
handling proposal to the Baseline compiler.

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

js/src/jit-test/lib/wasm-binary.js

@@ -152,7 +152,7 @@ const MozPrefix = 0xff;
 
 const definedOpcodes =
     [0x00, 0x01, 0x02, 0x03, 0x04, 0x05,
-     ...(wasmExceptionsEnabled() ? [0x06, 0x07, 0x08, 0x09] : []),
+     ...(wasmExceptionsEnabled() ? [0x06, 0x07, 0x08, 0x09, 0x0a] : []),
      0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
      0x10, 0x11,
      ...(wasmExceptionsEnabled() ? [0x18, 0x19] : []),

js/src/jit-test/tests/wasm/exceptions/instructions.js

@@ -989,3 +989,123 @@ assertEq(
   ).exports.f(),
   42
 );
+
+// Test try-unwind blocks.
+assertEq(
+  wasmEvalText(
+    `(module
+       (func (export "f") (result i32)
+         try (result i32)
+           i32.const 42
+         unwind
+         end))`
+  ).exports.f(),
+  42
+);
+
+assertEq(
+  wasmEvalText(
+    `(module
+       (func (export "f") (result i32)
+         try (result i32)
+           i32.const 42
+           br 0
+         unwind
+         end))`
+  ).exports.f(),
+  42
+);
+
+assertEq(
+  wasmEvalText(
+    `(module
+       (event $exn)
+       (func (export "f") (result i32)
+         try (result i32)
+           throw $exn
+         unwind
+           i32.const 42
+           br 0
+         end))`
+  ).exports.f(),
+  42
+);
+
+assertEq(
+  wasmEvalText(
+    `(module
+       (event $exn)
+       (func (export "f") (result i32)
+         try (result i32)
+           throw $exn
+         unwind
+           i32.const 42
+           return
+         end))`
+  ).exports.f(),
+  42
+);
+
+assertEq(
+  wasmEvalText(
+    `(module
+       (type (func))
+       (event $exn (type 0))
+       (func (export "f") (result i32) (local i32)
+         try
+           try
+             throw $exn
+           unwind
+             i32.const 1
+             local.set 0
+           end
+         catch $exn
+         end
+         local.get 0))`
+  ).exports.f(),
+  1
+);
+
+// Test the interaction between delegate and unwind.
+assertEq(
+  wasmEvalText(
+    `(module
+       (event $exn)
+       (func (export "f") (result i32) (local i32)
+         try
+           try
+             try
+               throw $exn
+             delegate 0
+           unwind
+             i32.const 42
+             local.set 0
+           end
+         catch_all
+         end
+         local.get 0))`
+  ).exports.f(),
+  42
+);
+
+assertEq(
+  wasmEvalText(
+    `(module
+       (type (func))
+       (event $exn (type 0))
+       (func (export "f") (result i32) (local i32)
+         try $l
+           try
+             try
+               throw $exn
+             delegate 1
+           unwind
+             i32.const 27
+             local.set 0
+           end
+         catch_all
+         end
+         local.get 0))`
+  ).exports.f(),
+  0
+);

js/src/jit-test/tests/wasm/exceptions/throw-to-js.js

@@ -242,3 +242,16 @@ assertWasmThrowsExn(() =>
          delegate 0))`
   ).exports.f()
 );
+
+// Test unwind rethrowing on exit.
+assertWasmThrowsExn(() =>
+  wasmEvalText(
+    `(module
+       (event $exn (param))
+       (func (export "f") (result i32)
+         try (result i32)
+           throw $exn
+         unwind
+         end))`
+  ).exports.f()
+);

js/src/jit-test/tests/wasm/exceptions/validation.js

@@ -641,6 +641,66 @@ function testValidateDelegate() {
   );
 }
 
+function testValidateUnwind() {
+  wasmValidateText(
+    `(module
+       (event $exn (param))
+       (func (local i32)
+         try
+           throw $exn
+         unwind
+           i32.const 1
+           local.set 0
+        end))`
+  );
+
+  wasmValidateText(
+    `(module
+       (event $exn (param))
+       (func (result i32)
+         try (result i32)
+           i32.const 1
+           br 0
+         unwind
+           i32.const 2
+           br 0
+         end))`
+  );
+
+  wasmFailValidateText(
+    `(module
+       (event $exn (param))
+       (func (export "f")
+         try (result i32)
+           (i32.const 1)
+         unwind
+         end))`,
+    /unused values not explicitly dropped by end of block/
+  );
+
+  wasmFailValidateText(
+    `(module
+       (event $exn (param))
+       (func (local i32)
+         try
+           throw $exn
+         unwind
+           (i32.const 1)
+         end))`,
+    /unused values not explicitly dropped by end of block/
+  );
+
+  wasmFailValidateText(
+    `(module (func unwind))`,
+    /unwind can only be used within a try/
+  );
+
+  wasmFailValidateText(
+    `(module (func try unwind rethrow 0 end))`,
+    /rethrow target was not a catch block/
+  );
+}
+
 testValidateDecode();
 testValidateThrow();
 testValidateTryCatch();
@@ -649,3 +709,4 @@ testValidateCatchAll();
 testValidateExnPayload();
 testValidateRethrow();
 testValidateDelegate();
+testValidateUnwind();

js/src/wasm/WasmBaselineCompile.cpp

@@ -8433,6 +8433,7 @@ class BaseCompiler final : public BaseCompilerInterface {
   [[nodiscard]] bool emitCatch();
   [[nodiscard]] bool emitCatchAll();
   [[nodiscard]] bool emitDelegate();
+  [[nodiscard]] bool emitUnwind();
   [[nodiscard]] bool emitThrow();
   [[nodiscard]] bool emitRethrow();
 #endif
@@ -8487,6 +8488,7 @@ class BaseCompiler final : public BaseCompilerInterface {
   [[nodiscard]] bool endIfThenElse(ResultType type);
 #ifdef ENABLE_WASM_EXCEPTIONS
   [[nodiscard]] bool endTryCatch(ResultType type);
+  [[nodiscard]] bool endTryUnwind(ResultType type);
 #endif
 
   void doReturn(ContinuationKind kind);
@@ -10400,6 +10402,11 @@ bool BaseCompiler::emitEnd() {
         return false;
       }
       break;
+    case LabelKind::Unwind:
+      if (!endTryUnwind(type)) {
+        return false;
+      }
+      break;
 #endif
   }
 
@@ -10908,6 +10915,57 @@ bool BaseCompiler::emitDelegate() {
   return pushBlockResults(resultType);
 }
 
+bool BaseCompiler::emitUnwind() {
+  ResultType resultType;
+  NothingVector unused_tryValues;
+
+  if (!iter_.readUnwind(&resultType, &unused_tryValues)) {
+    return false;
+  }
+
+  Control& tryUnwind = controlItem();
+
+  // End the try branch like the first catch block in a try-catch.
+  if (deadCode_) {
+    fr.resetStackHeight(tryUnwind.stackHeight, resultType);
+    popValueStackTo(tryUnwind.stackSize);
+  } else {
+    MOZ_ASSERT(stk_.length() == tryUnwind.stackSize + resultType.length());
+    popBlockResults(resultType, tryUnwind.stackHeight, ContinuationKind::Jump);
+    freeResultRegisters(resultType);
+    masm.jump(&tryUnwind.label);
+    MOZ_ASSERT(!tryUnwind.bceSafeOnExit);
+    MOZ_ASSERT(!tryUnwind.deadOnArrival);
+  }
+
+  deadCode_ = tryUnwind.deadOnArrival;
+
+  if (deadCode_) {
+    return true;
+  }
+
+  bceSafe_ = 0;
+
+  // Create an exception landing pad for the unwind instructions.
+  //
+  // We bind otherLabel so that `delegate` can jump here as well.
+  masm.bind(&tryUnwind.otherLabel);
+  fr.setStackHeight(tryUnwind.stackHeight);
+
+  // Push the exception reference so that the end of the `unwind` can throw.
+  RegRef exn = RegRef(WasmExceptionReg);
+  needRef(exn);
+  pushRef(exn);
+
+  WasmTryNoteVector& tryNotes = masm.tryNotes();
+  WasmTryNote& tryNote = tryNotes[tryUnwind.tryNoteIndex];
+  tryNote.end = masm.currentOffset();
+  tryNote.entryPoint = tryNote.end;
+  tryNote.framePushed = masm.framePushed();
+
+  return true;
+}
+
 bool BaseCompiler::endTryCatch(ResultType type) {
   uint32_t lineOrBytecode = readCallSiteLineOrBytecode();
 
@@ -11016,6 +11074,45 @@ bool BaseCompiler::endTryCatch(ResultType type) {
   return pushBlockResults(type);
 }
 
+bool BaseCompiler::endTryUnwind(ResultType type) {
+  uint32_t lineOrBytecode = readCallSiteLineOrBytecode();
+  Control& tryUnwind = controlItem();
+
+  if (deadCode_) {
+    fr.resetStackHeight(tryUnwind.stackHeight, type);
+    popValueStackTo(tryUnwind.stackSize);
+  } else {
+    // An unwind should have no results so only the exception reference will
+    // remain here on top of the original stack.
+    MOZ_ASSERT(stk_.length() == tryUnwind.stackSize + 1);
+
+    RegRef exn = popRef();
+
+    if (!throwFrom(exn, lineOrBytecode)) {
+      return false;
+    }
+    MOZ_ASSERT(stk_.length() == tryUnwind.stackSize);
+    MOZ_ASSERT(!tryUnwind.bceSafeOnExit);
+    MOZ_ASSERT(!tryUnwind.deadOnArrival);
+  }
+
+  deadCode_ = tryUnwind.deadOnArrival;
+
+  if (deadCode_) {
+    return true;
+  }
+
+  // The try branch may jump here.
+  if (tryUnwind.label.used()) {
+    masm.bind(&tryUnwind.label);
+  }
+
+  captureResultRegisters(type);
+  bceSafe_ = tryUnwind.bceSafeOnExit;
+
+  return pushBlockResults(type);
+}
+
 bool BaseCompiler::emitThrow() {
   uint32_t lineOrBytecode = readCallSiteLineOrBytecode();
   uint32_t exnIndex;
@@ -16038,6 +16135,11 @@ bool BaseCompiler::emitBody() {
         CHECK(emitDelegate());
         iter_.popDelegate();
         NEXT();
+      case uint16_t(Op::Unwind):
+        if (!moduleEnv_.exceptionsEnabled()) {
+          return iter_.unrecognizedOpcode(&op);
+        }
+        CHECK_NEXT(emitUnwind());
       case uint16_t(Op::Throw):
         if (!moduleEnv_.exceptionsEnabled()) {
           return iter_.unrecognizedOpcode(&op);

js/src/wasm/WasmConstants.h

@@ -220,6 +220,7 @@ enum class Op {
   Catch = 0x07,
   Throw = 0x08,
   Rethrow = 0x09,
+  Unwind = 0x0a,
 #endif
   End = 0x0b,
   Br = 0x0c,

js/src/wasm/WasmIonCompile.cpp

@@ -2478,6 +2478,9 @@ static bool EmitEnd(FunctionCompiler& f) {
     case LabelKind::CatchAll:
       MOZ_CRASH("NYI");
       break;
+    case LabelKind::Unwind:
+      MOZ_CRASH("NYI");
+      break;
 #endif
   }
 
@@ -2603,6 +2606,17 @@ static bool EmitDelegate(FunctionCompiler& f) {
   MOZ_CRASH("NYI");
 }
 
+static bool EmitUnwind(FunctionCompiler& f) {
+  ResultType resultType;
+  DefVector tryValues;
+
+  if (!f.iter().readUnwind(&resultType, &tryValues)) {
+    return false;
+  }
+
+  MOZ_CRASH("NYI");
+}
+
 static bool EmitThrow(FunctionCompiler& f) {
   uint32_t exnIndex;
   DefVector argValues;
@@ -4555,6 +4569,11 @@ static bool EmitBodyExprs(FunctionCompiler& f) {
           return false;
         }
         break;
+      case uint16_t(Op::Unwind):
+        if (!f.moduleEnv().exceptionsEnabled()) {
+          return f.iter().unrecognizedOpcode(&op);
+        }
+        CHECK(EmitUnwind(f));
       case uint16_t(Op::Throw):
         if (!f.moduleEnv().exceptionsEnabled()) {
           return f.iter().unrecognizedOpcode(&op);

js/src/wasm/WasmOpIter.cpp

@@ -277,6 +277,8 @@ OpKind wasm::Classify(OpBytes op) {
       WASM_EXN_OP(OpKind::CatchAll);
     case Op::Delegate:
       WASM_EXN_OP(OpKind::Delegate);
+    case Op::Unwind:
+      WASM_EXN_OP(OpKind::Unwind);
     case Op::Throw:
       WASM_EXN_OP(OpKind::Throw);
     case Op::Rethrow:

js/src/wasm/WasmOpIter.h

@@ -42,7 +42,8 @@ enum class LabelKind : uint8_t {
 #ifdef ENABLE_WASM_EXCEPTIONS
   Try,
   Catch,
-  CatchAll
+  CatchAll,
+  Unwind
 #endif
 };
 
@@ -205,6 +206,7 @@ enum class OpKind {
   Catch,
   CatchAll,
   Delegate,
+  Unwind,
   Throw,
   Rethrow,
   Try,
@@ -279,6 +281,12 @@ class ControlStackEntry {
     kind_ = LabelKind::CatchAll;
     polymorphicBase_ = false;
   }
+
+  void switchToUnwind() {
+    MOZ_ASSERT(kind() == LabelKind::Try);
+    kind_ = LabelKind::Unwind;
+    polymorphicBase_ = false;
+  }
 #endif
 };
 
@@ -488,6 +496,8 @@ class MOZ_STACK_CLASS OpIter : private Policy {
                                   ResultType* resultType,
                                   ValueVector* tryResults);
   void popDelegate();
+  [[nodiscard]] bool readUnwind(ResultType* resultType,
+                                ValueVector* tryResults);
   [[nodiscard]] bool readThrow(uint32_t* eventIndex, ValueVector* argValues);
   [[nodiscard]] bool readRethrow(uint32_t* relativeDepth);
 #endif
@@ -1231,11 +1241,32 @@ inline bool OpIter<Policy>::readEnd(LabelKind* kind, ResultType* type,
                                     ValueVector* resultsForEmptyElse) {
   MOZ_ASSERT(Classify(op_) == OpKind::End);
 
+  Control& block = controlStack_.back();
+
+#ifdef ENABLE_WASM_EXCEPTIONS
+  if (block.kind() == LabelKind::Try) {
+    return fail("try without catch or unwind not allowed");
+  }
+
+  // Unwind blocks require an empty result type rather than the try's type.
+  if (block.kind() == LabelKind::Unwind) {
+    if (valueStack_.length() != block.valueStackBase()) {
+      return fail("unused values not explicitly dropped by end of block");
+    }
+    *type = block.type().results();
+    // As the `end` for an `unwind` always rethrows the exception, we need
+    // to push the try block's type here to the value stack.
+    if (!push(*type)) {
+      return false;
+    }
+  } else if (!checkStackAtEndOfBlock(type, results)) {
+    return false;
+  }
+#else
   if (!checkStackAtEndOfBlock(type, results)) {
     return false;
   }
-
-  Control& block = controlStack_.back();
+#endif
 
   if (block.kind() == LabelKind::Then) {
     ResultType params = block.type().params();
@@ -1258,12 +1289,6 @@ inline bool OpIter<Policy>::readEnd(LabelKind* kind, ResultType* type,
     elseParamStack_.shrinkBy(nparams);
   }
 
-#ifdef ENABLE_WASM_EXCEPTIONS
-  if (block.kind() == LabelKind::Try) {
-    return fail("try without catch or unwind not allowed");
-  }
-#endif
-
   *kind = block.kind();
   return true;
 }
@@ -1572,6 +1597,26 @@ inline void OpIter<Policy>::popDelegate() {
   controlStack_.popBack();
 }
 
+template <typename Policy>
+inline bool OpIter<Policy>::readUnwind(ResultType* resultType,
+                                       ValueVector* tryResults) {
+  MOZ_ASSERT(Classify(op_) == OpKind::Unwind);
+
+  Control& block = controlStack_.back();
+  if (block.kind() != LabelKind::Try) {
+    return fail("unwind can only be used within a try");
+  }
+
+  if (!checkStackAtEndOfBlock(resultType, tryResults)) {
+    return false;
+  }
+
+  valueStack_.shrinkTo(block.valueStackBase());
+  block.switchToUnwind();
+
+  return true;
+}
+
 template <typename Policy>
 inline bool OpIter<Policy>::readThrow(uint32_t* eventIndex,
                                       ValueVector* argValues) {

js/src/wasm/WasmValidate.cpp

@@ -1510,6 +1510,11 @@ static bool DecodeFunctionBodyExprs(const ModuleEnvironment& env,
         iter.popDelegate();
         break;
       }
+      case uint16_t(Op::Unwind):
+        if (!env.exceptionsEnabled()) {
+          return iter.unrecognizedOpcode(&op);
+        }
+        CHECK(iter.readUnwind(&unusedType, &nothings));
       case uint16_t(Op::Throw): {
         if (!env.exceptionsEnabled()) {
           return iter.unrecognizedOpcode(&op);