[Wasm GC] DeadArgumentElimination: Do not refine return types of tail callees (#4025)

If there is a tail call, we can't change the return type of the function, as it
must match in the functions doing a tail call of it.

2 files changed, 26 insertions, 3 deletions

diff --git a/src/passes/DeadArgumentElimination.cpp b/src/passes/DeadArgumentElimination.cpp
index ddba3bee8..7ba1be9f8 100644
--- a/src/passes/DeadArgumentElimination.cpp
+++ b/src/passes/DeadArgumentElimination.cpp
@@ -327,8 +327,14 @@ struct DAE : public Pass {
       // affect whether an argument is used or not, it just refines the type
       // where possible.
       refineArgumentTypes(func, calls, module);
-      // Refine return types as well.
-      if (refineReturnTypes(func, calls, module)) {
+      // Refine return types as well, if we can. We cannot do so if this
+      // function is tail-called, because then the return type must match that
+      // of the function doing a tail call of it - we cannot change just one of
+      // them.
+      //
+      // TODO: Try to optimize in a more holistic manner, see the TODO in
+      //       refineReturnTypes() about missing a global optimum.
+      if (!tailCallees.count(name) && refineReturnTypes(func, calls, module)) {
         refinedReturnTypes = true;
       }
       // Check if all calls pass the same constant for a particular argument.
@@ -434,7 +440,7 @@ struct DAE : public Pass {
         if (infoMap[name].hasTailCalls) {
           continue;
         }
-        if (tailCallees.find(name) != tailCallees.end()) {
+        if (tailCallees.count(name)) {
           continue;
         }
         auto iter = allCalls.find(name);
diff --git a/test/lit/passes/dae-gc.wast b/test/lit/passes/dae-gc.wast
index f80673cf2..230f757fc 100644
--- a/test/lit/passes/dae-gc.wast
+++ b/test/lit/passes/dae-gc.wast
@@ -541,4 +541,21 @@
    (i32.const 1)
   )
  )
+
+ ;; This function does a return call of the one after it. The one after it
+ ;; returns a ref.func of this one. They both begin by returning a funcref; if
+ ;; we refine the return type of the latter (which ref.func allows us to do)
+ ;; then the return type would not match, and we would get a validation error.
+ ;; CHECK:      (func $do-return-call (result funcref)
+ ;; CHECK-NEXT:  (return_call $return-ref-func)
+ ;; CHECK-NEXT: )
+ (func $do-return-call (result funcref)
+  (return_call $return-ref-func)
+ )
+ ;; CHECK:      (func $return-ref-func (result funcref)
+ ;; CHECK-NEXT:  (ref.func $do-return-call)
+ ;; CHECK-NEXT: )
+ (func $return-ref-func (result funcref)
+  (ref.func $do-return-call)
+ )
 )