DeadArgumentElimination argument subtyping: Add fixups if the param is used (#4319)

Before, if we saw a param is written, that prevented us from subtyping it:

function foo(x : oldType) {
  ..
  x = someValue;
  ..
}

Even if all calls to foo send some specific struct type that we'd like to subtype
to, seeing that write stopped us. To handle such a write we need to do some
extra handling for the case in which it is written a less-specific type (that is,
if someValue is of type oldType, something like this:

function foo(x : newType) {
  var x_old : oldType;
  x_old = x; // copy the param to x_old, and use x_old everywhere
  ..
  x_old = someValue;
  ..
}

That is, still refine the param type, but inside the function use a new local that
has the old type, and is guaranteed to validate. This PR implements that logic
so that we can optimize more cases.

To allow that, this PR avoids trying to both refine a type and remove a param as
being unused - that has annoying corner cases. If it is unused, we can simply
remove it anyhow.

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