Initial support for exact heap types

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