Fix bug in UniqueRecGroups

src/ir/type-updating.cpp

@@ -41,9 +41,9 @@ GlobalTypeRewriter::GlobalTypeRewriter(Module& wasm)
   for (auto& [type, info] : typeInfo) {
     if (info.visibility == ModuleUtils::Visibility::Public) {
       auto group = type.getRecGroup();
-      if (seenGroups.insert(type.getRecGroup()).second) {
-        std::vector<HeapType> groupTypes(group.begin(), group.end());
-        publicGroups.insert(std::move(groupTypes));
+      if (seenGroups.insert(group).second) {
+        [[maybe_unused]] RecGroup unique = publicGroups.insert(group);
+        assert(unique == group);
       }
     }
   }
@@ -199,12 +199,12 @@ GlobalTypeRewriter::rebuildTypes(std::vector<HeapType> types) {
   }
 #endif
   // Ensure the new types are different from any public rec group.
-  const auto& newTypes = publicGroups.insert(*buildResults);
+  RecGroup newGroup = publicGroups.insert((*buildResults)[0].getRecGroup());
 
   // Map the old types to the new ones.
   TypeMap oldToNewTypes;
   for (auto [type, index] : typeIndices) {
-    oldToNewTypes[type] = newTypes[index];
+    oldToNewTypes[type] = newGroup[index];
   }
   mapTypeNamesAndIndices(oldToNewTypes);
   return oldToNewTypes;

src/passes/TypeSSA.cpp

@@ -61,33 +61,31 @@ namespace wasm {
 
 namespace {
 
-// Ensure there are no conflicts between the newly built types and any existing
-// types.
-std::vector<HeapType> ensureTypesAreInNewRecGroup(std::vector<HeapType>&& types,
-                                                  Module& wasm) {
+// Ensure there are no conflicts between the newly built types in the given
+// RecGroup and any existing types.
+std::vector<HeapType> ensureRecGroupIsUnique(RecGroup group, Module& wasm) {
   std::unordered_set<RecGroup> existing;
   for (auto type : ModuleUtils::collectHeapTypes(wasm)) {
     existing.insert(type.getRecGroup());
   }
 
   UniqueRecGroups unique(wasm.features);
   for (auto group : existing) {
-    std::vector<HeapType> types(group.begin(), group.end());
     // N.B. we use `insertOrGet` rather than `insert` because some passes (DAE,
     // BlockMerging) can create multiple types with the same shape, so we can't
     // assume all the rec groups are already unique. The rec groups will have
     // different types, but their shapes will match considering how exactness,
     // etc. will be erased by the binary writer.
-    unique.insertOrGet(std::move(types));
+    unique.insertOrGet(group);
   }
 
-  auto num = types.size();
-  std::vector<HeapType> uniqueTypes = unique.insert(std::move(types));
-  if (uniqueTypes.size() != num) {
+  RecGroup uniqueGroup = unique.insert(group);
+  std::vector<HeapType> uniqueTypes(uniqueGroup.begin(), uniqueGroup.end());
+  if (uniqueTypes.size() != group.size()) {
     // Remove the brand type, which we do not need to consider further.
     uniqueTypes.pop_back();
   }
-  assert(uniqueTypes.size() == num);
+  assert(uniqueTypes.size() == group.size());
   return uniqueTypes;
 }
 
@@ -355,7 +353,7 @@ struct TypeSSA : public Pass {
     assert(newTypes.size() == num);
 
     // Make sure this is actually a new rec group.
-    newTypes = ensureTypesAreInNewRecGroup(std::move(newTypes), *module);
+    newTypes = ensureRecGroupIsUnique(newTypes[0].getRecGroup(), *module);
 
     // Success: we can apply the new types.
 

src/wasm-type-shape.h

@@ -164,20 +164,19 @@ struct BrandTypeIterator {
 struct UniqueRecGroups {
   using Groups = std::list<std::vector<HeapType>>;
   Groups groups;
-  std::unordered_map<RecGroupShape, Groups::iterator> shapes;
+  std::unordered_map<RecGroupShape, RecGroup> shapes;
 
   FeatureSet features;
 
   UniqueRecGroups(FeatureSet features) : features(features) {}
 
-  // Insert a rec group. If it is already unique, return the original types.
-  // Otherwise rebuild the group to make it unique and return the rebuilt types,
-  // including the brand.
-  const std::vector<HeapType>& insert(std::vector<HeapType> group);
+  // Insert a rec group. If it is already unique, return the original group.
+  // Otherwise rebuild the group to make it unique and return the rebuilt group.
+  RecGroup insert(RecGroup group);
 
   // If the group is unique, insert it and return the types. Otherwise, return
   // the types that already have this shape.
-  const std::vector<HeapType>& insertOrGet(std::vector<HeapType> group);
+  RecGroup insertOrGet(RecGroup group);
 };
 
 } // namespace wasm

src/wasm/wasm-type-shape.cpp

@@ -373,51 +373,51 @@ bool ComparableRecGroupShape::operator>(const RecGroupShape& other) const {
   return GT == compareComparable(*this, other);
 }
 
-const std::vector<HeapType>&
-UniqueRecGroups::insert(std::vector<HeapType> types) {
-  auto groupIt = groups.emplace(groups.end(), std::move(types));
-  auto& group = *groupIt;
-  if (shapes.emplace(RecGroupShape(group, features), groupIt).second) {
+RecGroup UniqueRecGroups::insert(RecGroup group) {
+  auto typesIt = groups.emplace(groups.end(), group.begin(), group.end());
+  auto& types = *typesIt;
+  if (shapes.emplace(RecGroupShape(types, features), group).second) {
     // The types are already unique.
     return group;
   }
   // There is a conflict. Find a brand that makes the group unique.
   BrandTypeIterator brand;
-  group.push_back(*brand);
-  while (!shapes.emplace(RecGroupShape(group, features), groupIt).second) {
-    group.back() = *++brand;
-  }
-  // Rebuild the rec group to include the brand. Map the old types (excluding
-  // the brand) to their corresponding new types to preserve recursions within
-  // the group.
-  Index size = group.size();
-  TypeBuilder builder(size);
-  std::unordered_map<HeapType, HeapType> newTypes;
-  for (Index i = 0; i < size - 1; ++i) {
-    newTypes[group[i]] = builder[i];
-  }
-  for (Index i = 0; i < size; ++i) {
-    builder[i].copy(group[i], [&](HeapType type) {
-      if (auto newType = newTypes.find(type); newType != newTypes.end()) {
-        return newType->second;
-      }
-      return type;
-    });
-  }
-  builder.createRecGroup(0, size);
-  group = *builder.build();
+  types.push_back(*brand);
+  Index size = types.size();
+  do {
+    types.back() = *brand;
+    ++brand;
+    TypeBuilder builder(size);
+    // Map the old types (excluding the brand) to their corresponding new types
+    // to preserve recursions within the group.
+    std::unordered_map<HeapType, HeapType> newTypes;
+    for (Index i = 0; i < size - 1; ++i) {
+      newTypes[group[i]] = builder[i];
+    }
+    for (Index i = 0; i < size; ++i) {
+      builder[i].copy(types[i], [&](HeapType type) {
+        if (auto newType = newTypes.find(type); newType != newTypes.end()) {
+          return newType->second;
+        }
+        return type;
+      });
+    }
+    builder.createRecGroup(0, size);
+    types = *builder.build();
+    group = types[0].getRecGroup();
+  } while (!shapes.emplace(RecGroupShape(types, features), group).second);
+
   return group;
 }
 
-const std::vector<HeapType>&
-UniqueRecGroups::insertOrGet(std::vector<HeapType> types) {
-  auto groupIt = groups.emplace(groups.end(), std::move(types));
+RecGroup UniqueRecGroups::insertOrGet(RecGroup group) {
+  auto typesIt = groups.emplace(groups.end(), group.begin(), group.end());
   auto [it, inserted] =
-    shapes.emplace(RecGroupShape(*groupIt, features), groupIt);
+    shapes.emplace(RecGroupShape(*typesIt, features), group);
   if (!inserted) {
-    groups.erase(groupIt);
+    groups.erase(typesIt);
   }
-  return *it->second;
+  return it->second;
 }
 
 } // namespace wasm

src/wasm/wasm-type.cpp

@@ -2705,9 +2705,9 @@ TypeBuilder::BuildResult TypeBuilder::build() {
     // If we are building multiple groups, make sure there will be no conflicts
     // after disallowed features are taken into account.
     if (groupSize > 0 && groupSize != entryCount) {
-      auto expectedFirst = (*built)[0];
-      auto& types = impl->unique.insertOrGet(*built);
-      if (types[0] != expectedFirst) {
+      auto group = (*built)[0].getRecGroup();
+      auto uniqueGroup = impl->unique.insertOrGet(group);
+      if (group != uniqueGroup) {
         return {TypeBuilder::Error{
           groupStart, TypeBuilder::ErrorReason::RecGroupCollision}};
       }