[Wasm GC] Refine return types of tail calling functions in DeadArgumentElimination (#4036)

To do this we need to look at tail calls and not just returns.

1 files changed, 18 insertions, 2 deletions

diff --git a/src/passes/DeadArgumentElimination.cpp b/src/passes/DeadArgumentElimination.cpp
index f3e3ae98c..601acfc5c 100644
--- a/src/passes/DeadArgumentElimination.cpp
+++ b/src/passes/DeadArgumentElimination.cpp
@@ -660,9 +660,25 @@ private:
     }
 
     // Scan the body and look at the returns.
+    auto processReturnType = [&](Type type) {
+      refinedType = Type::getLeastUpperBound(refinedType, type);
+      // Return whether we still look ok to do the optimization. If this is
+      // false then we can stop here.
+      return refinedType != originalType;
+    };
     for (auto* ret : FindAll<Return>(func->body).list) {
-      refinedType = Type::getLeastUpperBound(refinedType, ret->value->type);
-      if (refinedType == originalType) {
+      if (!processReturnType(ret->value->type)) {
+        return false;
+      }
+    }
+    for (auto* call : FindAll<Call>(func->body).list) {
+      if (call->isReturn &&
+          !processReturnType(module->getFunction(call->target)->getResults())) {
+        return false;
+      }
+    }
+    for (auto* call : FindAll<CallIndirect>(func->body).list) {
+      if (call->isReturn && !processReturnType(call->sig.results)) {
         return false;
       }
     }