Binary encoding and decoding of exact casts (#7347)

The binary format for casts needs to be extended to support exact
references. In particular, add new opcodes for `ref.test` and `ref.cast`
that take reference type immediates instead of heap type immediates. Use
two more bits in the flags immediate of `br_on_cast` and
`br_on_cast_fail` to encode the source and destination exactness.

Author: tlively

src/wasm-binary.h

@@ -304,6 +304,13 @@ enum SegmentFlag {
   UsesExpressions = 1 << 2
 };
 
+enum BrOnCastFlag {
+  InputNullable = 1 << 0,
+  OutputNullable = 1 << 1,
+  InputExact = 1 << 2,
+  OutputExact = 1 << 3,
+};
+
 enum EncodedType {
   // value types
   i32 = -0x1,  // 0x7f
@@ -1126,6 +1133,8 @@ enum ASTNodes {
   I31GetS = 0x1d,
   I31GetU = 0x1e,
   RefI31Shared = 0x1f,
+  RefTestRT = 0x20,
+  RefCastRT = 0x21,
 
   // Shared GC Opcodes
 

src/wasm/wasm-binary.cpp

@@ -4251,16 +4251,30 @@ Result<> WasmBinaryReader::readInst() {
           return builder.makeRefTest(Type(getHeapType(), NonNullable));
         case BinaryConsts::RefTestNull:
           return builder.makeRefTest(Type(getHeapType(), Nullable));
+        case BinaryConsts::RefTestRT:
+          return builder.makeRefTest(getType());
         case BinaryConsts::RefCast:
           return builder.makeRefCast(Type(getHeapType(), NonNullable));
         case BinaryConsts::RefCastNull:
           return builder.makeRefCast(Type(getHeapType(), Nullable));
+        case BinaryConsts::RefCastRT:
+          return builder.makeRefCast(getType());
         case BinaryConsts::BrOnCast:
         case BinaryConsts::BrOnCastFail: {
           auto flags = getInt8();
           auto label = getU32LEB();
-          auto in = Type(getHeapType(), (flags & 1) ? Nullable : NonNullable);
-          auto cast = Type(getHeapType(), (flags & 2) ? Nullable : NonNullable);
+          auto srcNull = (flags & BinaryConsts::BrOnCastFlag::InputNullable)
+                           ? Nullable
+                           : NonNullable;
+          auto dstNull = (flags & BinaryConsts::BrOnCastFlag::OutputNullable)
+                           ? Nullable
+                           : NonNullable;
+          auto srcExact =
+            (flags & BinaryConsts::BrOnCastFlag::InputExact) ? Exact : Inexact;
+          auto dstExact =
+            (flags & BinaryConsts::BrOnCastFlag::OutputExact) ? Exact : Inexact;
+          auto in = Type(getHeapType(), srcNull, srcExact);
+          auto cast = Type(getHeapType(), dstNull, dstExact);
           auto kind = op == BinaryConsts::BrOnCast ? BrOnCast : BrOnCastFail;
           return builder.makeBrOn(label, kind, in, cast);
         }

src/wasm/wasm-stack.cpp

@@ -2260,22 +2260,38 @@ void BinaryInstWriter::visitCallRef(CallRef* curr) {
 
 void BinaryInstWriter::visitRefTest(RefTest* curr) {
   o << int8_t(BinaryConsts::GCPrefix);
-  if (curr->castType.isNullable()) {
-    o << U32LEB(BinaryConsts::RefTestNull);
+  if (curr->castType.isExact() &&
+      parent.getModule()->features.hasCustomDescriptors()) {
+    // Fall back to the general form with a reftype immediate.
+    o << U32LEB(BinaryConsts::RefTestRT);
+    parent.writeType(curr->castType);
   } else {
-    o << U32LEB(BinaryConsts::RefTest);
+    // Use the special-case form with heap type immediate.
+    if (curr->castType.isNullable()) {
+      o << U32LEB(BinaryConsts::RefTestNull);
+    } else {
+      o << U32LEB(BinaryConsts::RefTest);
+    }
+    parent.writeHeapType(curr->castType.getHeapType());
   }
-  parent.writeHeapType(curr->castType.getHeapType());
 }
 
 void BinaryInstWriter::visitRefCast(RefCast* curr) {
   o << int8_t(BinaryConsts::GCPrefix);
-  if (curr->type.isNullable()) {
-    o << U32LEB(BinaryConsts::RefCastNull);
+  if (curr->type.isExact() &&
+      parent.getModule()->features.hasCustomDescriptors()) {
+    // Fall back to the general form with a reftype immediate.
+    o << U32LEB(BinaryConsts::RefCastRT);
+    parent.writeType(curr->type);
   } else {
-    o << U32LEB(BinaryConsts::RefCast);
+    // Use the special-case form with heap type immediate.
+    if (curr->type.isNullable()) {
+      o << U32LEB(BinaryConsts::RefCastNull);
+    } else {
+      o << U32LEB(BinaryConsts::RefCast);
+    }
+    parent.writeHeapType(curr->type.getHeapType());
   }
-  parent.writeHeapType(curr->type.getHeapType());
 }
 
 void BinaryInstWriter::visitBrOn(BrOn* curr) {
@@ -2298,8 +2314,24 @@ void BinaryInstWriter::visitBrOn(BrOn* curr) {
       }
       assert(curr->ref->type.isRef());
       assert(Type::isSubType(curr->castType, curr->ref->type));
-      uint8_t flags = (curr->ref->type.isNullable() ? 1 : 0) |
-                      (curr->castType.isNullable() ? 2 : 0);
+      uint8_t flags = 0;
+      if (curr->ref->type.isNullable()) {
+        flags |= BinaryConsts::BrOnCastFlag::InputNullable;
+      }
+      if (curr->castType.isNullable()) {
+        flags |= BinaryConsts::BrOnCastFlag::OutputNullable;
+      }
+      if (parent.getModule()->features.hasCustomDescriptors()) {
+        // If custom descriptors (and therefore exact references) are not
+        // enabled, then these flags wouldn't be recognized, and we will be
+        // generalizing all exact references to be non-exact anyway.
+        if (curr->ref->type.isExact()) {
+          flags |= BinaryConsts::BrOnCastFlag::InputExact;
+        }
+        if (curr->castType.isExact()) {
+          flags |= BinaryConsts::BrOnCastFlag::OutputExact;
+        }
+      }
       o << flags;
       o << U32LEB(getBreakIndex(curr->name));
       parent.writeHeapType(curr->ref->type.getHeapType());