[EH] Fix pop enclosed within a block in DCE (#6922)

#6400 fixed this case but that handled only when a `pop` is an immediate child of the current expression, but a `pop` can be nested deeper down. We conservatively run the EH fixup whenever we have a `pop` and create `block`s in the optimization. We considered using `FindAll<Pop>` to make it precise, but we decided the quadratic time plexity was not worth it. Fixes #6918.
WebAssembly · Sep 11, 2024 · 432ed1c · 432ed1c
1 parent 1a2d26f
commit 432ed1c
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 15 deletions.
diff --git a/src/passes/DeadCodeElimination.cpp b/src/passes/DeadCodeElimination.cpp
@@ -56,6 +56,10 @@ struct DeadCodeElimination
   // as we remove code, we must keep the types of other nodes valid
   TypeUpdater typeUpdater;
 
+  // Information used to decide whether we need EH fixups at the end
+  bool hasPop = false;     // Do we have a 'pop' in this function?
+  bool addedBlock = false; // Have we added blocks in this function?
+
   Expression* replaceCurrent(Expression* expression) {
     auto* old = getCurrent();
     if (old == expression) {
@@ -73,6 +77,10 @@ struct DeadCodeElimination
   }
 
   void visitExpression(Expression* curr) {
+    if (curr->is<Pop>()) {
+      hasPop = true;
+    }
+
     if (!Properties::isControlFlowStructure(curr)) {
       // Control flow structures require special handling, but others are
       // simple.
@@ -90,11 +98,7 @@ 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;
@@ -109,12 +113,8 @@ struct DeadCodeElimination
           if (remainingChildren.size() == 1) {
             replaceCurrent(remainingChildren[0]);
           } else {
+            addedBlock = true;
             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;
-            }
           }
         }
       }
@@ -196,10 +196,11 @@ struct DeadCodeElimination
     }
   }
 
-  bool needEHFixups = false;
-
   void visitFunction(Function* curr) {
-    if (needEHFixups) {
+    // We conservatively run the EH pop fixup if this function has a 'pop' and
+    // if we have ever added blocks in the optimization, which may have moved
+    // pops into the blocks.
+    if (hasPop && addedBlock) {
       EHUtils::handleBlockNestedPops(curr, *getModule());
     }
   }

diff --git a/test/lit/passes/dce-eh-legacy.wast b/test/lit/passes/dce-eh-legacy.wast
@@ -1,14 +1,28 @@
-;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
 ;; RUN: wasm-opt %s --dce -all -S -o - | filecheck %s
 
 ;; If either try body or catch body is reachable, the whole try construct is
 ;; reachable
 (module
+  ;; CHECK:      (type $0 (func))
+
+  ;; CHECK:      (type $1 (func (param i32)))
+
+  ;; CHECK:      (type $2 (func (param eqref)))
+
+  ;; CHECK:      (type $3 (func (param i32 i32)))
+
+  ;; CHECK:      (type $4 (func (result i32)))
+
+  ;; CHECK:      (type $struct (struct (field (mut eqref))))
+  (type $struct (struct (field (mut (ref null eq)))))
+
   ;; CHECK:      (tag $e)
   (tag $e)
-
   ;; CHECK:      (tag $e-i32 (param i32))
   (tag $e-i32 (param i32))
+  ;; CHECK:      (tag $e-eqref (param eqref))
+  (tag $e-eqref (param (ref null eq)))
 
   ;; CHECK:      (func $foo (type $0)
   ;; CHECK-NEXT:  (nop)
@@ -149,11 +163,63 @@
     )
   )
 
-  ;; CHECK:      (func $helper-i32-i32 (type $2) (param $0 i32) (param $1 i32)
+  ;; CHECK:      (func $helper-i32-i32 (type $3) (param $0 i32) (param $1 i32)
   ;; CHECK-NEXT:  (nop)
   ;; CHECK-NEXT: )
   (func $helper-i32-i32 (param $0 i32) (param $1 i32)
    (nop)
   )
+
+  ;; CHECK:      (func $pop-within-block (type $4) (result i32)
+  ;; CHECK-NEXT:  (local $0 eqref)
+  ;; CHECK-NEXT:  (try (result i32)
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch $e-eqref
+  ;; CHECK-NEXT:    (local.set $0
+  ;; CHECK-NEXT:     (pop eqref)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (block
+  ;; CHECK-NEXT:     (block
+  ;; CHECK-NEXT:      (drop
+  ;; CHECK-NEXT:       (struct.new $struct
+  ;; CHECK-NEXT:        (local.get $0)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (unreachable)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $pop-within-block (result i32)
+    (try (result i32)
+      (do
+        (i32.const 0)
+      )
+      (catch $e-eqref
+        (drop
+          ;; Optimization moves the 'pop' inside a block, which needs to be
+          ;; extracted out of the block at the end.
+          ;; (block
+          ;;   (drop
+          ;;     (struct.new $struct.0
+          ;;       (pop eqref)
+          ;;     )
+          ;;   )
+          ;;   (unreachable)
+          ;; )
+          (ref.eq
+            (struct.new $struct
+              (pop (ref null eq))
+            )
+            (unreachable)
+          )
+        )
+        (i32.const 0)
+      )
+    )
+  )
 )