From 538d334aa9f75b82f4468473bc49db41c5449469 Mon Sep 17 00:00:00 2001
From: Thomas Lively <tlively@google.com>
Date: Fri, 22 Dec 2023 13:16:54 -0800
Subject: [PATCH] [Parser] Parse br_on_cast{_fail} input annotations

And validate in IRBuilder both that the input annotation is valid and that the
input matches it.
---
 src/parser/contexts.h          | 12 ++++++++----
 src/parser/parsers.h           |  8 +++++---
 src/wasm-ir-builder.h          |  2 +-
 src/wasm/wasm-ir-builder.cpp   | 12 ++++++++++--
 test/lit/wat-kitchen-sink.wast |  4 ++--
 5 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/src/parser/contexts.h b/src/parser/contexts.h
index 2565a4817dd..1cefb288da3 100644
--- a/src/parser/contexts.h
+++ b/src/parser/contexts.h
@@ -443,7 +443,8 @@ struct NullInstrParserCtx {
 
   Result<> makeBrOn(Index, LabelIdxT, BrOnOp) { return Ok{}; }
 
-  template<typename TypeT> Result<> makeBrOn(Index, LabelIdxT, BrOnOp, TypeT) {
+  template<typename TypeT>
+  Result<> makeBrOn(Index, LabelIdxT, BrOnOp, TypeT, TypeT) {
     return Ok{};
   }
 
@@ -1690,9 +1691,12 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
     return withLoc(pos, irBuilder.makeRefCast(type));
   }
 
-  Result<>
-  makeBrOn(Index pos, Index label, BrOnOp op, Type castType = Type::none) {
-    return withLoc(pos, irBuilder.makeBrOn(label, op, castType));
+  Result<> makeBrOn(Index pos,
+                    Index label,
+                    BrOnOp op,
+                    Type in = Type::none,
+                    Type out = Type::none) {
+    return withLoc(pos, irBuilder.makeBrOn(label, op, in, out));
   }
 
   Result<> makeStructNew(Index pos, HeapType type) {
diff --git a/src/parser/parsers.h b/src/parser/parsers.h
index 9302f4fc9a9..1d786f363d1 100644
--- a/src/parser/parsers.h
+++ b/src/parser/parsers.h
@@ -1594,9 +1594,11 @@ template<typename Ctx> Result<> makeBrOnNull(Ctx& ctx, Index pos, bool onFail) {
 template<typename Ctx> Result<> makeBrOnCast(Ctx& ctx, Index pos, bool onFail) {
   auto label = labelidx(ctx);
   CHECK_ERR(label);
-  auto type = reftype(ctx);
-  CHECK_ERR(type);
-  return ctx.makeBrOn(pos, *label, onFail ? BrOnCastFail : BrOnCast, *type);
+  auto in = reftype(ctx);
+  CHECK_ERR(in);
+  auto out = reftype(ctx);
+  CHECK_ERR(out);
+  return ctx.makeBrOn(pos, *label, onFail ? BrOnCastFail : BrOnCast, *in, *out);
 }
 
 template<typename Ctx>
diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h
index dcb25efa785..5add042a63b 100644
--- a/src/wasm-ir-builder.h
+++ b/src/wasm-ir-builder.h
@@ -168,7 +168,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   [[nodiscard]] Result<> makeRefTest(Type type);
   [[nodiscard]] Result<> makeRefCast(Type type);
   [[nodiscard]] Result<>
-  makeBrOn(Index label, BrOnOp op, Type castType = Type::none);
+  makeBrOn(Index label, BrOnOp op, Type in = Type::none, Type out = Type::none);
   [[nodiscard]] Result<> makeStructNew(HeapType type);
   [[nodiscard]] Result<> makeStructNewDefault(HeapType type);
   [[nodiscard]] Result<>
diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index 3cfe0ea30c2..43a989c6e0f 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -1420,12 +1420,20 @@ Result<> IRBuilder::makeRefCast(Type type) {
   return Ok{};
 }
 
-Result<> IRBuilder::makeBrOn(Index label, BrOnOp op, Type castType) {
+Result<> IRBuilder::makeBrOn(Index label, BrOnOp op, Type in, Type out) {
   BrOn curr;
   CHECK_ERR(visitBrOn(&curr));
+  if (out != Type::none) {
+    if (!Type::isSubType(out, in)) {
+      return Err{"output type is not a subtype of the input type"};
+    }
+    if (!Type::isSubType(curr.ref->type, in)) {
+      return Err{"expected input to match input type annotation"};
+    }
+  }
   auto name = getLabelName(label);
   CHECK_ERR(name);
-  push(builder.makeBrOn(op, *name, curr.ref, castType));
+  push(builder.makeBrOn(op, *name, curr.ref, out));
   return Ok{};
 }
 
diff --git a/test/lit/wat-kitchen-sink.wast b/test/lit/wat-kitchen-sink.wast
index 303cfc53b0c..775a0a139db 100644
--- a/test/lit/wat-kitchen-sink.wast
+++ b/test/lit/wat-kitchen-sink.wast
@@ -3549,7 +3549,7 @@
   block (result i31ref)
    block (result (ref any))
     local.get 0
-    br_on_cast 1 i31ref
+    br_on_cast 1 anyref i31ref
    end
    unreachable
   end
@@ -3574,7 +3574,7 @@
   block (result (ref any))
    block (result i31ref)
     local.get 0
-    br_on_cast_fail 1 i31ref
+    br_on_cast_fail 1 anyref i31ref
    end
    unreachable
   end