3 files changed, 108 insertions, 90 deletions

diff --git a/src/ir/type-updating.cpp b/src/ir/type-updating.cpp
index 316dbd93f..0d579367f 100644
--- a/src/ir/type-updating.cpp
+++ b/src/ir/type-updating.cpp
@@ -17,6 +17,7 @@
 #include "type-updating.h"
 #include "find_all.h"
 #include "ir/module-utils.h"
+#include "ir/utils.h"
 #include "wasm-type.h"
 #include "wasm.h"
 
@@ -316,6 +317,99 @@ Expression* fixLocalGet(LocalGet* get, Module& wasm) {
   return get;
 }
 
+void updateParamTypes(Function* func,
+                      const std::vector<Type>& newParamTypes,
+                      Module& wasm) {
+  // Before making this update, we must be careful if the param was "reused",
+  // specifically, if it is assigned a less-specific type in the body then
+  // we'd get a validation error when we refine it. To handle that, if a less-
+  // specific type is assigned simply switch to a new local, that is, we can
+  // do a fixup like this:
+  //
+  // function foo(x : oldType) {
+  //   ..
+  //   x = (oldType)val;
+  //
+  // =>
+  //
+  // function foo(x : newType) {
+  //   var x_oldType = x; // assign the param immediately to a fixup var
+  //   ..
+  //   x_oldType = (oldType)val; // fixup var is used throughout the body
+  //
+  // Later optimization passes may be able to remove the extra var, and can
+  // take advantage of the refined argument type while doing so.
+
+  // A map of params that need a fixup to the new fixup var used for it.
+  std::unordered_map<Index, Index> paramFixups;
+
+  FindAll<LocalSet> sets(func->body);
+
+  for (auto* set : sets.list) {
+    auto index = set->index;
+    if (func->isParam(index) && !paramFixups.count(index) &&
+        !Type::isSubType(set->value->type, newParamTypes[index])) {
+      paramFixups[index] = Builder::addVar(func, func->getLocalType(index));
+    }
+  }
+
+  FindAll<LocalGet> gets(func->body);
+
+  // Apply the fixups we identified that we need.
+  if (!paramFixups.empty()) {
+    // Write the params immediately to the fixups.
+    Builder builder(wasm);
+    std::vector<Expression*> contents;
+    for (Index index = 0; index < func->getNumParams(); index++) {
+      auto iter = paramFixups.find(index);
+      if (iter != paramFixups.end()) {
+        auto fixup = iter->second;
+        contents.push_back(builder.makeLocalSet(
+          fixup, builder.makeLocalGet(index, newParamTypes[index])));
+      }
+    }
+    contents.push_back(func->body);
+    func->body = builder.makeBlock(contents);
+
+    // Update gets and sets using the param to use the fixup.
+    for (auto* get : gets.list) {
+      auto iter = paramFixups.find(get->index);
+      if (iter != paramFixups.end()) {
+        get->index = iter->second;
+      }
+    }
+    for (auto* set : sets.list) {
+      auto iter = paramFixups.find(set->index);
+      if (iter != paramFixups.end()) {
+        set->index = iter->second;
+      }
+    }
+  }
+
+  // Update local.get/local.tee operations that use the modified param type.
+  for (auto* get : gets.list) {
+    auto index = get->index;
+    if (func->isParam(index)) {
+      get->type = newParamTypes[index];
+    }
+  }
+  for (auto* set : sets.list) {
+    auto index = set->index;
+    if (func->isParam(index) && set->isTee()) {
+      set->type = newParamTypes[index];
+      set->finalize();
+    }
+  }
+
+  // Propagate the new get and set types outwards.
+  ReFinalize().walkFunctionInModule(func, &wasm);
+
+  if (!paramFixups.empty()) {
+    // We have added locals, and must handle non-nullability of them.
+    TypeUpdating::handleNonDefaultableLocals(func, wasm);
+  }
+}
+
 } // namespace TypeUpdating
 
 } // namespace wasm
diff --git a/src/ir/type-updating.h b/src/ir/type-updating.h
index d3be7c3ca..8794cd0da 100644
--- a/src/ir/type-updating.h
+++ b/src/ir/type-updating.h
@@ -364,6 +364,18 @@ Type getValidLocalType(Type type, FeatureSet features);
 // ref.as_non_null around it, if the local should be non-nullable but is not).
 Expression* fixLocalGet(LocalGet* get, Module& wasm);
 
