Remove unused types during type optimizations (#5361)

The type rewriting utility in type-updating.cpp gathers all the used heap types,
then rewrites them to newly built and possibly modified heap types. The problem
is that for the isorecursive type system, the set of "used" heap types was
overly broad because it also included unused heap types that are in a rec group
with used types. In the context of emitting a binary, it is important to treat
these types as used because failing to emit them would change the identity of
the used types, but in the context of type optimizations it is ok to treat them
as truly unused because we are changing type identities anyway.

Update the type rewriting utility to only include truly used types in the set of
output types. This causes all existing type optimizations to implicitly drop
unused types, but only if they find any other optimizations to do and actually
run the rewriter utitility. Their output will also still include unused types
that were used before their optimizations were applied.

To overcome these limitations and better match the optimizing power of nominal
mode, which never includes unused types in the output, add a new type
optimization pass that removes unused types and does nothing else and run it
near the end of the global optimization pipeline.

5 files changed, 84 insertions, 10 deletions

diff --git a/src/ir/module-utils.cpp b/src/ir/module-utils.cpp
index f2c90dcfd..f3d9201b7 100644
--- a/src/ir/module-utils.cpp
+++ b/src/ir/module-utils.cpp
@@ -84,7 +84,9 @@ struct CodeScanner
       // not written in the binary format, so it doesn't need to be counted, but
       // it does need to be taken into account in the IR (this may be the only
       // place this type appears in the entire binary, and we must scan all
-      // types as the analyses that use us depend on that).
+      // types as the analyses that use us depend on that). TODO: This is kind
+      // of a hack, so it would be nice to remove. If we could remove it, we
+      // could also remove some of the pruning logic in getHeapTypeCounts below.
       counts.include(get->type);
     } else if (auto* set = curr->dynCast<StructSet>()) {
       counts.note(set->ref->type);
@@ -105,7 +107,10 @@ struct CodeScanner
   }
 };
 
-Counts getHeapTypeCounts(Module& wasm) {
+// Count the number of times each heap type that would appear in the binary is
+// referenced. If `prune`, exclude types that are never referenced, even though
+// a binary would be invalid without them.
+Counts getHeapTypeCounts(Module& wasm, bool prune = false) {
   // Collect module-level info.
   Counts counts;
   CodeScanner(wasm, counts).walkModuleCode(&wasm);
@@ -141,6 +146,19 @@ Counts getHeapTypeCounts(Module& wasm) {
     }
   }
 
+  if (prune) {
+    // Remove types that are not actually used.
+    auto it = counts.begin();
+    while (it != counts.end()) {
+      if (it->second == 0) {
+        auto deleted = it++;
+        counts.erase(deleted);
+      } else {
+        ++it;
+      }
+    }
+  }
+
   // Recursively traverse each reference type, which may have a child type that
   // is itself a reference type. This reflects an appearance in the binary
   // format that is in the type section itself. As we do this we may find more
@@ -178,12 +196,14 @@ Counts getHeapTypeCounts(Module& wasm) {
     }
 
     // Make sure we've noted the complete recursion group of each type as well.
-    auto recGroup = ht.getRecGroup();
-    if (includedGroups.insert(recGroup).second) {
-      for (auto type : recGroup) {
-        if (!counts.count(type)) {
-          newTypes.insert(type);
-          counts.include(type);
+    if (!prune) {
+      auto recGroup = ht.getRecGroup();
+      if (includedGroups.insert(recGroup).second) {
+        for (auto type : recGroup) {
+          if (!counts.count(type)) {
+            newTypes.insert(type);
+            counts.include(type);
+          }
         }
       }
     }
@@ -297,11 +317,11 @@ std::vector<HeapType> getPublicHeapTypes(Module& wasm) {
 }
 
 std::vector<HeapType> getPrivateHeapTypes(Module& wasm) {
-  auto allTypes = getHeapTypeCounts(wasm);
+  auto allTypes = getHeapTypeCounts(wasm, true);
   auto publicTypes = getPublicTypeSet(wasm);
   std::vector<HeapType> types;
   types.reserve(allTypes.size() - publicTypes.size());
-  for (auto [type, _] : allTypes) {
+  for (auto& [type, _] : allTypes) {
     if (!publicTypes.count(type)) {
       types.push_back(type);
     }
diff --git a/src/passes/CMakeLists.txt b/src/passes/CMakeLists.txt
index e97f5e388..f29e7db9b 100644
--- a/src/passes/CMakeLists.txt
+++ b/src/passes/CMakeLists.txt
@@ -94,6 +94,7 @@ set(passes_SOURCES
   RemoveUnusedBrs.cpp
   RemoveUnusedNames.cpp
   RemoveUnusedModuleElements.cpp
+  RemoveUnusedTypes.cpp
   ReorderFunctions.cpp
   ReorderGlobals.cpp
   ReorderLocals.cpp
diff --git a/src/passes/RemoveUnusedTypes.cpp b/src/passes/RemoveUnusedTypes.cpp
new file mode 100644
index 000000000..d346bc585
--- /dev/null
+++ b/src/passes/RemoveUnusedTypes.cpp
@@ -0,0 +1,48 @@
+/*
+ * Copyright 2022 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//
+// Remove unused types from private rec groups. This is done implicitly by other
+// type optimizations as well, but only if they find anything else they want to
+// optimize.
+//
+
+#include "ir/module-utils.h"
+#include "ir/type-updating.h"
+#include "pass.h"
+
+namespace wasm {
+
+namespace {
+
+struct RemoveUnusedTypes : Pass {
+  void run(Module* module) override {
+    if (!module->features.hasGC()) {
+      // This pass only does anything with GC types.
+      return;
+    }
+    // We're not changing the contents of any of the types, so we just round
+    // trip them throgh GlobalTypeRewriter, which will put all the private types
+    // in a single new rec group and leave out all the unused types.
+    GlobalTypeRewriter(*module).update();
+  }
+};
+
+} // anonymous namespace
+
+Pass* createRemoveUnusedTypesPass() { return new RemoveUnusedTypes(); }
+
+} // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index 21501c62f..732c2e8e2 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -369,6 +369,9 @@ void PassRegistry::registerPasses() {
   registerPass("remove-unused-names",
                "removes names from locations that are never branched to",
                createRemoveUnusedNamesPass);
+  registerPass("remove-unused-types",
+               "remove unused private GC types",
+               createRemoveUnusedTypesPass);
   registerPass("reorder-functions",
                "sorts functions by access frequency",
                createReorderFunctionsPass);
@@ -620,6 +623,7 @@ void PassRunner::addDefaultGlobalOptimizationPrePasses() {
     }
     addIfNoDWARFIssues("remove-unused-module-elements");
     if (options.closedWorld) {
+      addIfNoDWARFIssues("remove-unused-types");
       addIfNoDWARFIssues("cfp");
       addIfNoDWARFIssues("gsi");
     }
diff --git a/src/passes/passes.h b/src/passes/passes.h
index 8133c867e..4d6da0ca9 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -116,6 +116,7 @@ Pass* createRemoveUnusedBrsPass();
 Pass* createRemoveUnusedModuleElementsPass();
 Pass* createRemoveUnusedNonFunctionModuleElementsPass();
 Pass* createRemoveUnusedNamesPass();
+Pass* createRemoveUnusedTypesPass();
 Pass* createReorderFunctionsPass();
 Pass* createReorderGlobalsPass();
 Pass* createReorderGlobalsAlwaysPass();