Update the binary format for exact heap types (#7502)

Although we now store exactness on `Type` instead of `HeapType`, the
binary format still encodes exactness as part of the heap type. Update
the binary parser to read exactness wherever a heap type (but not just a
type index) is expected and update the binary writer similarly.

Author: tlively

src/wasm-binary.h

@@ -307,8 +307,6 @@ enum SegmentFlag {
 enum BrOnCastFlag {
   InputNullable = 1 << 0,
   OutputNullable = 1 << 1,
-  InputExact = 1 << 2,
-  OutputExact = 1 << 3,
 };
 
 enum EncodedType {
@@ -334,7 +332,6 @@ enum EncodedType {
   eqref = -0x13,        // 0x6d
   nonnullable = -0x1c,  // 0x64
   nullable = -0x1d,     // 0x63
-  exact = -0x1e,        // 0x62
   contref = -0x18,      // 0x68
   nullcontref = -0x0b,  // 0x75
   // exception handling
@@ -351,6 +348,8 @@ enum EncodedType {
   SubFinal = 0x4f,
   Shared = 0x65,
   SharedLEB = -0x1b, // Also 0x65 as an SLEB128
+  Exact = 0x62,
+  ExactLEB = -0x1e, // Also 0x62 as an SLEB128
   Rec = 0x4e,
   Descriptor = 0x4d,
   Describes = 0x4c,
@@ -1135,8 +1134,6 @@ enum ASTNodes {
   I31GetS = 0x1d,
   I31GetU = 0x1e,
   RefI31Shared = 0x1f,
-  RefTestRT = 0x20,
-  RefCastRT = 0x21,
 
   // Shared GC Opcodes
 
@@ -1378,7 +1375,7 @@ class WasmBinaryWriter {
 
   // Writes an arbitrary heap type, which may be indexed or one of the
   // basic types like funcref.
-  void writeHeapType(HeapType type);
+  void writeHeapType(HeapType type, Exactness exact);
   // Writes an indexed heap type. Note that this is encoded differently than a
   // general heap type because it does not allow negative values for basic heap
   // types.
@@ -1509,8 +1506,7 @@ class WasmBinaryReader {
   Type getType();
   // Get a type given the initial S32LEB has already been read, and is provided.
   Type getType(int code);
-  Type getTypeNoExact(int code);
-  HeapType getHeapType();
+  std::pair<HeapType, Exactness> getHeapType();
   HeapType getIndexedHeapType();
 
   Type getConcreteType();

src/wasm-type.h

@@ -382,7 +382,7 @@ class Type {
   bool isExact() const {
     return isRef() && !getHeapType().isBasic() && (id & ExactMask);
   }
-  bool isInexact() const { return isRef() && !(id & ExactMask); }
+  bool isInexact() const { return isRef() && !isExact(); }
   HeapType getHeapType() const {
     assert(isRef());
     HeapType masked(id & ~NullMask);

src/wasm/wasm-binary.cpp

@@ -279,7 +279,7 @@ void WasmBinaryWriter::writeTypes() {
       }
       if (super) {
         o << U32LEB(1);
-        writeHeapType(*super);
+        writeHeapType(*super, Inexact);
       } else {
         o << U32LEB(0);
       }
@@ -289,11 +289,11 @@ void WasmBinaryWriter::writeTypes() {
     }
     if (auto desc = type.getDescribedType()) {
       o << uint8_t(BinaryConsts::EncodedType::Describes);
-      writeHeapType(*desc);
+      writeHeapType(*desc, Inexact);
     }
     if (auto desc = type.getDescriptorType()) {
       o << uint8_t(BinaryConsts::EncodedType::Descriptor);
-      writeHeapType(*desc);
+      writeHeapType(*desc, Inexact);
     }
     switch (type.getKind()) {
       case HeapTypeKind::Func: {
@@ -322,7 +322,7 @@ void WasmBinaryWriter::writeTypes() {
         break;
       case HeapTypeKind::Cont:
         o << uint8_t(BinaryConsts::EncodedType::Cont);
-        writeHeapType(type.getContinuation().type);
+        writeHeapType(type.getContinuation().type, Inexact);
         break;
       case HeapTypeKind::Basic:
         WASM_UNREACHABLE("unexpected kind");
@@ -1574,12 +1574,6 @@ void WasmBinaryWriter::writeInlineBuffer(const char* data, size_t size) {
 
 void WasmBinaryWriter::writeType(Type type) {
   if (type.isRef()) {
-    // Exact references are introduced by the custom descriptors feature, but
-    // can be used internally even when it is not enabled. In that case, we have
-    // to generalize the types to be inexact before writing them.
-    if (!wasm->features.hasCustomDescriptors()) {
-      type = Type(type.getHeapType(), type.getNullability(), Inexact);
-    }
     // The only reference types allowed without GC are funcref, externref, and
     // exnref. We internally use more refined versions of those types, but we
     // cannot emit those without GC.
@@ -1596,12 +1590,6 @@ void WasmBinaryWriter::writeType(Type type) {
         type = Type(type.getHeapType().getTop(), Nullable);
       }
     }
-    // If the type is exact, emit the exact prefix and continue on without
-    // considering exactness.
-    if (type.isExact()) {
-      o << S32LEB(BinaryConsts::EncodedType::exact);
-      type = Type(type.getHeapType(), type.getNullability(), Inexact);
-    }
     auto heapType = type.getHeapType();
     if (type.isNullable() && heapType.isBasic() && !heapType.isShared()) {
       switch (heapType.getBasic(Unshared)) {
@@ -1657,7 +1645,7 @@ void WasmBinaryWriter::writeType(Type type) {
     } else {
       o << S32LEB(BinaryConsts::EncodedType::nonnullable);
     }
-    writeHeapType(type.getHeapType());
+    writeHeapType(type.getHeapType(), type.getExactness());
     return;
   }
   int ret = 0;
@@ -1688,14 +1676,20 @@ void WasmBinaryWriter::writeType(Type type) {
   o << S32LEB(ret);
 }
 
-void WasmBinaryWriter::writeHeapType(HeapType type) {
+void WasmBinaryWriter::writeHeapType(HeapType type, Exactness exactness) {
   // ref.null always has a bottom heap type in Binaryen IR, but those types are
   // only actually valid with GC. Otherwise, emit the corresponding valid top
   // types instead.
+  if (!wasm->features.hasCustomDescriptors()) {
+    exactness = Inexact;
+  }
   if (!wasm->features.hasGC()) {
     type = type.getTop();
   }
-
+  assert(!type.isBasic() || exactness == Inexact);
+  if (exactness == Exact) {
+    o << uint8_t(BinaryConsts::EncodedType::Exact);
+  }
   if (!type.isBasic()) {
     o << S64LEB(getTypeIndex(type)); // TODO: Actually s33
     return;
@@ -2190,43 +2184,41 @@ Signature WasmBinaryReader::getBlockType() {
   return Signature(Type::none, getType(code));
 }
 
-Type WasmBinaryReader::getTypeNoExact(int code) {
+Type WasmBinaryReader::getType(int code) {
   Type type;
   if (getBasicType(code, type)) {
     return type;
   }
+  auto [heapType, exactness] = getHeapType();
   switch (code) {
     case BinaryConsts::EncodedType::nullable:
-      return Type(getHeapType(), Nullable);
+      return Type(heapType, Nullable, exactness);
     case BinaryConsts::EncodedType::nonnullable:
-      return Type(getHeapType(), NonNullable);
+      return Type(heapType, NonNullable, exactness);
     default:
       throwError("invalid wasm type: " + std::to_string(code));
   }
   WASM_UNREACHABLE("unexpected type");
 }
 
-Type WasmBinaryReader::getType(int code) {
-  if (code == BinaryConsts::EncodedType::exact) {
-    auto type = getTypeNoExact(getS32LEB());
-    if (!type.isRef()) {
-      throwError("invalid exact prefix on non-reference type");
-    }
-    return Type(type.getHeapType(), type.getNullability(), Exact);
-  }
-  return getTypeNoExact(code);
-}
-
 Type WasmBinaryReader::getType() { return getType(getS32LEB()); }
 
-HeapType WasmBinaryReader::getHeapType() {
+std::pair<HeapType, Exactness> WasmBinaryReader::getHeapType() {
   auto type = getS64LEB(); // TODO: Actually s33
+  auto exactness = Inexact;
+  if (type == BinaryConsts::EncodedType::ExactLEB) {
+    exactness = Exact;
+    type = getS64LEB(); // TODO: Actually s33
+  }
   // Single heap types are negative; heap type indices are non-negative
   if (type >= 0) {
     if (size_t(type) >= types.size()) {
-      throwError("invalid signature index: " + std::to_string(type));
+      throwError("invalid type index: " + std::to_string(type));
     }
-    return types[type];
+    return {types[type], exactness};
+  }
+  if (exactness == Exact) {
+    throwError("invalid type index: " + std::to_string(type));
   }
   auto share = Unshared;
   if (type == BinaryConsts::EncodedType::SharedLEB) {
@@ -2235,11 +2227,9 @@ HeapType WasmBinaryReader::getHeapType() {
   }
   HeapType ht;
   if (getBasicHeapType(type, ht)) {
-    return ht.getBasic(share);
-  } else {
-    throwError("invalid wasm heap type: " + std::to_string(type));
+    return {ht.getBasic(share), Inexact};
   }
-  WASM_UNREACHABLE("unexpected type");
+  throwError("invalid wasm heap type: " + std::to_string(type));
 }
 
 HeapType WasmBinaryReader::getIndexedHeapType() {
@@ -2361,23 +2351,34 @@ void WasmBinaryReader::readMemories() {
 void WasmBinaryReader::readTypes() {
   TypeBuilder builder(getU32LEB());
 
-  auto readHeapType = [&]() -> HeapType {
+  auto readHeapType = [&]() -> std::pair<HeapType, Exactness> {
     int64_t htCode = getS64LEB(); // TODO: Actually s33
+    auto exactness = Inexact;
+    if (htCode == BinaryConsts::EncodedType::ExactLEB) {
+      exactness = Exact;
+      htCode = getS64LEB(); // TODO: Actually s33
+    }
+    if (htCode >= 0) {
+      if (size_t(htCode) >= builder.size()) {
+        throwError("invalid type index: " + std::to_string(htCode));
+      }
+      return {builder.getTempHeapType(size_t(htCode)), exactness};
+    }
+    if (exactness == Exact) {
+      throwError("invalid type index: " + std::to_string(htCode));
+    }
     auto share = Unshared;
     if (htCode == BinaryConsts::EncodedType::SharedLEB) {
       share = Shared;
       htCode = getS64LEB(); // TODO: Actually s33
     }
     HeapType ht;
     if (getBasicHeapType(htCode, ht)) {
-      return ht.getBasic(share);
-    }
-    if (size_t(htCode) >= builder.size()) {
-      throwError("invalid type index: " + std::to_string(htCode));
+      return {ht.getBasic(share), Inexact};
     }
-    return builder.getTempHeapType(size_t(htCode));
+    throwError("invalid wasm heap type: " + std::to_string(htCode));
   };
-  auto makeTypeNoExact = [&](int32_t typeCode) {
+  auto makeType = [&](int32_t typeCode) {
     Type type;
     if (getBasicType(typeCode, type)) {
       return type;
@@ -2390,29 +2391,18 @@ void WasmBinaryReader::readTypes() {
                              ? Nullable
                              : NonNullable;
 
-        HeapType ht = readHeapType();
+        auto [ht, exactness] = readHeapType();
         if (ht.isBasic()) {
-          return Type(ht, nullability);
+          return Type(ht, nullability, exactness);
         }
 
-        return builder.getTempRefType(ht, nullability);
+        return builder.getTempRefType(ht, nullability, exactness);
       }
       default:
         throwError("unexpected type index: " + std::to_string(typeCode));
     }
     WASM_UNREACHABLE("unexpected type");
   };
-  auto makeType = [&](int32_t typeCode) {
-    if (typeCode == BinaryConsts::EncodedType::exact) {
-      auto type = makeTypeNoExact(getS32LEB());
-      if (!type.isRef()) {
-        throwError("unexpected exact prefix on non-reference type");
-      }
-      return builder.getTempRefType(
-        type.getHeapType(), type.getNullability(), Exact);
-    }
-    return makeTypeNoExact(typeCode);
-  };
   auto readType = [&]() { return makeType(getS32LEB()); };
 
   auto readSignatureDef = [&]() {
@@ -2431,7 +2421,10 @@ void WasmBinaryReader::readTypes() {
   };
 
   auto readContinuationDef = [&]() {
-    HeapType ht = readHeapType();
+    auto [ht, exactness] = readHeapType();
+    if (exactness != Inexact) {
+      throw ParseException("invalid exact type in cont definition");
+    }
     if (!ht.isSignature()) {
       throw ParseException("cont types must be built from function types");
     }
@@ -3046,8 +3039,12 @@ Result<> WasmBinaryReader::readInst() {
       return builder.visitCatchAll();
     case BinaryConsts::Delegate:
       return builder.visitDelegate(getU32LEB());
-    case BinaryConsts::RefNull:
-      return builder.makeRefNull(getHeapType());
+    case BinaryConsts::RefNull: {
+      auto [heapType, exactness] = getHeapType();
+      // Exactness is allowed but doesn't matter, since we always use the bottom
+      // heap type.
+      return builder.makeRefNull(heapType);
+    }
     case BinaryConsts::RefIsNull:
       return builder.makeRefIsNull();
     case BinaryConsts::RefFunc:
@@ -4282,34 +4279,36 @@ Result<> WasmBinaryReader::readInst() {
           return builder.makeI31Get(true);
         case BinaryConsts::I31GetU:
           return builder.makeI31Get(false);
-        case BinaryConsts::RefTest:
-          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::RefTest: {
+          auto [heapType, exactness] = getHeapType();
+          return builder.makeRefTest(Type(heapType, NonNullable, exactness));
+        }
+        case BinaryConsts::RefTestNull: {
+          auto [heapType, exactness] = getHeapType();
+          return builder.makeRefTest(Type(heapType, Nullable, exactness));
+        }
+        case BinaryConsts::RefCast: {
+          auto [heapType, exactness] = getHeapType();
+          return builder.makeRefCast(Type(heapType, NonNullable, exactness));
+        }
+        case BinaryConsts::RefCastNull: {
+          auto [heapType, exactness] = getHeapType();
+          return builder.makeRefCast(Type(heapType, Nullable, exactness));
+        }
         case BinaryConsts::BrOnCast:
         case BinaryConsts::BrOnCastFail: {
           auto flags = getInt8();
-          auto label = getU32LEB();
           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 label = getU32LEB();
+          auto [srcType, srcExact] = getHeapType();
+          auto [dstType, dstExact] = getHeapType();
+          auto in = Type(srcType, srcNull, srcExact);
+          auto cast = Type(dstType, dstNull, dstExact);
           auto kind = op == BinaryConsts::BrOnCast ? BrOnCast : BrOnCastFail;
           return builder.makeBrOn(label, kind, in, cast);
         }