+// Applies new types of parameters to a function. This does all the necessary
+// changes aside from altering the function type, which the caller is expected
+// to do (the caller might simply change the type, but in other cases the caller
+// might be rewriting the types and need to preserve their identity in terms of
+// nominal typing, so we don't change the type here). The specific things this
+// function does are to update the types of local.get/tee operations,
+// refinalize, etc., basically all operations necessary to ensure validation
+// with the new types.
+void updateParamTypes(Function* func,
+                      const std::vector<Type>& newParamTypes,
+                      Module& wasm);
+
 } // namespace TypeUpdating
 
 } // namespace wasm
diff --git a/src/passes/DeadArgumentElimination.cpp b/src/passes/DeadArgumentElimination.cpp
index d7b4d7db2..e997c591f 100644
--- a/src/passes/DeadArgumentElimination.cpp
+++ b/src/passes/DeadArgumentElimination.cpp
@@ -589,98 +589,10 @@ private:
     }
 
     // We can do this!
+    TypeUpdating::updateParamTypes(func, newParamTypes, *module);
 
-    // Before making this update, we must be careful if the param was "reused",
-    // specifically, if it is assigned a less-specific type in the body then
-    // we'd get a validation error when we refine it. To handle that, if a less-
-    // specific type is assigned simply switch to a new local, that is, we can
-    // do a fixup like this:
-    //
-    // function foo(x : oldType) {
-    //   ..
-    //   x = (oldType)val;
-    //
-    // =>
-    //
-    // function foo(x : newType) {
-    //   var x_oldType = x; // assign the param immediately to a fixup var
-    //   ..
-    //   x_oldType = (oldType)val; // fixup var is used throughout the body
-    //
-    // Later optimization passes may be able to remove the extra var, and can
-    // take advantage of the refined argument type while doing so.
-
-    // A map of params that need a fixup to the new fixup var used for it.
-    std::unordered_map<Index, Index> paramFixups;
-
-    FindAll<LocalSet> sets(func->body);
-
-    for (auto* set : sets.list) {
-      auto index = set->index;
-      if (func->isParam(index) && !paramFixups.count(index) &&
-          !Type::isSubType(set->value->type, newParamTypes[index])) {
-        paramFixups[index] = Builder::addVar(func, func->getLocalType(index));
-      }
-    }
-
-    FindAll<LocalGet> gets(func->body);
-
-    // Apply the fixups we identified that we need.
-    if (!paramFixups.empty()) {
-      // Write the params immediately to the fixups.
-      Builder builder(*module);
-      std::vector<Expression*> contents;
-      for (Index index = 0; index < func->getNumParams(); index++) {
-        auto iter = paramFixups.find(index);
-        if (iter != paramFixups.end()) {
-          auto fixup = iter->second;
-          contents.push_back(builder.makeLocalSet(
-            fixup, builder.makeLocalGet(index, newParamTypes[index])));
-        }
-      }
-      contents.push_back(func->body);
-      func->body = builder.makeBlock(contents);
-
-      // Update gets and sets using the param to use the fixup.
-      for (auto* get : gets.list) {
-        auto iter = paramFixups.find(get->index);
-        if (iter != paramFixups.end()) {
-          get->index = iter->second;
-        }
-      }
-      for (auto* set : sets.list) {
-        auto iter = paramFixups.find(set->index);
-        if (iter != paramFixups.end()) {
-          set->index = iter->second;
-        }
-      }
-    }
-
-    // Now that fixups are done, we can apply the new types.
+    // Also update the function's type.
     func->setParams(newParams);
-
-    // Update local.get/local.tee operations that use the modified param type.
-    for (auto* get : gets.list) {
-      auto index = get->index;
-      if (func->isParam(index)) {
-        get->type = func->getLocalType(index);
-      }
-    }
-    for (auto* set : sets.list) {
-      auto index = set->index;
-      if (func->isParam(index) && set->isTee()) {
-        set->type = func->getLocalType(index);
-        set->finalize();
-      }
-    }
-
-    // Propagate the new get and set types outwards.
-    ReFinalize().walkFunctionInModule(func, module);
-
-    if (!paramFixups.empty()) {
-      // We have added locals, and must handle non-nullability of them.
-      TypeUpdating::handleNonDefaultableLocals(func, *module);
-    }
   }
 
   // See if the types returned from a function allow us to define a more refined