Fix DeadArgumentElimination return value opts on nesting+recursion (#5701)

The testcase here has a recursive call that is also nested in itself, something
like

(return
  (call $me
    (return
      ..

This found a bug in our return value removal logic. When we remove a return
value we both modify call sites (to add drops) and modify returns
(to remove their values). One of those uses pointers into the IR which the other
invalidated, so the order of the two matters. This PR just reorders that code to
fix the bug.

2 files changed, 50 insertions, 13 deletions

diff --git a/src/passes/DeadArgumentElimination.cpp b/src/passes/DeadArgumentElimination.cpp
index f739e9263..ab4c8b8eb 100644
--- a/src/passes/DeadArgumentElimination.cpp
+++ b/src/passes/DeadArgumentElimination.cpp
@@ -317,7 +317,21 @@ private:
   void
   removeReturnValue(Function* func, std::vector<Call*>& calls, Module* module) {
     func->setResults(Type::none);
-    Builder builder(*module);
+    // Remove the drops on the calls. Note that we must do this before updating
+    // returns in ReturnUpdater, as there may be recursive calls of this
+    // function to itself. So we first use the information in allDroppedCalls
+    // before the ReturnUpdater potentially invalidates that information as it
+    // modifies the function.
+    for (auto* call : calls) {
+      auto iter = allDroppedCalls.find(call);
+      assert(iter != allDroppedCalls.end());
+      Expression** location = iter->second;
+      *location = call;
+      // Update the call's type.
+      if (call->type != Type::unreachable) {
+        call->type = Type::none;
+      }
+    }
     // Remove any return values.
     struct ReturnUpdater : public PostWalker<ReturnUpdater> {
       Module* module;
@@ -334,18 +348,7 @@ private:
     } returnUpdater(func, module);
     // Remove any value flowing out.
     if (func->body->type.isConcrete()) {
-      func->body = builder.makeDrop(func->body);
-    }
-    // Remove the drops on the calls.
-    for (auto* call : calls) {
-      auto iter = allDroppedCalls.find(call);
-      assert(iter != allDroppedCalls.end());
-      Expression** location = iter->second;
-      *location = call;
-      // Update the call's type.
-      if (call->type != Type::unreachable) {
-        call->type = Type::none;
-      }
+      func->body = Builder(*module).makeDrop(func->body);
     }
   }
 
diff --git a/test/lit/passes/dae_all-features.wast b/test/lit/passes/dae_all-features.wast
index 8bfd3c9d0..3952747f7 100644
--- a/test/lit/passes/dae_all-features.wast
+++ b/test/lit/passes/dae_all-features.wast
@@ -658,3 +658,37 @@
   )
  )
 )
+
+(module
+ ;; CHECK:      (type $i32_=>_none (func (param i32)))
+
+ ;; CHECK:      (func $0 (type $i32_=>_none) (param $0 i32)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (call $0
+ ;; CHECK-NEXT:    (block
+ ;; CHECK-NEXT:     (drop
+ ;; CHECK-NEXT:      (i32.const 1)
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:     (return)
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (return)
+ ;; CHECK-NEXT: )
+ (func $0 (param $0 i32) (result i32)
+  ;; The result of this function can be removed, which makes us modify the
+  ;; returns (we should not return a value any more) and also the calls (the
+  ;; calls must be dropped). The returns here are nested in each other, and one
+  ;; is a recursive call to this function itself, which makes this a corner case
+  ;; we might emit invalid code for.
+  (return
+   (drop
+    (call $0
+     (return
+      (i32.const 1)
+     )
+    )
+   )
+  )
+ )
+)