From d08be6cbf7eec520e97e45d6f459439bd321003c Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Thu, 14 Mar 2024 09:53:45 -0700
Subject: [PATCH] DCE: Fix old EH on a pop that gets moved in a catch body

---
 src/passes/DeadCodeElimination.cpp | 32 +++++++++++++----
 test/lit/passes/dce-eh-old.wast    | 56 ++++++++++++++++++++++++++++--
 2 files changed, 79 insertions(+), 9 deletions(-)

diff --git a/src/passes/DeadCodeElimination.cpp b/src/passes/DeadCodeElimination.cpp
index 0b382f46574..996133ad589 100644
--- a/src/passes/DeadCodeElimination.cpp
+++ b/src/passes/DeadCodeElimination.cpp
@@ -28,13 +28,14 @@
 // have no side effects.
 //
 
-#include <ir/iteration.h>
-#include <ir/properties.h>
-#include <ir/type-updating.h>
-#include <pass.h>
-#include <vector>
-#include <wasm-builder.h>
-#include <wasm.h>
+#include "ir/eh-utils.h"
+#include "ir/iteration.h"
+#include "ir/properties.h"
+#include "ir/type-updating.h"
+#include "pass.h"
+#include "vector"
+#include "wasm-builder.h"
+#include "wasm.h"
 
 namespace wasm {
 
@@ -89,7 +90,11 @@ struct DeadCodeElimination
           Builder builder(*getModule());
           std::vector<Expression*> remainingChildren;
           bool afterUnreachable = false;
+          bool hasPop = false;
           for (auto* child : ChildIterator(curr)) {
+            if (child->is<Pop>()) {
+              hasPop = true;
+            }
             if (afterUnreachable) {
               typeUpdater.noteRecursiveRemoval(child);
               continue;
@@ -105,6 +110,11 @@ struct DeadCodeElimination
             replaceCurrent(remainingChildren[0]);
           } else {
             replaceCurrent(builder.makeBlock(remainingChildren));
+            if (hasPop) {
+              // We are moving a pop into a new block we just created, which
+              // means we may need to fix things up here.
+              needEHFixups = true;
+            }
           }
         }
       }
@@ -179,6 +189,14 @@ struct DeadCodeElimination
       WASM_UNREACHABLE("unimplemented DCE control flow structure");
     }
   }
+
+  bool needEHFixups = false;
+
+  void visitFunction(Function* curr) {
+    if (needEHFixups) {
+      EHUtils::handleBlockNestedPops(curr, *getModule());
+    }
+  }
 };
 
 Pass* createDeadCodeEliminationPass() { return new DeadCodeElimination(); }
diff --git a/test/lit/passes/dce-eh-old.wast b/test/lit/passes/dce-eh-old.wast
index 123f7a42d72..b6e44bd6298 100644
--- a/test/lit/passes/dce-eh-old.wast
+++ b/test/lit/passes/dce-eh-old.wast
@@ -5,13 +5,15 @@
 ;; reachable
 (module
   ;; CHECK:      (tag $e)
+  (tag $e)
+
+  ;; CHECK:      (tag $e-i32 (param i32))
+  (tag $e-i32 (param i32))
 
   ;; CHECK:      (func $foo (type $0)
   ;; CHECK-NEXT:  (nop)
   ;; CHECK-NEXT: )
   (func $foo)
-  (tag $e)
-
 
   ;; CHECK:      (func $try_unreachable (type $0)
   ;; CHECK-NEXT:  (try $try
@@ -130,4 +132,54 @@
       )
     )
   )
+
+  ;; CHECK:      (func $call-pop-catch (type $0)
+  ;; CHECK-NEXT:  (local $0 i32)
+  ;; CHECK-NEXT:  (block $label
+  ;; CHECK-NEXT:   (try $try
+  ;; CHECK-NEXT:    (do
+  ;; CHECK-NEXT:     (nop)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (catch $e-i32
+  ;; CHECK-NEXT:     (local.set $0
+  ;; CHECK-NEXT:      (pop i32)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (block
+  ;; CHECK-NEXT:      (block
+  ;; CHECK-NEXT:       (drop
+  ;; CHECK-NEXT:        (local.get $0)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (br $label)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $call-pop-catch
+    (block $label
+     (try
+       (do
+         (nop)
+       )
+       (catch $e-i32
+         ;; This call is unreachable and can be removed. The pop, however, must
+         ;; be carefully handled while we do so, to not break validation.
+         (call $helper-i32-i32
+           (pop i32)
+           (br $label)
+         )
+         (nop)
+       )
+      )
+    )
+  )
+
+  ;; CHECK:      (func $helper-i32-i32 (type $2) (param $0 i32) (param $1 i32)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $helper-i32-i32 (param $0 i32) (param $1 i32)
+   (nop)
+  )
 )
+