Avoid conflicts with public rec groups in MinimizeRecGroups (#6911)

As with all type optimizations, MinimizeRecGroups only changes private
types, which are the only types that are safe to modify. However, it is
important for correctness that MinimimizeRecGroups maintain separate
type identities for all types, whether public or private, to ensure that
casts that should differentiate two types cannot change behavior.

Previously the pass worked exclusively on private types, so there was
nothing preventing it from constructing a minimial rec group that
happened to have the same shape, and therefore type identity, as a
public rec group. #6886 exhibits a fuzzer test case where this happens
and changes the behavior of the program.

Fix the bug by recording all public rec group shapes and resolve
conflicts with these shapes by updating the shape of the conflicting
non-public type.

Fixes #6886.

1 files changed, 73 insertions, 7 deletions

diff --git a/src/passes/MinimizeRecGroups.cpp b/src/passes/MinimizeRecGroups.cpp
index 3a61e1e0b..8b38f9e4a 100644
--- a/src/passes/MinimizeRecGroups.cpp
+++ b/src/passes/MinimizeRecGroups.cpp
@@ -350,6 +350,14 @@ struct MinimizeRecGroups : Pass {
   // For each shape of rec group we have created, its index in `groups`.
   std::unordered_map<RecGroupShape, Index> groupShapeIndices;
 
+  // A sentinel index for public group shapes, which are recorded in
+  // `groupShapeIndices` but not in `groups`.
+  static constexpr Index PublicGroupIndex = -1;
+
+  // Since we won't store public groups in `groups`, we need this separate place
+  // to store their types referred to by their `RecGroupShape`s.
+  std::vector<std::vector<HeapType>> publicGroupTypes;
+
   // When we find that two groups are isomorphic to (i.e. permutations of) each
   // other, we combine their equivalence classes and choose a new class
   // representative to use to disambiguate the groups and any further groups
@@ -370,9 +378,32 @@ struct MinimizeRecGroups : Pass {
 
     initBrandOptions();
 
-    types = ModuleUtils::getPrivateHeapTypes(*module);
-    for (auto type : ModuleUtils::collectHeapTypes(*module)) {
-      typeIndices.insert({type, typeIndices.size()});
+    auto typeInfo = ModuleUtils::collectHeapTypeInfo(
+      *module,
+      ModuleUtils::TypeInclusion::AllTypes,
+      ModuleUtils::VisibilityHandling::FindVisibility);
+
+    types.reserve(typeInfo.size());
+
+    // We cannot optimize public types, but we do need to make sure we don't
+    // generate new groups with the same shape.
+    std::unordered_set<RecGroup> publicGroups;
+    for (auto& [type, info] : typeInfo) {
+      if (info.visibility == ModuleUtils::Visibility::Private) {
+        // We can optimize private types.
+        types.push_back(type);
+        typeIndices.insert({type, typeIndices.size()});
+      } else {
+        publicGroups.insert(type.getRecGroup());
+      }
+    }
+
+    publicGroupTypes.reserve(publicGroups.size());
+    for (auto group : publicGroups) {
+      publicGroupTypes.emplace_back(group.begin(), group.end());
+      [[maybe_unused]] auto [_, inserted] = groupShapeIndices.insert(
+        {RecGroupShape(publicGroupTypes.back()), PublicGroupIndex});
+      assert(inserted);
     }
 
     // The number of types to optimize is an upper bound on the number of
@@ -428,7 +459,15 @@ struct MinimizeRecGroups : Pass {
       return;
     }
 
-    // We have a conflict. There are five possibilities:
+    // We have a conflict. There are seven possibilities:
+    //
+    //   0. We have a conflict with a public group and...
+    //
+    //     A. We are trying to insert the next permutation of an existing
+    //        equivalence class.
+    //
+    //     B. We are inserting a new group not yet affiliated with a nontrivial
+    //        equivalence class.
     //
     //   1. We are trying to insert the next permutation of an existing
     //      equivalence class and have found that...
@@ -454,12 +493,39 @@ struct MinimizeRecGroups : Pass {
     //
     // These possibilities are handled in order below.
 
+    auto& groupInfo = groups[group];
+    Index groupRep = equivalenceClasses.getRoot(group);
+
     Index other = it->second;
 
-    auto& groupInfo = groups[group];
-    auto& otherInfo = groups[other];
+    if (other == PublicGroupIndex) {
+      // The group currently has the same shape as some public group. We can't
+      // change the public group, so we need to change the shape of the current
+      // group.
+      //
+      // Case 0A: Since we already have a nontrivial equivalence class, we can
+      // try the next permutation to avoid conflict with the public group.
+      if (groups[groupRep].classInfo) {
+        // We have everything we need to generate the next permutation of this
+        // group.
+        auto& classInfo = *groups[groupRep].classInfo;
+        classInfo.advance();
+        classInfo.permute(groupInfo);
+        shapesToUpdate.push_back(group);
+        return;
+      }
 
-    Index groupRep = equivalenceClasses.getRoot(group);
+      // Case 0B: There is no nontrivial equivalence class yet. Create one so
+      // we have all the information we need to try a different permutation.
+      assert(group == groupRep);
+      groupInfo.permutation = getCanonicalPermutation(groupInfo.group);
+      groupInfo.classInfo.emplace(groupInfo);
+      groupInfo.classInfo->permute(groupInfo);
+      shapesToUpdate.push_back(group);
+      return;
+    }
+
+    auto& otherInfo = groups[other];
     Index otherRep = equivalenceClasses.getRoot(other);
 
     // Case 1A: We have found a permutation of this group that has the same