3 files changed, 52 insertions, 17 deletions

diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 9540bd8f3..1de5e4bd1 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -3783,7 +3783,7 @@ BinaryConsts::ASTNodes WasmBinaryBuilder::readExpression(Expression*& curr) {
       auto call = allocator.alloc<CallRef>();
       call->isReturn = true;
       curr = call;
-      visitCallRef(call);
+      visitCallRef(call, getTypeByIndex(getU32LEB()));
       break;
     }
     case BinaryConsts::AtomicPrefix: {
@@ -6777,22 +6777,32 @@ void WasmBinaryBuilder::visitRethrow(Rethrow* curr) {
   curr->finalize();
 }
 
-void WasmBinaryBuilder::visitCallRef(CallRef* curr) {
+void WasmBinaryBuilder::visitCallRef(CallRef* curr,
+                                     std::optional<HeapType> maybeType) {
   BYN_TRACE("zz node: CallRef\n");
   curr->target = popNonVoidExpression();
-  auto type = curr->target->type;
-  if (type == Type::unreachable) {
-    // If our input is unreachable, then we cannot even find out how many inputs
-    // we have, and just set ourselves to unreachable as well.
-    curr->finalize(type);
-    return;
-  }
-  if (!type.isRef()) {
-    throwError("Non-ref type for a call_ref: " + type.toString());
+  HeapType heapType;
+  if (maybeType) {
+    heapType = *maybeType;
+    if (!Type::isSubType(curr->target->type, Type(heapType, Nullable))) {
+      throwError("Call target has invalid type: " +
+                 curr->target->type.toString());
+    }
+  } else {
+    auto type = curr->target->type;
+    if (type == Type::unreachable) {
+      // If our input is unreachable, then we cannot even find out how many
+      // inputs we have, and just set ourselves to unreachable as well.
+      curr->finalize(type);
+      return;
+    }
+    if (!type.isRef()) {
+      throwError("Non-ref type for a call_ref: " + type.toString());
+    }
+    heapType = type.getHeapType();
   }
-  auto heapType = type.getHeapType();
   if (!heapType.isSignature()) {
-    throwError("Invalid reference type for a call_ref: " + type.toString());
+    throwError("Invalid reference type for a call_ref: " + heapType.toString());
   }
   auto sig = heapType.getSignature();
   auto num = sig.params.size();
@@ -6800,7 +6810,11 @@ void WasmBinaryBuilder::visitCallRef(CallRef* curr) {
   for (size_t i = 0; i < num; i++) {
     curr->operands[num - i - 1] = popNonVoidExpression();
   }
-  curr->finalize(sig.results);
+  if (maybeType) {
+    curr->finalize();
+  } else {
+    curr->finalize(sig.results);
+  }
 }
 
 bool WasmBinaryBuilder::maybeVisitI31New(Expression*& out, uint32_t code) {
diff --git a/src/wasm/wasm-s-parser.cpp b/src/wasm/wasm-s-parser.cpp
index ef41dec93..cf4824323 100644
--- a/src/wasm/wasm-s-parser.cpp
+++ b/src/wasm/wasm-s-parser.cpp
@@ -2830,9 +2830,24 @@ Expression* SExpressionWasmBuilder::makeTupleExtract(Element& s) {
 }
 
 Expression* SExpressionWasmBuilder::makeCallRef(Element& s, bool isReturn) {
+  Index operandsStart = 1;
+  HeapType sigType;
+  if (isReturn) {
+    sigType = parseHeapType(*s[1]);
+    operandsStart = 2;
+  }
   std::vector<Expression*> operands;
-  parseOperands(s, 1, s.size() - 1, operands);
+  parseOperands(s, operandsStart, s.size() - 1, operands);
   auto* target = parseExpression(s[s.size() - 1]);
+
+  if (isReturn) {
+    if (!sigType.isSignature()) {
+      throw ParseException(
+        "return_call_ref type annotation should be a signature", s.line, s.col);
+    }
+    return Builder(wasm).makeCallRef(
+      target, operands, sigType.getSignature().results, isReturn);
+  }
   return ValidatingBuilder(wasm, s.line, s.col)
     .validateAndMakeCallRef(target, operands, isReturn);
 }
diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp
index d88138075..243b2810b 100644
--- a/src/wasm/wasm-stack.cpp
+++ b/src/wasm/wasm-stack.cpp
@@ -2013,8 +2013,14 @@ void BinaryInstWriter::visitI31Get(I31Get* curr) {
 }
 
 void BinaryInstWriter::visitCallRef(CallRef* curr) {
-  o << int8_t(curr->isReturn ? BinaryConsts::RetCallRef
-                             : BinaryConsts::CallRef);
+  if (curr->isReturn) {
+    assert(curr->target->type != Type::unreachable);
+    // TODO: `emitUnreachable` if target has bottom type.
+    o << int8_t(BinaryConsts::RetCallRef);
+    parent.writeIndexedHeapType(curr->target->type.getHeapType());
+    return;
+  }
+  o << int8_t(BinaryConsts::CallRef);
 }
 
 void BinaryInstWriter::visitRefTest(RefTest* curr) {