Bug 1690965 - part 2: add rethrow wasm instruction r=rhunt

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

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

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

@@ -55,6 +55,7 @@ const BlockCode        = 0x02;
 const TryCode          = 0x06;
 const CatchCode        = 0x07;
 const ThrowCode        = 0x08;
+const RethrowCode      = 0x09;
 const EndCode          = 0x0b;
 const ReturnCode       = 0x0f;
 const CallCode         = 0x10;
@@ -150,7 +151,7 @@ const MozPrefix = 0xff;
 
 const definedOpcodes =
     [0x00, 0x01, 0x02, 0x03, 0x04, 0x05,
-     ...(wasmExceptionsEnabled() ? [0x06, 0x07, 0x08] : []),
+     ...(wasmExceptionsEnabled() ? [0x06, 0x07, 0x08, 0x09] : []),
      0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
      0x10, 0x11,
      ...(wasmExceptionsEnabled() ? [0x19] : []),

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

@@ -294,6 +294,40 @@ assertEq(
   2
 );
 
+assertEq(
+  wasmEvalText(
+    `(module
+       (type (func (param)))
+       (event $exn (type 0))
+       (func (export "f") (result i32)
+         try $l (result i32)
+           (throw $exn)
+         catch $exn
+           (i32.const 2)
+           (br $l)
+           rethrow 0
+         end))`
+  ).exports.f(),
+  2
+);
+
+assertEq(
+  wasmEvalText(
+    `(module
+       (type (func (param)))
+       (event $exn (type 0))
+       (func (export "f") (result i32)
+         try $l (result i32)
+           (throw $exn)
+         catch_all
+           (i32.const 2)
+           (br $l)
+           rethrow 0
+         end))`
+  ).exports.f(),
+  2
+);
+
 // Test br branching out of a catch block.
 assertEq(
   wasmEvalText(
@@ -664,3 +698,143 @@ assertErrorMessage(
     102
   );
 }
+
+// Test simple rethrow.
+assertEq(
+  wasmEvalText(
+    `(module
+       (type (func))
+       (event $exn (type 0))
+       (func (export "f") (result i32)
+         try (result i32)
+           try
+             throw $exn
+           catch $exn
+             rethrow 0
+           end
+           i32.const 1
+         catch $exn
+           i32.const 27
+         end))`
+  ).exports.f(),
+  27
+);
+
+assertEq(
+  wasmEvalText(
+    `(module
+       (type (func))
+       (event $exn (type 0))
+       (func (export "f") (result i32)
+         try (result i32)
+           try
+             throw $exn
+           catch_all
+             rethrow 0
+           end
+           i32.const 1
+         catch $exn
+           i32.const 27
+         end))`
+  ).exports.f(),
+  27
+);
+
+// Test rethrows in nested blocks.
+assertEq(
+  wasmEvalText(
+    `(module
+       (type (func))
+       (event $exn (type 0))
+       (func (export "f") (result i32)
+         try (result i32)
+           try
+             throw $exn
+           catch $exn
+             block
+               rethrow 1
+             end
+           end
+           i32.const 1
+         catch $exn
+           i32.const 27
+         end))`
+  ).exports.f(),
+  27
+);
+
+assertEq(
+  wasmEvalText(
+    `(module
+       (type (func))
+       (event $exn (type 0))
+       (func (export "f") (result i32)
+         try (result i32)
+           try
+             throw $exn
+           catch_all
+             block
+               rethrow 1
+             end
+           end
+           i32.const 1
+         catch $exn
+           i32.const 27
+         end))`
+  ).exports.f(),
+  27
+);
+
+assertEq(
+  wasmEvalText(
+    `(module
+       (type (func))
+       (event $exn1 (type 0))
+       (event $exn2 (type 0))
+       (func (export "f") (result i32)
+         try (result i32)
+           try
+             throw $exn1
+           catch $exn1
+             try
+               throw $exn2
+             catch $exn2
+               rethrow 1
+             end
+           end
+           i32.const 0
+         catch $exn1
+           i32.const 1
+         catch $exn2
+           i32.const 2
+         end))`
+  ).exports.f(),
+  1
+);
+
+assertEq(
+  wasmEvalText(
+    `(module
+       (type (func))
+       (event $exn1 (type 0))
+       (event $exn2 (type 0))
+       (func (export "f") (result i32)
+         try (result i32)
+           try
+             throw $exn1
+           catch $exn1
+             try
+               throw $exn2
+             catch_all
+               rethrow 1
+             end
+           end
+           i32.const 0
+         catch $exn1
+           i32.const 1
+         catch $exn2
+           i32.const 2
+         end))`
+  ).exports.f(),
+  1
+);

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

@@ -49,6 +49,29 @@ assertThrowsValue(
   42
 );
 
+// Like previous test, but using a rethrow instruction instead.
+assertThrowsValue(
+  () =>
+    wasmEvalText(
+      `(module
+         (import "m" "import" (func $import))
+         (func (export "f")
+           try
+             (call $import)
+           catch_all
+             (rethrow 0)
+           end))`,
+      {
+        m: {
+          import: () => {
+            throw 42;
+          },
+        },
+      }
+    ).exports.f(),
+  42
+);
+
 // Test for throwing to JS and then back to Wasm.
 {
   var wasmThrower;

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

@@ -99,6 +99,28 @@ function testValidateDecode() {
     ]),
     /try without catch or unwind not allowed/
   );
+
+  // Rethrow must have a depth argument.
+  wasmInvalid(
+    moduleWithSections([
+      sigSection([emptyType]),
+      declSection([0]),
+      eventSection([{ type: 0 }]),
+      bodySection([
+        funcBody(
+          {
+            locals: [],
+            body: [
+              RethrowCode,
+              // Index missing.
+            ],
+          },
+          (withEndCode = false)
+        ),
+      ]),
+    ]),
+    /unable to read rethrow depth/
+  );
 }
 
 function testValidateThrow() {
@@ -392,9 +414,120 @@ function testValidateExnPayload() {
   wasmInvalid(invalid1, /event index out of range/);
 }
 
+function testValidateRethrow() {
+  wasmValidateText(
+    `(module
+       (event $exn (param))
+       (func
+         try
+           nop
+         catch $exn
+           rethrow 0
+         end))`
+  );
+
+  wasmValidateText(
+    `(module
+       (event $exn (param))
+       (func
+         try
+           nop
+         catch_all
+           rethrow 0
+         end))`
+  );
+
+  wasmValidateText(
+    `(module
+       (func (result i32)
+         try (result i32)
+           (i32.const 1)
+         catch_all
+           rethrow 0
+         end))`
+  );
+
+  wasmValidateText(
+    `(module
+       (event $exn (param))
+       (func
+         try
+           nop
+         catch $exn
+           block
+             try
+             catch $exn
+               rethrow 0
+             end
+           end
+         end))`
+  );
+
+  wasmValidateText(
+    `(module
+       (event $exn (param))
+       (func
+         try
+           nop
+         catch $exn
+           block
+             try
+             catch $exn
+               rethrow 2
+             end
+           end
+         end))`
+  );
+
+  wasmFailValidateText(
+    `(module
+       (event $exn (param))
+       (func
+         try
+           nop
+         catch $exn
+           block
+             try
+             catch $exn
+               rethrow 1
+             end
+           end
+         end))`,
+    /rethrow target was not a catch block/
+  );
+
+  wasmFailValidateText(
+    `(module (func rethrow 0))`,
+    /rethrow target was not a catch block/
+  );
+
+  wasmFailValidateText(
+    `(module (func try rethrow 0 catch_all end))`,
+    /rethrow target was not a catch block/
+  );
+
+  wasmFailValidateText(
+    `(module
+       (event $exn (param))
+       (func
+         try
+           nop
+         catch $exn
+           block
+             try
+             catch $exn
+               rethrow 4
+             end
+           end
+         end))`,
+    /rethrow depth exceeds current nesting level/
+  );
+}
+
 testValidateDecode();
 testValidateThrow();
 testValidateTryCatch();
 testValidateCatch();
 testValidateCatchAll();
 testValidateExnPayload();
+testValidateRethrow();

js/src/wasm/WasmBaselineCompile.cpp

@@ -4086,6 +4086,13 @@ class BaseCompiler final : public BaseCompilerInterface {
     }
   }
 
