[NFC] Refactor GlobalTypeRewriter to split out the type mapping logic (#5295)

The logic that is split out into mapTypes gets a map of old type => new type and then
updates the module to replace the old with the new.

This will be useful in a future pass to merge types. We will tell it to map the types to
be merged with the types we want to merge them into.

This removes an assert on all types being in the map to allow some of them not to be.
(the new pass that will use this will only want to map some types).

2 files changed, 30 insertions, 13 deletions

diff --git a/src/ir/type-updating.cpp b/src/ir/type-updating.cpp
index 2d75f0951..08c04e8d4 100644
--- a/src/ir/type-updating.cpp
+++ b/src/ir/type-updating.cpp
@@ -97,21 +97,33 @@ void GlobalTypeRewriter::update() {
   // Map the old types to the new ones. This uses the fact that type indices
   // are the same in the old and new types, that is, we have not added or
   // removed types, just modified them.
-  using OldToNewTypes = std::unordered_map<HeapType, HeapType>;
-  OldToNewTypes oldToNewTypes;
+  TypeMap oldToNewTypes;
   for (Index i = 0; i < indexedTypes.types.size(); i++) {
     oldToNewTypes[indexedTypes.types[i]] = newTypes[i];
   }
 
+  // Update type names (doing it before mapTypes can help debugging there, but
+  // has no other effect; mapTypes does not look at type names).
+  for (auto& [old, new_] : oldToNewTypes) {
+    if (wasm.typeNames.count(old)) {
+      wasm.typeNames[new_] = wasm.typeNames[old];
+    }
+  }
+
+  mapTypes(oldToNewTypes);
+}
+
+void GlobalTypeRewriter::mapTypes(const TypeMap& oldToNewTypes) {
+
   // Replace all the old types in the module with the new ones.
   struct CodeUpdater
     : public WalkerPass<
         PostWalker<CodeUpdater, UnifiedExpressionVisitor<CodeUpdater>>> {
     bool isFunctionParallel() override { return true; }
 
-    OldToNewTypes& oldToNewTypes;
+    const TypeMap& oldToNewTypes;
 
-    CodeUpdater(OldToNewTypes& oldToNewTypes) : oldToNewTypes(oldToNewTypes) {}
+    CodeUpdater(const TypeMap& oldToNewTypes) : oldToNewTypes(oldToNewTypes) {}
 
     std::unique_ptr<Pass> create() override {
       return std::make_unique<CodeUpdater>(oldToNewTypes);
@@ -136,8 +148,10 @@ void GlobalTypeRewriter::update() {
         return type;
       }
       if (type.isFunction() || type.isData()) {
-        assert(oldToNewTypes.count(type));
-        return oldToNewTypes[type];
+        auto iter = oldToNewTypes.find(type);
+        if (iter != oldToNewTypes.end()) {
+          return iter->second;
+        }
       }
       return type;
     }
@@ -228,13 +242,6 @@ void GlobalTypeRewriter::update() {
   for (auto& tag : wasm.tags) {
     tag->sig = updater.getNew(tag->sig);
   }
-
-  // Update type names.
-  for (auto& [old, new_] : oldToNewTypes) {
-    if (wasm.typeNames.count(old)) {
-      wasm.typeNames[new_] = wasm.typeNames[old];
-    }
-  }
 }
 
 Type GlobalTypeRewriter::getTempType(Type type) {
diff --git a/src/ir/type-updating.h b/src/ir/type-updating.h
index 50ba9efca..8da84ceb6 100644
--- a/src/ir/type-updating.h
+++ b/src/ir/type-updating.h
@@ -337,6 +337,16 @@ public:
   // the module.
   void update();
 
+  using TypeMap = std::unordered_map<HeapType, HeapType>;
+
+  // Given a map of old type => new type to use instead, this rewrites all type
+  // uses in the module to apply that map. This is used internally in update()
+  // but may be useful by itself as well.
+  //
+  // The input map does not need to contain all the types. Whenever a type does
+  // not appear, it is mapped to itself.
+  void mapTypes(const TypeMap& oldToNewTypes);
+
   // Subclasses can implement these methods to modify the new set of types that
   // we map to. By default, we simply copy over the types, and these functions
   // are the hooks to apply changes through. The methods receive as input the