1 files changed, 81 insertions, 15 deletions

diff --git a/src/passes/DeadArgumentElimination.cpp b/src/passes/DeadArgumentElimination.cpp
index 033e766b6..d7b4d7db2 100644
--- a/src/passes/DeadArgumentElimination.cpp
+++ b/src/passes/DeadArgumentElimination.cpp
@@ -327,7 +327,7 @@ struct DAE : public Pass {
       // Refine argument types before doing anything else. This does not
       // affect whether an argument is used or not, it just refines the type
       // where possible.
-      refineArgumentTypes(func, calls, module);
+      refineArgumentTypes(func, calls, module, infoMap[name]);
       // Refine return types as well.
       if (refineReturnTypes(func, calls, module)) {
         refinedReturnTypes = true;
@@ -339,6 +339,7 @@ struct DAE : public Pass {
           assert(call->target == name);
           assert(call->operands.size() == numParams);
           auto* operand = call->operands[i];
+          // TODO: refnull etc.
           if (auto* c = operand->dynCast<Const>()) {
             if (value.type == Type::none) {
               // This is the first value seen.
@@ -544,7 +545,8 @@ private:
   // is not exported or called from the table or by reference.
   void refineArgumentTypes(Function* func,
                            const std::vector<Call*>& calls,
-                           Module* module) {
+                           Module* module,
+                           const DAEFunctionInfo& info) {
     if (!module->features.hasGC()) {
       return;
     }
@@ -553,7 +555,13 @@ private:
     newParamTypes.reserve(numParams);
     for (Index i = 0; i < numParams; i++) {
       auto originalType = func->getLocalType(i);
-      if (!originalType.isRef()) {
+      // If the parameter type is not a reference, there is nothing to refine.
+      // And if it is unused, also do nothing, as we can leave it to the other
+      // parts of this pass to optimize it properly, which avoids having to
+      // think about corner cases involving refining the type of an unused
+      // param (in particular, unused params are turned into locals, which means
+      // we'd need to think about defaultability etc.).
+      if (!originalType.isRef() || info.unusedParams.has(i)) {
         newParamTypes.push_back(originalType);
         continue;
       }
@@ -575,31 +583,84 @@ private:
 
     // Check if we are able to optimize here before we do the work to scan the
     // function body.
-    if (Type(newParamTypes) == func->getParams()) {
+    auto newParams = Type(newParamTypes);
+    if (newParams == func->getParams()) {
       return;
     }
 
-    // In terms of parameters, we can do this. However, we must also check
-    // local operations in the body, as if the parameter is reused and written
-    // to, then those types must be taken into account as well.
+    // We can do this!
+
+    // 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) &&
+      if (func->isParam(index) && !paramFixups.count(index) &&
           !Type::isSubType(set->value->type, newParamTypes[index])) {
-        // TODO: we could still optimize here, by creating a new local.
-        newParamTypes[index] = func->getLocalType(index);
+        paramFixups[index] = Builder::addVar(func, func->getLocalType(index));
       }
     }
 
-    auto newParams = Type(newParamTypes);
-    if (newParams == func->getParams()) {
-      return;
+    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;
+        }
+      }
     }
 
-    // We can do this! Update the types, including the types of gets and tees.
+    // Now that fixups are done, we can apply the new types.
     func->setParams(newParams);
-    for (auto* get : FindAll<LocalGet>(func->body).list) {
+
+    // 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);
@@ -615,6 +676,11 @@ private:
 
     // 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