Add passes to finalize or unfinalize types (#5944)

TypeFinalization finalizes all types that we can, that is, all private types that have no
children. TypeUnFinalization unfinalizes (opens) all (private) types.

These could be used by first opening all types, optimizing, and then finalizing, as that
might find more opportunities.

Fixes #5933

6 files changed, 108 insertions, 0 deletions

diff --git a/src/ir/type-updating.cpp b/src/ir/type-updating.cpp
index cca0d2360..59b77cca1 100644
--- a/src/ir/type-updating.cpp
+++ b/src/ir/type-updating.cpp
@@ -115,6 +115,9 @@ GlobalTypeRewriter::TypeMap GlobalTypeRewriter::rebuildTypes() {
         typeBuilder[i].subTypeOf(*super);
       }
     }
+
+    modifyTypeBuilderEntry(typeBuilder, i, type);
+
     i++;
   }
 
diff --git a/src/ir/type-updating.h b/src/ir/type-updating.h
index 31fd77c69..ddeb74b31 100644
--- a/src/ir/type-updating.h
+++ b/src/ir/type-updating.h
@@ -373,6 +373,14 @@ public:
   virtual void modifyArray(HeapType oldType, Array& array) {}
   virtual void modifySignature(HeapType oldType, Signature& sig) {}
 
+  // This additional hook is called after modify* and other operations, and
+  // allows the caller to do things like typeBuilder[i].setOpen(false);
+  //
+  // This is provided the builder, the index we are on, and the old heap type
+  // for that index.
+  virtual void
+  modifyTypeBuilderEntry(TypeBuilder& typeBuilder, Index i, HeapType oldType) {}
+
   // Subclasses can override this method to modify supertypes. The new
   // supertype, if any, must be a supertype (or the same as) the original
   // supertype.
diff --git a/src/passes/CMakeLists.txt b/src/passes/CMakeLists.txt
index a47333cf6..bd29bfa0b 100644
--- a/src/passes/CMakeLists.txt
+++ b/src/passes/CMakeLists.txt
@@ -114,6 +114,7 @@ set(passes_SOURCES
   StripEH.cpp
   SSAify.cpp
   TupleOptimization.cpp
+  TypeFinalizing.cpp
   Untee.cpp
   Vacuum.cpp
   ${CMAKE_CURRENT_BINARY_DIR}/WasmIntrinsics.cpp
diff --git a/src/passes/TypeFinalizing.cpp b/src/passes/TypeFinalizing.cpp
new file mode 100644
index 000000000..ebf04318a
--- /dev/null
+++ b/src/passes/TypeFinalizing.cpp
@@ -0,0 +1,88 @@
+/*
+ * Copyright 2023 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.
+ */
+
+//
+// Make all types final or unfinal (open). A typical workflow might be to un-
+// finalize all types, optimize a lot, and then finalize at the end.
+//
+
+#include "ir/module-utils.h"
+#include "ir/subtypes.h"
+#include "ir/type-updating.h"
+#include "pass.h"
+#include "wasm-type.h"
+#include "wasm.h"
+
+namespace wasm {
+
+namespace {
+
+struct TypeFinalizing : public Pass {
+  // Whether we should finalize or unfinalize types.
+  bool finalize;
+
+  // We will only modify types that we are allowed to.
+  std::unordered_set<HeapType> modifiableTypes;
+
+  TypeFinalizing(bool finalize) : finalize(finalize) {}
+
+  void run(Module* module) override {
+    if (!module->features.hasGC()) {
+      return;
+    }
+
+    // To make a type final, it must have no subtypes.
+    std::optional<SubTypes> subTypes;
+    if (finalize) {
+      subTypes = SubTypes(*module);
+    }
+
+    auto privateTypes = ModuleUtils::getPrivateHeapTypes(*module);
+    for (auto type : privateTypes) {
+      // If we are finalizing types then we can only do that to leaf types. If
+      // we are unfinalizing, we can do that unconditionally.
+      if (!finalize || subTypes->getImmediateSubTypes(type).empty()) {
+        modifiableTypes.insert(type);
+      }
+    }
+
+    class TypeRewriter : public GlobalTypeRewriter {
+      TypeFinalizing& parent;
+
+    public:
+      TypeRewriter(Module& wasm, TypeFinalizing& parent)
+        : GlobalTypeRewriter(wasm), parent(parent) {}
+
+      void modifyTypeBuilderEntry(TypeBuilder& typeBuilder,
+                                  Index i,
+                                  HeapType oldType) override {
+        if (parent.modifiableTypes.count(oldType)) {
+          typeBuilder[i].setOpen(!parent.finalize);
+        }
+      }
+    };
+
+    TypeRewriter(*module, *this).update();
+  }
+};
+
+} // anonymous namespace
+
+Pass* createTypeFinalizingPass() { return new TypeFinalizing(true); }
+
+Pass* createTypeUnFinalizingPass() { return new TypeFinalizing(false); }
+
+} // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index 15e8c1f26..3880519da 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -483,12 +483,18 @@ void PassRegistry::registerPasses() {
   registerPass("tuple-optimization",
                "optimize trivial tuples away",
                createTupleOptimizationPass);
+  registerPass("type-finalizing",
+               "mark all leaf types as final",
+               createTypeFinalizingPass);
   registerPass("type-merging",
                "merge types to their supertypes where possible",
                createTypeMergingPass);
   registerPass("type-ssa",
                "create new nominal types to help other optimizations",
                createTypeSSAPass);
+  registerPass("type-unfinalizing",
+               "mark all types as non-final (open)",
+               createTypeUnFinalizingPass);
   registerPass("untee",
                "removes local.tees, replacing them with sets and gets",
                createUnteePass);
diff --git a/src/passes/passes.h b/src/passes/passes.h
index 1cd2a5672..f84cfcaab 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -155,8 +155,10 @@ Pass* createTrapModeClamp();
 Pass* createTrapModeJS();
 Pass* createTupleOptimizationPass();
 Pass* createTypeRefiningPass();
+Pass* createTypeFinalizingPass();
 Pass* createTypeMergingPass();
 Pass* createTypeSSAPass();
+Pass* createTypeUnFinalizingPass();
 Pass* createUnteePass();
 Pass* createVacuumPass();