Add a TypeBuilder API for copying a heap type

src/ir/type-updating.cpp

@@ -88,55 +88,47 @@ GlobalTypeRewriter::TypeMap GlobalTypeRewriter::rebuildTypes(
 
   // Create the temporary heap types.
   i = 0;
+  auto map = [&](HeapType type) -> HeapType {
+    if (auto it = typeIndices.find(type); it != typeIndices.end()) {
+      return typeBuilder[it->second];
+    }
+    return type;
+  };
   for (auto [type, _] : typeIndices) {
-    typeBuilder[i].setOpen(type.isOpen());
-    typeBuilder[i].setShared(type.getShared());
-    if (type.isSignature()) {
-      auto sig = type.getSignature();
-      TypeList newParams, newResults;
-      for (auto t : sig.params) {
-        newParams.push_back(getTempType(t));
+    typeBuilder[i].copy(type, map);
+    switch (type.getKind()) {
+      case HeapTypeKind::Func: {
+        auto newSig = HeapType(typeBuilder[i]).getSignature();
+        modifySignature(type, newSig);
+        typeBuilder[i] = newSig;
+        break;
       }
-      for (auto t : sig.results) {
-        newResults.push_back(getTempType(t));
+      case HeapTypeKind::Struct: {
+        auto newStruct = HeapType(typeBuilder[i]).getStruct();
+        modifyStruct(type, newStruct);
+        typeBuilder[i] = newStruct;
+        break;
       }
-      Signature newSig(typeBuilder.getTempTupleType(newParams),
-                       typeBuilder.getTempTupleType(newResults));
-      modifySignature(type, newSig);
-      typeBuilder[i] = newSig;
-    } else if (type.isStruct()) {
-      auto struct_ = type.getStruct();
-      // Start with a copy to get mutability/packing/etc.
-      auto newStruct = struct_;
-      for (auto& field : newStruct.fields) {
-        field.type = getTempType(field.type);
+      case HeapTypeKind::Array: {
+        auto newArray = HeapType(typeBuilder[i]).getArray();
+        modifyArray(type, newArray);
+        typeBuilder[i] = newArray;
+        break;
       }
-      modifyStruct(type, newStruct);
-      typeBuilder[i] = newStruct;
-    } else if (type.isArray()) {
-      auto array = type.getArray();
-      // Start with a copy to get mutability/packing/etc.
-      auto newArray = array;
-      newArray.element.type = getTempType(newArray.element.type);
-      modifyArray(type, newArray);
-      typeBuilder[i] = newArray;
-    } else {
-      WASM_UNREACHABLE("bad type");
+      case HeapTypeKind::Cont:
+        WASM_UNREACHABLE("TODO: cont");
+      case HeapTypeKind::Basic:
+        WASM_UNREACHABLE("unexpected kind");
     }
 
-    // Apply a super, if there is one
     if (auto super = getDeclaredSuperType(type)) {
-      if (auto it = typeIndices.find(*super); it != typeIndices.end()) {
-        assert(it->second < i);
-        typeBuilder[i].subTypeOf(typeBuilder[it->second]);
-      } else {
-        typeBuilder[i].subTypeOf(*super);
-      }
+      typeBuilder[i].subTypeOf(map(*super));
+    } else {
+      typeBuilder[i].subTypeOf(std::nullopt);
     }
 
     modifyTypeBuilderEntry(typeBuilder, i, type);
-
-    i++;
+    ++i;
   }
 
   auto buildResults = typeBuilder.build();
@@ -316,8 +308,7 @@ Type GlobalTypeRewriter::getTempType(Type type) {
     return type;
   }
   if (type.isTuple()) {
-    auto& tuple = type.getTuple();
-    auto newTuple = tuple;
+    auto newTuple = type.getTuple();
     for (auto& t : newTuple) {
       t = getTempType(t);
     }

src/passes/TypeSSA.cpp

@@ -98,20 +98,7 @@ std::vector<HeapType> ensureTypesAreInNewRecGroup(RecGroup recGroup,
       // Make a builder and add a slot for the hash.
       TypeBuilder builder(num + 1);
       for (Index i = 0; i < num; i++) {
-        auto type = types[i];
-        if (type.isStruct()) {
-          builder[i] = type.getStruct();
-        } else {
-          // Atm this pass only needs struct and array types. If we refactor
-          // this function to be general purpose we'd need to extend that. TODO
-          assert(type.isArray());
-          builder[i] = type.getArray();
-        }
-        if (auto super = type.getDeclaredSuperType()) {
-          builder[i].subTypeOf(*super);
-        }
-        builder[i].setOpen(type.isOpen());
-        builder[i].setShared(type.getShared());
+        builder[i].copy(types[i]);
       }
 
       // Implement the hash as a struct with "random" fields, and add it.

src/wasm-type.h

@@ -660,6 +660,67 @@ struct TypeBuilder {
   void setHeapType(size_t i, Struct&& struct_);
   void setHeapType(size_t i, Array array);
 
+  // Sets the heap type at index `i` to be a copy of the given heap type with
+  // its referenced HeapTypes to be replaced according to the provided mapping
+  // function.
+  template<typename F> void copyHeapType(size_t i, HeapType type, F map) {
+    assert(!type.isBasic());
+    if (auto super = type.getDeclaredSuperType()) {
+      setSubType(i, map(*super));
+    }
+    setOpen(i, type.isOpen());
+    setShared(i, type.getShared());
+
+    auto copySingleType = [&](Type t) -> Type {
+      if (t.isBasic()) {
+        return t;
+      }
+      assert(t.isRef());
+      return getTempRefType(map(t.getHeapType()), t.getNullability());
+    };
+    auto copyType = [&](Type t) -> Type {
+      if (t.isTuple()) {
+        std::vector<Type> elems;
+        elems.reserve(t.size());
+        for (auto elem : t) {
+          elems.push_back(copySingleType(elem));
+        }
+        return getTempTupleType(elems);
+      }
+      return copySingleType(t);
+    };
+    switch (type.getKind()) {
+      case HeapTypeKind::Func: {
+        auto sig = type.getSignature();
+        setHeapType(i, Signature(copyType(sig.params), copyType(sig.results)));
+        return;
+      }
+      case HeapTypeKind::Struct: {
+        const auto& struct_ = type.getStruct();
+        std::vector<Field> fields;
+        fields.reserve(struct_.fields.size());
+        for (auto field : struct_.fields) {
+          field.type = copyType(field.type);
+          fields.push_back(field);
+        }
+        setHeapType(i, Struct(fields));
+        return;
+      }
+      case HeapTypeKind::Array: {
+        auto elem = type.getArray().element;
+        elem.type = copyType(elem.type);
+        // MSVC gets confused without this disambiguation.
+        setHeapType(i, wasm::Array(elem));
+        return;
+      }
+      case HeapTypeKind::Cont:
+        setHeapType(i, Continuation(map(type.getContinuation().type)));
+        return;
+      case HeapTypeKind::Basic:
+        WASM_UNREACHABLE("unexpected kind");
+    }
+  }
+
   // Gets the temporary HeapType at index `i`. This HeapType should only be used
   // to construct temporary Types using the methods below.
   HeapType getTempHeapType(size_t i);
@@ -672,7 +733,7 @@ struct TypeBuilder {
 
   // Declare the HeapType being built at index `i` to be an immediate subtype of
   // the given HeapType.
-  void setSubType(size_t i, HeapType super);
+  void setSubType(size_t i, std::optional<HeapType> super);
 
   // Create a new recursion group covering slots [i, i + length). Groups must
   // not overlap or go out of bounds.
@@ -746,7 +807,7 @@ struct TypeBuilder {
       builder.setHeapType(index, array);
       return *this;
     }
-    Entry& subTypeOf(HeapType other) {
+    Entry& subTypeOf(std::optional<HeapType> other) {
       builder.setSubType(index, other);
       return *this;
     }
@@ -758,6 +819,13 @@ struct TypeBuilder {
       builder.setShared(index, share);
       return *this;
     }
+    template<typename F> Entry& copy(HeapType type, F map) {
+      builder.copyHeapType(index, type, map);
+      return *this;
+    }
+    Entry& copy(HeapType type) {
+      return copy(type, [](HeapType t) { return t; });
+    }
   };
 
   Entry operator[](size_t i) { return Entry{*this, i}; }

src/wasm/wasm-type.cpp

@@ -2532,10 +2532,10 @@ Type TypeBuilder::getTempRefType(HeapType type, Nullability nullable) {
   return markTemp(impl->typeStore.insert(TypeInfo(type, nullable)));
 }
 
-void TypeBuilder::setSubType(size_t i, HeapType super) {
+void TypeBuilder::setSubType(size_t i, std::optional<HeapType> super) {
   assert(i < size() && "index out of bounds");
   HeapTypeInfo* sub = impl->entries[i].info.get();
-  sub->supertype = getHeapTypeInfo(super);
+  sub->supertype = super ? getHeapTypeInfo(*super) : nullptr;
 }
 
 void TypeBuilder::createRecGroup(size_t index, size_t length) {