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