2 files changed, 31 insertions, 11 deletions

diff --git a/src/ir/type-updating.cpp b/src/ir/type-updating.cpp
index 12a8c7c36..760d71b9c 100644
--- a/src/ir/type-updating.cpp
+++ b/src/ir/type-updating.cpp
@@ -28,7 +28,10 @@ namespace wasm {
 
 GlobalTypeRewriter::GlobalTypeRewriter(Module& wasm) : wasm(wasm) {}
 
-void GlobalTypeRewriter::update() { mapTypes(rebuildTypes()); }
+void GlobalTypeRewriter::update(
+  const std::vector<HeapType>& additionalPrivateTypes) {
+  mapTypes(rebuildTypes(additionalPrivateTypes));
+}
 
 GlobalTypeRewriter::TypeMap GlobalTypeRewriter::rebuildTypes(
   const std::vector<HeapType>& additionalPrivateTypes) {
@@ -40,8 +43,17 @@ GlobalTypeRewriter::TypeMap GlobalTypeRewriter::rebuildTypes(
   Index i = 0;
   auto privateTypes = ModuleUtils::getPrivateHeapTypes(wasm);
 
-  for (auto t : additionalPrivateTypes) {
-    privateTypes.push_back(t);
+  if (!additionalPrivateTypes.empty()) {
+    // Only add additional private types that are not already in the list.
+    std::unordered_set<HeapType> privateTypesSet(privateTypes.begin(),
+                                                 privateTypes.end());
+
+    for (auto t : additionalPrivateTypes) {
+      if (!privateTypesSet.count(t)) {
+        privateTypes.push_back(t);
+        privateTypesSet.insert(t);
+      }
+    }
   }
 
   // Topological sort to have supertypes first, but we have to account for the
diff --git a/src/ir/type-updating.h b/src/ir/type-updating.h
index 1626b40bd..60b92e585 100644
--- a/src/ir/type-updating.h
+++ b/src/ir/type-updating.h
@@ -352,7 +352,12 @@ public:
   // Main entry point. This performs the entire process of creating new heap
   // types and calling the hooks below, then applies the new types throughout
   // the module.
-  void update();
+  //
+  // This only operates on private types (so as not to modify the module's
+  // external ABI). It takes as a parameter a list of public types to consider
+  // private, which allows more flexibility (e.g. in closed world if a pass
+  // knows a type is safe to modify despite being public, it can add it).
+  void update(const std::vector<HeapType>& additionalPrivateTypes = {});
 
   using TypeMap = std::unordered_map<HeapType, HeapType>;
 
@@ -398,7 +403,10 @@ public:
 
   // Helper for the repeating pattern of just updating Signature types using a
   // map of old heap type => new Signature.
-  static void updateSignatures(const SignatureUpdates& updates, Module& wasm) {
+  static void
+  updateSignatures(const SignatureUpdates& updates,
+                   Module& wasm,
+                   const std::vector<HeapType>& additionalPrivateTypes = {}) {
     if (updates.empty()) {
       return;
     }
@@ -407,9 +415,11 @@ public:
       const SignatureUpdates& updates;
 
     public:
-      SignatureRewriter(Module& wasm, const SignatureUpdates& updates)
+      SignatureRewriter(Module& wasm,
+                        const SignatureUpdates& updates,
+                        const std::vector<HeapType>& additionalPrivateTypes)
         : GlobalTypeRewriter(wasm), updates(updates) {
-        update();
+        update(additionalPrivateTypes);
       }
 
       void modifySignature(HeapType oldSignatureType, Signature& sig) override {
@@ -419,7 +429,7 @@ public:
           sig.results = getTempType(iter->second.results);
         }
       }
-    } rewriter(wasm, updates);
+    } rewriter(wasm, updates, additionalPrivateTypes);
   }
 
 protected:
@@ -427,9 +437,7 @@ protected:
   // returns a map from the old types to the modified types. Used internally in
   // update().
   //
-  // This only operates on private types (so as not to modify the module's
-  // external ABI). It takes as a parameter a list of public types to consider
-  // private, which allows more flexibility.
+  // See above regarding private types.
   TypeMap
   rebuildTypes(const std::vector<HeapType>& additionalPrivateTypes = {});