+  // Duplicate the reference at the specified depth and load it into a register.
+  void dupRefAt(uint32_t depth, RegRef dest) {
+    MOZ_ASSERT(depth < stk_.length());
+    Stk& src = peek(stk_.length() - depth - 1);
+    loadRef(src, dest);
+  }
+
   // Flush all local and register value stack elements to memory.
   //
   // TODO / OPTIMIZE: As this is fairly expensive and causes worse
@@ -5199,6 +5206,31 @@ class BaseCompiler final : public BaseCompilerInterface {
     }
   }
 
+#ifdef ENABLE_WASM_EXCEPTIONS
+  // This function is similar to popBlockResults, but additionally handles the
+  // implicit exception pointer that is pushed to the value stack on entry to
+  // a catch handler by dropping it appropriately.
+  void popCatchResults(ResultType type, StackHeight stackBase) {
+    if (!type.empty()) {
+      ABIResultIter iter(type);
+      popRegisterResults(iter);
+      if (!iter.done()) {
+        popStackResults(iter, stackBase);
+        // Since popStackResults clobbers the stack, we only need to free the
+        // exception off of the value stack.
+        popValueStackBy(1);
+      } else {
+        // If there are no stack results, we have to adjust the stack by
+        // dropping the exception reference that's now on the stack.
+        dropValue();
+      }
+    } else {
+      dropValue();
+    }
+    fr.popStackBeforeBranch(stackBase, type);
+  }
+#endif
+
   Stk captureStackResult(const ABIResult& result, StackHeight resultsBase,
                          uint32_t stackResultBytes) {
     MOZ_ASSERT(result.onStack());
@@ -8398,6 +8430,7 @@ class BaseCompiler final : public BaseCompilerInterface {
   [[nodiscard]] bool emitCatch();
   [[nodiscard]] bool emitCatchAll();
   [[nodiscard]] bool emitThrow();
+  [[nodiscard]] bool emitRethrow();
 #endif
   [[nodiscard]] bool emitEnd();
   [[nodiscard]] bool emitBr();
@@ -10574,9 +10607,17 @@ void BaseCompiler::emitCatchSetup(LabelKind kind, Control& tryCatch,
     fr.resetStackHeight(tryCatch.stackHeight, resultType);
     popValueStackTo(tryCatch.stackSize);
   } else {
-    MOZ_ASSERT(stk_.length() == tryCatch.stackSize + resultType.length());
+    // If the previous block is a catch, we need to handle the extra exception
+    // reference on the stack (for rethrow) and thus the stack size is 1 more.
+    MOZ_ASSERT(stk_.length() == tryCatch.stackSize + resultType.length() +
+                                    (kind == LabelKind::Try ? 0 : 1));
     // Try jumps to the end of the try-catch block unless a throw is done.
-    popBlockResults(resultType, tryCatch.stackHeight, ContinuationKind::Jump);
+    if (kind == LabelKind::Try) {
+      popBlockResults(resultType, tryCatch.stackHeight, ContinuationKind::Jump);
+    } else {
+      popCatchResults(resultType, tryCatch.stackHeight);
+    }
+    MOZ_ASSERT(stk_.length() == tryCatch.stackSize);
     freeResultRegisters(resultType);
     MOZ_ASSERT(!tryCatch.deadOnArrival);
   }
@@ -10671,7 +10712,9 @@ bool BaseCompiler::emitCatch() {
 #  endif
   masm.loadPtr(Address(refs, NativeObject::offsetOfElements()), refs);
 
-  freeRef(exn);
+  // This reference is pushed onto the stack because a potential rethrow
+  // may need to access it. It is always popped at the end of the block.
+  pushRef(exn);
 
   masm.loadPtr(Address(values, dataOffset), values);
   size_t argOffset = 0;
@@ -10764,6 +10807,14 @@ bool BaseCompiler::emitCatchAll() {
 
   masm.bind(&tryCatch.catchInfos.back().label);
 
+  // The code in the landing pad guarantees us that the exception reference
+  // is live in this register.
+  RegRef exn = RegRef(WasmExceptionReg);
+  needRef(exn);
+  // This reference is pushed onto the stack because a potential rethrow
+  // may need to access it. It is always popped at the end of the block.
+  pushRef(exn);
+
   return true;
 }
 
@@ -10776,10 +10827,13 @@ bool BaseCompiler::endTryCatch(ResultType type) {
     fr.resetStackHeight(tryCatch.stackHeight, type);
     popValueStackTo(tryCatch.stackSize);
   } else {
-    MOZ_ASSERT(stk_.length() == tryCatch.stackSize + type.length());
+    // Since the previous block is a catch, we must handle the extra exception
+    // reference on the stack (for rethrow) and thus the stack size is 1 more.
+    MOZ_ASSERT(stk_.length() == tryCatch.stackSize + type.length() + 1);
     // Assume we have a control join, so place results in block result
     // allocations and also handle the implicit exception reference.
-    popBlockResults(type, tryCatch.stackHeight, ContinuationKind::Jump);
+    popCatchResults(type, tryCatch.stackHeight);
+    MOZ_ASSERT(stk_.length() == tryCatch.stackSize);
     // Since we will emit a landing pad after this and jump over it to get to
     // the control join, we free these here and re-capture at the join.
     freeResultRegisters(type);
@@ -10987,6 +11041,25 @@ bool BaseCompiler::emitThrow() {
 
   return throwFrom(exn, lineOrBytecode);
 }
+
+bool BaseCompiler::emitRethrow() {
+  uint32_t lineOrBytecode = readCallSiteLineOrBytecode();
+
+  uint32_t relativeDepth;
+  if (!iter_.readRethrow(&relativeDepth)) {
+    return false;
+  }
+
+  if (deadCode_) {
+    return true;
+  }
+
+  Control& tryCatch = controlItem(relativeDepth);
+  RegRef exn = needRef();
+  dupRefAt(tryCatch.stackSize, exn);
+
+  return throwFrom(exn, lineOrBytecode);
+}
 #endif
 
 bool BaseCompiler::emitDrop() {
@@ -15871,6 +15944,11 @@ bool BaseCompiler::emitBody() {
           return iter_.unrecognizedOpcode(&op);
         }
         CHECK_NEXT(emitThrow());
+      case uint16_t(Op::Rethrow):
+        if (!moduleEnv_.exceptionsEnabled()) {
+          return iter_.unrecognizedOpcode(&op);
+        }
+        CHECK_NEXT(emitRethrow());
 #endif
       case uint16_t(Op::Br):
         CHECK_NEXT(emitBr());

js/src/wasm/WasmConstants.h

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

js/src/wasm/WasmIonCompile.cpp

@@ -2600,6 +2600,15 @@ static bool EmitThrow(FunctionCompiler& f) {
 
   MOZ_CRASH("NYI");
 }
+
+static bool EmitRethrow(FunctionCompiler& f) {
+  uint32_t relativeDepth;
+  if (!f.iter().readRethrow(&relativeDepth)) {
+    return false;
+  }
+
+  MOZ_CRASH("NYI");
+}
 #endif
 
 static bool EmitCallArgs(FunctionCompiler& f, const FuncType& funcType,
@@ -4531,6 +4540,11 @@ static bool EmitBodyExprs(FunctionCompiler& f) {
           return f.iter().unrecognizedOpcode(&op);
         }
         CHECK(EmitThrow(f));
+      case uint16_t(Op::Rethrow):
+        if (!f.moduleEnv().exceptionsEnabled()) {
+          return f.iter().unrecognizedOpcode(&op);
+        }
+        CHECK(EmitRethrow(f));
 #endif
       case uint16_t(Op::Br):
         CHECK(EmitBr(f));

js/src/wasm/WasmOpIter.cpp

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

js/src/wasm/WasmOpIter.h

@@ -205,6 +205,7 @@ enum class OpKind {
   Catch,
   CatchAll,
   Throw,
+  Rethrow,
   Try,
 #  endif
 };
@@ -483,6 +484,7 @@ class MOZ_STACK_CLASS OpIter : private Policy {
                                   ResultType* resultType,
                                   ValueVector* tryResults);
   [[nodiscard]] bool readThrow(uint32_t* eventIndex, ValueVector* argValues);
+  [[nodiscard]] bool readRethrow(uint32_t* relativeDepth);
 #endif
   [[nodiscard]] bool readUnreachable();
   [[nodiscard]] bool readDrop();
@@ -664,6 +666,11 @@ class MOZ_STACK_CLASS OpIter : private Policy {
         .controlItem();
   }
 
+  // Return the LabelKind of an element in the control stack.
+  LabelKind controlKind(uint32_t relativeDepth) {
+    return controlStack_[controlStack_.length() - 1 - relativeDepth].kind();
+  }
+
   // Return a reference to the outermost element on the control stack.
   ControlItem& controlOutermost() { return controlStack_[0].controlItem(); }
 
@@ -1538,6 +1545,26 @@ inline bool OpIter<Policy>::readThrow(uint32_t* eventIndex,
   afterUnconditionalBranch();
   return true;
 }
+
+template <typename Policy>
+inline bool OpIter<Policy>::readRethrow(uint32_t* relativeDepth) {
+  MOZ_ASSERT(Classify(op_) == OpKind::Rethrow);
+
+  if (!readVarU32(relativeDepth)) {
+    return fail("unable to read rethrow depth");
+  }
+
+  if (*relativeDepth >= controlStack_.length()) {
+    return fail("rethrow depth exceeds current nesting level");
+  }
+  LabelKind kind = controlKind(*relativeDepth);
+  if (kind != LabelKind::Catch && kind != LabelKind::CatchAll) {
+    return fail("rethrow target was not a catch block");
+  }
+
+  afterUnconditionalBranch();
+  return true;
+}
 #endif
 
 template <typename Policy>

js/src/wasm/WasmValidate.cpp

@@ -1506,6 +1506,13 @@ static bool DecodeFunctionBodyExprs(const ModuleEnvironment& env,
         uint32_t unusedIndex;
         CHECK(iter.readThrow(&unusedIndex, &nothings));
       }
+      case uint16_t(Op::Rethrow): {
+        if (!env.exceptionsEnabled()) {
+          return iter.unrecognizedOpcode(&op);
+        }
+        uint32_t unusedDepth;
+        CHECK(iter.readRethrow(&unusedDepth));
+      }
 #endif
       case uint16_t(Op::ThreadPrefix): {
         if (env.sharedMemoryEnabled() == Shareable::False) {