Initial support for exact heap types (#7396)

The custom descriptors proposal has moved exactness from reference types
to defined (but not abstract) heap types. Since we only use a bit in the
heap type representation to represent sharedness for abstract heap
types, we can conveniently reuse the same bit to represent exactness for
heap types.

Implement basic support for representing exact heap types and taking
them into account in canonicalization. Also ensure that other operations
like getting the rec group of a heap type or looking up its structure
work properly on exact heap types.

Author: tlively

src/wasm-type.h

@@ -62,6 +62,7 @@ using Tuple = TypeList;
 
 enum Nullability { NonNullable, Nullable };
 enum Mutability { Immutable, Mutable };
+enum Exactness { Inexact, Exact };
 
 // HeapType name information used for printing.
 struct TypeNames {
@@ -98,10 +99,12 @@ class HeapType {
   static constexpr int TypeBits = 2;
   static constexpr int UsedBits = TypeBits + 1;
   static constexpr int SharedMask = 1 << TypeBits;
+  static constexpr int ExactMask = SharedMask;
 
 public:
-  // Bits 0-1 are used by the Type representation, so need to be left free.
-  // Bit 2 determines whether the basic heap type is shared (1) or unshared (0).
+  // Bits 0-1 are used by the Type representation, so need to be left free. Bit
+  // 2 determines whether a basic heap type is shared (1) or unshared (0). For
+  // non-basic heap types, bit 2 determines whether the type is exact instead.
   enum BasicHeapType : uint32_t {
     ext = 1 << UsedBits,
     func = 2 << UsedBits,
@@ -126,7 +129,7 @@ class HeapType {
   constexpr HeapType(BasicHeapType id) : id(id) {}
 
   // But converting raw TypeID is more dangerous, so make it explicit
-  explicit HeapType(TypeID id) : id(id) {}
+  explicit constexpr HeapType(TypeID id) : id(id) {}
 
   // Choose an arbitrary heap type as the default.
   constexpr HeapType() : HeapType(func) {}
@@ -167,8 +170,12 @@ class HeapType {
   bool isBottom() const;
   bool isOpen() const;
   bool isShared() const { return getShared() == Shared; }
+  bool isExact() const { return getExactness() == Exact; }
 
   Shareability getShared() const;
+  Exactness getExactness() const {
+    return !isBasic() && (id & ExactMask) ? Exact : Inexact;
+  }
 
   // Check if the type is a given basic heap type, while ignoring whether it is
   // shared or not.
@@ -217,15 +224,29 @@ class HeapType {
   // Get the index of this non-basic type within its recursion group.
   size_t getRecGroupIndex() const;
 
-  constexpr TypeID getID() const { return id; }
-
   // Get the shared or unshared version of this basic heap type.
   constexpr BasicHeapType getBasic(Shareability share) const {
     assert(isBasic());
     return BasicHeapType(share == Shared ? (id | SharedMask)
                                          : (id & ~SharedMask));
   }
 
+  constexpr HeapType with(Exactness exactness) const {
+    assert((!isBasic() || exactness == Inexact) &&
+           "abstract types cannot be exact");
+    return HeapType(exactness == Exact ? (id | ExactMask) : (id & ~ExactMask));
+  }
+
+  // The ID is the numeric representation of the heap type and can be used in
+  // FFI or hashing applications. The "raw" ID is the numeric representation of
+  // the plain version of the type without exactness or any other attributes we
+  // might add in the future. It's useful in contexts where all heap types using
+  // the same type definition need to be treated identically.
+  constexpr TypeID getID() const { return id; }
+  constexpr TypeID getRawID() const {
+    return isBasic() ? id : with(Inexact).id;
+  }
+
   // (In)equality must be defined for both HeapType and BasicHeapType because it
   // is otherwise ambiguous whether to convert both this and other to int or
   // convert other to HeapType.

src/wasm/wasm-type.cpp

@@ -228,7 +228,7 @@ namespace {
 
 HeapTypeInfo* getHeapTypeInfo(HeapType ht) {
   assert(!ht.isBasic());
-  return (HeapTypeInfo*)ht.getID();
+  return (HeapTypeInfo*)(ht.getRawID());
 }
 
 HeapType asHeapType(std::unique_ptr<HeapTypeInfo>& info) {
@@ -1247,7 +1247,7 @@ RecGroup HeapType::getRecGroup() const {
   } else {
     // Mark the low bit to signify that this is a trivial recursion group and
     // points to a heap type info rather than a vector of heap types.
-    return RecGroup(id | 1);
+    return RecGroup(getRawID() | 1);
   }
 }
 
@@ -1608,14 +1608,20 @@ bool SubTyper::isSubType(const Array& a, const Array& b) {
 }
 
 void TypePrinter::printHeapTypeName(HeapType type) {
+  if (type.isExact()) {
+    os << "(exact ";
+  }
   if (type.isBasic()) {
     print(type);
-    return;
-  }
-  generator(type).name.print(os);
+  } else {
+    generator(type).name.print(os);
 #if TRACE_CANONICALIZATION
-  os << "(;" << ((type.getID() >> 4) % 1000) << ";) ";
+    os << "(;" << ((type.getID() >> 4) % 1000) << ";) ";
 #endif
+  }
+  if (type.isExact()) {
+    os << ')';
+  }
 }
 
 std::ostream& TypePrinter::print(Type type) {
@@ -1942,8 +1948,10 @@ size_t RecGroupHasher::hash(HeapType type) const {
     wasm::rehash(digest, type.getID());
     return digest;
   }
+  wasm::rehash(digest, type.isExact());
   wasm::rehash(digest, type.getRecGroupIndex());
   auto currGroup = type.getRecGroup();
+  wasm::rehash(digest, currGroup != group);
   if (currGroup != group) {
     wasm::rehash(digest, currGroup.getID());
   }
@@ -2073,6 +2081,9 @@ bool RecGroupEquator::eq(HeapType a, HeapType b) const {
   if (a.isBasic() || b.isBasic()) {
     return a == b;
   }
+  if (a.getExactness() != b.getExactness()) {
+    return false;
+  }
   if (a.getRecGroupIndex() != b.getRecGroupIndex()) {
     return false;
   }
@@ -2456,15 +2467,18 @@ void updateReferencedHeapTypes(
       isTopLevel = false;
       if (type->isRef()) {
         auto ht = type->getHeapType();
+        auto exact = ht.getExactness();
+        ht = ht.with(Inexact);
         if (auto it = canonicalized.find(ht); it != canonicalized.end()) {
-          *type = Type(it->second, type->getNullability());
+          *type = Type(it->second.with(exact), type->getNullability());
         }
       } else if (type->isTuple()) {
         TypeGraphWalkerBase<ChildUpdater>::scanType(type);
       }
     }
 
     void scanHeapType(HeapType* type) {
+      assert(!type->isExact() && "unexpected exact type in definition");
       if (isTopLevel) {
         isTopLevel = false;
         TypeGraphWalkerBase<ChildUpdater>::scanHeapType(type);
@@ -2529,7 +2543,8 @@ buildRecGroup(std::unique_ptr<RecGroupInfo>&& groupInfo,
   for (size_t i = 0; i < typeInfos.size(); ++i) {
     auto type = asHeapType(typeInfos[i]);
     for (auto child : type.getHeapTypeChildren()) {
-      if (isTemp(child) && !seenTypes.count(child)) {
+      HeapType rawChild(child.getRawID());
+      if (isTemp(rawChild) && !seenTypes.count(rawChild)) {
         return {TypeBuilder::Error{
           i, TypeBuilder::ErrorReason::ForwardChildReference}};
       }

test/gtest/type-builder.cpp

@@ -428,6 +428,93 @@ TEST_F(TypeTest, CanonicalizeUses) {
   EXPECT_NE(built[4], built[6]);
 }
 
+TEST_F(TypeTest, CanonicalizeExactHeapTypes) {
+  TypeBuilder builder(8);
+
+  HeapType inexact = HeapType(builder[0]).with(Inexact);
+  HeapType exact = HeapType(builder[1]).with(Exact);
+
+  Type inexactRef = builder.getTempRefType(inexact, Nullable);
+  Type exactRef = builder.getTempRefType(exact, Nullable);
+
+  // Types that vary in exactness of the referenced heap type are different.
+  builder[0] = Struct({Field(inexactRef, Mutable)});
+  builder[1] = Struct({Field(exactRef, Mutable)});
+  builder[2] = Signature(Type({inexactRef, exactRef}), Type::none);
+  builder[3] = Signature(Type::none, Type({exactRef, inexactRef}));
+
+  auto translate = [&](HeapType t) {
+    for (int i = 0; i < 4; ++i) {
+      if (t.with(Inexact) == builder[i]) {
+        return HeapType(builder[4 + i]).with(t.getExactness());
+      }
+    }
+    WASM_UNREACHABLE("unexpected type");
+  };
+
+  builder[4].copy(builder[0], translate);
+  builder[5].copy(builder[1], translate);
+  builder[6].copy(builder[2], translate);
+  builder[7].copy(builder[3], translate);
+
+  auto result = builder.build();
+  ASSERT_TRUE(result);
+  auto built = *result;
+
+  // Different types should be different.
+  EXPECT_NE(built[0], built[1]);
+  EXPECT_NE(built[0], built[2]);
+  EXPECT_NE(built[0], built[3]);
+  EXPECT_NE(built[1], built[2]);
+  EXPECT_NE(built[1], built[3]);
+  EXPECT_NE(built[2], built[3]);
+
+  // Copies of the types should match.
+  EXPECT_EQ(built[0], built[4]);
+  EXPECT_EQ(built[1], built[5]);
+  EXPECT_EQ(built[2], built[6]);
+  EXPECT_EQ(built[3], built[7]);
+
+  // A type is inexact by default.
+  EXPECT_EQ(built[0], built[0].with(Inexact));
+  EXPECT_EQ(built[1], built[1].with(Inexact));
+  EXPECT_EQ(built[2], built[2].with(Inexact));
+  EXPECT_EQ(built[3], built[3].with(Inexact));
+
+  // We can freely convert between exact and inexact.
+  EXPECT_EQ(built[0], built[0].with(Exact).with(Inexact));
+  EXPECT_EQ(built[0].with(Exact),
+            built[0].with(Exact).with(Inexact).with(Exact));
+
+  // Conversions are idempotent.
+  EXPECT_EQ(built[0].with(Exact), built[0].with(Exact).with(Exact));
+  EXPECT_EQ(built[0], built[0].with(Inexact));
+
+  // An exact version of a type is not the same as its inexact version.
+  EXPECT_NE(built[0].with(Exact), built[0].with(Inexact));
+
+  // But they have the same rec group.
+  EXPECT_EQ(built[0].with(Exact).getRecGroup(),
+            built[0].with(Inexact).getRecGroup());
+
+  // Looking up the inner structure works either way.
+  ASSERT_TRUE(built[0].with(Exact).isStruct());
+  ASSERT_TRUE(built[0].with(Inexact).isStruct());
+  EXPECT_EQ(built[0].with(Exact).getStruct(),
+            built[0].with(Inexact).getStruct());
+
+  // The exactness of children types is preserved.
+  EXPECT_EQ(built[0], built[0].getStruct().fields[0].type.getHeapType());
+  EXPECT_EQ(built[1].with(Exact),
+            built[1].getStruct().fields[0].type.getHeapType());
+  EXPECT_EQ(built[0], built[2].getSignature().params[0].getHeapType());
+  EXPECT_EQ(built[1].with(Exact),
+            built[2].getSignature().params[1].getHeapType());
+  EXPECT_EQ(built[0], built[3].getSignature().results[1].getHeapType());
+  EXPECT_EQ(built[1].with(Exact),
+            built[3].getSignature().results[0].getHeapType());
+}
+
 TEST_F(TypeTest, CanonicalizeSelfReferences) {
   TypeBuilder builder(5);
   // Single self-reference