[custom-descriptors] ref.cast_desc (#7630)

tlively · web-flow · commit b3a14c2ceaff · 2025-05-29T07:39:00.000-07:00
Add an optional `desc` child to `RefCast` for use in the new
`ref.cast_desc` instruction. Add support for parsing, printing, and
validating the new instruction. Fix a few minor issues left over from
adding the branching descriptor casts along the way.
diff --git a/scripts/gen-s-parser.py b/scripts/gen-s-parser.py
@@ -609,7 +609,8 @@
     ("i31.get_s",            "makeI31Get(true)"),
     ("i31.get_u",            "makeI31Get(false)"),
     ("ref.test",             "makeRefTest()"),
-    ("ref.cast",             "makeRefCast()"),
+    ("ref.cast",             "makeRefCast(false)"),
+    ("ref.cast_desc",        "makeRefCast(true)"),
     ("ref.get_desc",         "makeRefGetDesc()"),
     ("br_on_null",           "makeBrOnNull()"),
     ("br_on_non_null",       "makeBrOnNull(true)"),
diff --git a/scripts/test/fuzzing.py b/scripts/test/fuzzing.py
@@ -115,6 +115,7 @@
     'custom-descriptors.wast',
     'br_on_cast_desc.wast',
     'ref.get_cast.wast',
+    'ref.cast_desc.wast',
     # TODO: fix split_wast() on tricky escaping situations like a string ending
     #       in \\" (the " is not escaped - there is an escaped \ before it)
     'string-lifting-section.wast',
diff --git a/src/gen-s-parser.inc b/src/gen-s-parser.inc
@@ -4807,12 +4807,23 @@ switch (buf[0]) {
               return Ok{};
             }
             goto parse_error;
-          case 'c':
-            if (op == "ref.cast"sv) {
-              CHECK_ERR(makeRefCast(ctx, pos, annotations));
-              return Ok{};
+          case 'c': {
+            switch (buf[8]) {
+              case '\0':
+                if (op == "ref.cast"sv) {
+                  CHECK_ERR(makeRefCast(ctx, pos, annotations, false));
+                  return Ok{};
+                }
+                goto parse_error;
+              case '_':
+                if (op == "ref.cast_desc"sv) {
+                  CHECK_ERR(makeRefCast(ctx, pos, annotations, true));
+                  return Ok{};
+                }
+                goto parse_error;
+              default: goto parse_error;
             }
-            goto parse_error;
+          }
           case 'e':
             if (op == "ref.eq"sv) {
               CHECK_ERR(makeRefEq(ctx, pos, annotations));
diff --git a/src/ir/child-typer.h b/src/ir/child-typer.h
@@ -852,9 +852,17 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
     note(&curr->ref, Type(top, Nullable));
   }
 
-  void visitRefCast(RefCast* curr) {
+  void visitRefCast(RefCast* curr, std::optional<Type> target = std::nullopt) {
     auto top = curr->type.getHeapType().getTop();
     note(&curr->ref, Type(top, Nullable));
+    if (curr->desc) {
+      if (!target) {
+        target = curr->type;
+      }
+      auto desc = target->getHeapType().getDescriptorType();
+      assert(desc);
+      note(&curr->desc, Type(*desc, Nullable, curr->type.getExactness()));
+    }
   }
 
   void visitRefGetDesc(RefGetDesc* curr,
@@ -881,8 +889,9 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
         auto top = target->getHeapType().getTop();
         note(&curr->ref, Type(top, Nullable));
         if (curr->op == BrOnCastDesc || curr->op == BrOnCastDescFail) {
-          auto descriptor = *target->getHeapType().getDescriptorType();
-          note(&curr->desc, Type(descriptor, Nullable));
+          auto desc = target->getHeapType().getDescriptorType();
+          assert(desc);
+          note(&curr->desc, Type(*desc, Nullable));
         }
         return;
       }
diff --git a/src/parser/contexts.h b/src/parser/contexts.h
@@ -733,7 +733,7 @@ struct NullInstrParserCtx {
     return Ok{};
   }
   template<typename TypeT>
-  Result<> makeRefCast(Index, const std::vector<Annotation>&, TypeT) {
+  Result<> makeRefCast(Index, const std::vector<Annotation>&, TypeT, bool) {
     return Ok{};
   }
   template<typename HeapTypeT>
@@ -2592,8 +2592,9 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {
 
   Result<> makeRefCast(Index pos,
                        const std::vector<Annotation>& annotations,
-                       Type type) {
-    return withLoc(pos, irBuilder.makeRefCast(type));
+                       Type type,
+                       bool isDesc) {
+    return withLoc(pos, irBuilder.makeRefCast(type, isDesc));
   }
 
   Result<> makeRefGetDesc(Index pos,
diff --git a/src/parser/parsers.h b/src/parser/parsers.h
@@ -231,7 +231,7 @@ Result<> makeI31Get(Ctx&, Index, const std::vector<Annotation>&, bool signed_);
 template<typename Ctx>
 Result<> makeRefTest(Ctx&, Index, const std::vector<Annotation>&);
 template<typename Ctx>
-Result<> makeRefCast(Ctx&, Index, const std::vector<Annotation>&);
+Result<> makeRefCast(Ctx&, Index, const std::vector<Annotation>&, bool isDesc);
 template<typename Ctx>
 Result<> makeRefGetDesc(Ctx&, Index, const std::vector<Annotation>&);
 template<typename Ctx>
@@ -2219,11 +2219,13 @@ makeRefTest(Ctx& ctx, Index pos, const std::vector<Annotation>& annotations) {
 }
 
 template<typename Ctx>
-Result<>
-makeRefCast(Ctx& ctx, Index pos, const std::vector<Annotation>& annotations) {
+Result<> makeRefCast(Ctx& ctx,
+                     Index pos,
+                     const std::vector<Annotation>& annotations,
+                     bool isDesc) {
   auto type = reftype(ctx);
   CHECK_ERR(type);
-  return ctx.makeRefCast(pos, annotations, *type);
+  return ctx.makeRefCast(pos, annotations, *type, isDesc);
 }
 
 template<typename Ctx>
diff --git a/src/passes/Print.cpp b/src/passes/Print.cpp
@@ -2208,7 +2208,11 @@ struct PrintExpressionContents
     printType(curr->castType);
   }
   void visitRefCast(RefCast* curr) {
-    printMedium(o, "ref.cast ");
+    if (curr->desc) {
+      printMedium(o, "ref.cast_desc ");
+    } else {
+      printMedium(o, "ref.cast ");
+    }
     printType(curr->type);
   }
   void visitRefGetDesc(RefGetDesc* curr) {
diff --git a/src/wasm-binary.h b/src/wasm-binary.h
@@ -1172,6 +1172,8 @@ enum ASTNodes {
   RefTestNull = 0x15,
   RefCast = 0x16,
   RefCastNull = 0x17,
+  RefCastDesc = 0x23,
+  RefCastDescNull = 0x24,
   BrOnCast = 0x18,
   BrOnCastFail = 0x19,
   BrOnCastDesc = 0x25,
diff --git a/src/wasm-builder.h b/src/wasm-builder.h
@@ -888,8 +888,12 @@ class Builder {
     return ret;
   }
   RefCast* makeRefCast(Expression* ref, Type type) {
+    return makeRefCast(ref, nullptr, type);
+  }
+  RefCast* makeRefCast(Expression* ref, Expression* desc, Type type) {
     auto* ret = wasm.allocator.alloc<RefCast>();
     ret->ref = ref;
+    ret->desc = desc;
     ret->type = type;
     ret->finalize();
     return ret;
diff --git a/src/wasm-delegations-fields.def b/src/wasm-delegations-fields.def
@@ -650,6 +650,7 @@ DELEGATE_FIELD_CHILD(RefTest, ref)
 DELEGATE_FIELD_CASE_END(RefTest)
 
 DELEGATE_FIELD_CASE_START(RefCast)
+DELEGATE_FIELD_OPTIONAL_IMMEDIATE_TYPED_CHILD(RefCast, desc)
 DELEGATE_FIELD_CHILD(RefCast, ref)
 DELEGATE_FIELD_CASE_END(RefCast)
 
diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h
@@ -208,7 +208,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
                        bool isReturn,
                        std::optional<std::uint8_t> inline_ = std::nullopt);
   Result<> makeRefTest(Type type);
-  Result<> makeRefCast(Type type);
+  Result<> makeRefCast(Type type, bool isDesc);
   Result<> makeRefGetDesc(HeapType type);
   Result<> makeBrOn(Index label,
                     BrOnOp op,
diff --git a/src/wasm.h b/src/wasm.h
@@ -1620,6 +1620,9 @@ class RefCast : public SpecificExpression<Expression::RefCastId> {
 
   Expression* ref;
 
+  // Used only for ref.cast_desc.
+  Expression* desc;
+
   void finalize();
 
   Type& getCastType() { return type; }
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
@@ -4505,11 +4505,22 @@ Result<> WasmBinaryReader::readInst() {
         }
         case BinaryConsts::RefCast: {
           auto [heapType, exactness] = getHeapType();
-          return builder.makeRefCast(Type(heapType, NonNullable, exactness));
+          return builder.makeRefCast(Type(heapType, NonNullable, exactness),
+                                     false);
         }
         case BinaryConsts::RefCastNull: {
           auto [heapType, exactness] = getHeapType();
-          return builder.makeRefCast(Type(heapType, Nullable, exactness));
+          return builder.makeRefCast(Type(heapType, Nullable, exactness),
+                                     false);
+        }
+        case BinaryConsts::RefCastDesc: {
+          auto [heapType, exactness] = getHeapType();
+          return builder.makeRefCast(Type(heapType, NonNullable, exactness),
+                                     true);
+        }
+        case BinaryConsts::RefCastDescNull: {
+          auto [heapType, exactness] = getHeapType();
+          return builder.makeRefCast(Type(heapType, Nullable, exactness), true);
         }
         case BinaryConsts::RefGetDesc: {
           auto type = getIndexedHeapType();
diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
@@ -2005,11 +2005,28 @@ Result<> IRBuilder::makeRefTest(Type type) {
   return Ok{};
 }
 
-Result<> IRBuilder::makeRefCast(Type type) {
+Result<> IRBuilder::makeRefCast(Type type, bool isDesc) {
+  std::optional<HeapType> descriptor;
+  if (isDesc) {
+    assert(type.isRef());
+    descriptor = type.getHeapType().getDescriptorType();
+    if (!descriptor) {
+      return Err{"cast target must have descriptor"};
+    }
+  }
+
   RefCast curr;
   curr.type = type;
+  // Placeholder value to differentiate ref.cast_desc.
+  curr.desc = isDesc ? &curr : nullptr;
   CHECK_ERR(visitRefCast(&curr));
-  push(builder.makeRefCast(curr.ref, type));
+
+  if (isDesc) {
+    CHECK_ERR(
+      validateTypeAnnotation(type.with(*descriptor).with(Nullable), curr.desc));
+  }
+
+  push(builder.makeRefCast(curr.ref, curr.desc, type));
   return Ok{};
 }
 
@@ -2026,6 +2043,15 @@ Result<> IRBuilder::makeRefGetDesc(HeapType type) {
 
 Result<> IRBuilder::makeBrOn(
   Index label, BrOnOp op, Type in, Type out, std::optional<bool> likely) {
+  std::optional<HeapType> descriptor;
+  if (op == BrOnCastDesc || op == BrOnCastDescFail) {
+    assert(out.isRef());
+    descriptor = out.getHeapType().getDescriptorType();
+    if (!descriptor) {
+      return Err{"cast target must have descriptor"};
+    }
+  }
+
   BrOn curr;
   curr.op = op;
   curr.castType = out;
@@ -2039,11 +2065,6 @@ Result<> IRBuilder::makeBrOn(
       break;
     case BrOnCastDesc:
     case BrOnCastDescFail: {
-      assert(out.isRef());
-      auto descriptor = out.getHeapType().getDescriptorType();
-      if (!descriptor) {
-        return Err{"cast target must have descriptor"};
-      }
       CHECK_ERR(validateTypeAnnotation(out.with(*descriptor).with(Nullable),
                                        curr.desc));
     }
diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp
@@ -81,7 +81,7 @@ void BinaryInstWriter::visitBreak(Break* curr) {
       // are casting, and to emit the proper thing.
       RefCast cast;
       cast.type = to;
-      cast.ref = nullptr;
+      cast.ref = cast.desc = nullptr;
       visitRefCast(&cast);
     };
 
@@ -2271,11 +2271,23 @@ void BinaryInstWriter::visitRefTest(RefTest* curr) {
 }
 
 void BinaryInstWriter::visitRefCast(RefCast* curr) {
+  if (curr->desc && curr->desc->type.isNull()) {
+    emitUnreachable();
+    return;
+  }
   o << int8_t(BinaryConsts::GCPrefix);
   if (curr->type.isNullable()) {
-    o << U32LEB(BinaryConsts::RefCastNull);
+    if (curr->desc) {
+      o << U32LEB(BinaryConsts::RefCastDescNull);
+    } else {
+      o << U32LEB(BinaryConsts::RefCastNull);
+    }
   } else {
-    o << U32LEB(BinaryConsts::RefCast);
+    if (curr->desc) {
+      o << U32LEB(BinaryConsts::RefCastDesc);
+    } else {
+      o << U32LEB(BinaryConsts::RefCast);
+    }
   }
   parent.writeHeapType(curr->type.getHeapType(), curr->type.getExactness());
 }
diff --git a/src/wasm/wasm-validator.cpp b/src/wasm/wasm-validator.cpp
@@ -2966,6 +2966,40 @@ void FunctionValidator::visitRefCast(RefCast* curr) {
                  "ref.cast to exact type requires custom descriptors "
                  "[--enable-custom-descriptors]");
   }
+
+  if (!curr->desc) {
+    return;
+  }
+
+  shouldBeTrue(getModule()->features.hasCustomDescriptors(),
+               curr,
+               "ref.cast_desc requires custom descriptors "
+               "[--enable-custom-descriptors]");
+  if (!shouldBeTrue(curr->desc && curr->desc->type.isRef(),
+                    curr,
+                    "ref.cast_desc descriptor must have ref type")) {
+    return;
+  }
+  auto descriptor = curr->desc->type.getHeapType();
+  if (descriptor.isBottom()) {
+    return;
+  }
+
+  auto described = descriptor.getDescribedType();
+  if (!shouldBeTrue(bool(described),
+                    curr,
+                    "ref.cast_desc descriptor should have a described type")) {
+    return;
+  }
+  shouldBeEqual(*described,
+                curr->type.getHeapType(),
+                curr,
+                "ref.cast_desc cast type should be described by descriptor");
+  shouldBeEqual(
+    curr->type.getExactness(),
+    curr->desc->type.getExactness(),
+    curr,
+    "ref.cast_desc cast exactness should match descriptor exactness");
 }
 
 void FunctionValidator::visitRefGetDesc(RefGetDesc* curr) {
diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp
@@ -1052,11 +1052,31 @@ void RefTest::finalize() {
 }
 
 void RefCast::finalize() {
-  if (ref->type == Type::unreachable) {
+  if (ref->type == Type::unreachable ||
+      (desc && desc->type == Type::unreachable)) {
     type = Type::unreachable;
     return;
   }
 
+  if (desc) {
+    if (desc->type.isNull()) {
+      // Cast will never be executed and the instruction will not be emitted.
+      // Model this with an uninhabitable cast type.
+      type = desc->type.with(NonNullable);
+      return;
+    }
+    // The cast heap type and exactness is determined by the descriptor's type.
+    // Its nullability can be improved if the input valus is non-nullable.
+    auto heapType = desc->type.getHeapType().getDescribedType();
+    assert(heapType);
+    auto exactness = desc->type.getExactness();
+    type = type.with(*heapType).with(exactness);
+    if (ref->type.isNonNullable()) {
+      type = type.with(NonNullable);
+    }
+    return;
+  }
+
   // We reach this before validation, so the input type might be totally wrong.
   // Return early in this case to avoid doing the wrong thing below.
   if (!ref->type.isRef()) {
diff --git a/test/lit/basic/custom-descriptors.wast b/test/lit/basic/custom-descriptors.wast
diff --git a/test/spec/br_on_cast_desc.wast b/test/spec/br_on_cast_desc.wast
diff --git a/test/spec/ref.cast_desc.wast b/test/spec/ref.cast_desc.wast

Original file line number	Diff line number	Diff line change
`@@ -733,7 +733,7 @@ struct NullInstrParserCtx {`
`733`	`733`	`return Ok{};`
`734`	`734`	`}`
`735`	`735`	`template<typename TypeT>`
`736`		`- Result<> makeRefCast(Index, const std::vector<Annotation>&, TypeT) {`
	`736`	`+ Result<> makeRefCast(Index, const std::vector<Annotation>&, TypeT, bool) {`
`737`	`737`	`return Ok{};`
`738`	`738`	`}`
`739`	`739`	`template<typename HeapTypeT>`
`@@ -2592,8 +2592,9 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {`
`2592`	`2592`
`2593`	`2593`	`Result<> makeRefCast(Index pos,`
`2594`	`2594`	`const std::vector<Annotation>& annotations,`
`2595`		`- Type type) {`
`2596`		`- return withLoc(pos, irBuilder.makeRefCast(type));`
	`2595`	`+ Type type,`
	`2596`	`+ bool isDesc) {`
	`2597`	`+ return withLoc(pos, irBuilder.makeRefCast(type, isDesc));`
`2597`	`2598`	`}`
`2598`	`2599`
`2599`	`2600`	`Result<> makeRefGetDesc(Index pos,`
Original file line number	Diff line number	Diff line change
`@@ -2208,7 +2208,11 @@ struct PrintExpressionContents`
`2208`	`2208`	`printType(curr->castType);`
`2209`	`2209`	`}`
`2210`	`2210`	`void visitRefCast(RefCast* curr) {`
`2211`		`- printMedium(o, "ref.cast ");`
	`2211`	`+ if (curr->desc) {`
	`2212`	`+ printMedium(o, "ref.cast_desc ");`
	`2213`	`+ } else {`
	`2214`	`+ printMedium(o, "ref.cast ");`
	`2215`	`+ }`
`2212`	`2216`	`printType(curr->type);`
`2213`	`2217`	`}`
`2214`	`2218`	`void visitRefGetDesc(RefGetDesc* curr) {`