StackIR: Optimize away a drop before an unreachable (#6719)

Anything else right before an unreachable is removed by the main DCE
pass anyhow, but because of the structured form of BinaryenIR we can't remove
a drop. That is, this is the difference between

(i32.eqz
  (i32.const 42)
  (unreachable)
)

and

(drop
  (call $foo)
)
(unreachable)

In both cases the unreachable is preceded by something we don't need,
but in the latter case it must remain in BinaryenIR for validation.

To optimize this, add a rule in StackIR.

Fixes #6715

2 files changed, 156 insertions, 1 deletions

diff --git a/src/wasm/wasm-stack-opts.cpp b/src/wasm/wasm-stack-opts.cpp
index ae9fee394..5e18cf26f 100644
--- a/src/wasm/wasm-stack-opts.cpp
+++ b/src/wasm/wasm-stack-opts.cpp
@@ -44,8 +44,9 @@ void StackIROptimizer::run() {
   vacuum();
 }
 
-// Remove unreachable code.
 void StackIROptimizer::dce() {
+  // Remove code after an unreachable instruction: anything after it, up to the
+  // next control flow barrier, can simply be removed.
   bool inUnreachableCode = false;
   for (Index i = 0; i < insts.size(); i++) {
     auto* inst = insts[i];
@@ -64,6 +65,42 @@ void StackIROptimizer::dce() {
       inUnreachableCode = true;
     }
   }
+
+  // Remove code before an Unreachable. Consider this:
+  //
+  //  (drop
+  //   ..
+  //  )
+  //  (unreachable)
+  //
+  // The drop is not needed, as the unreachable puts the stack in the
+  // polymorphic state anyhow. Note that we don't need to optimize anything
+  // other than a drop here, as in general the Binaryen IR DCE pass will handle
+  // everything else. A drop followed by an unreachable is the only thing that
+  // pass cannot handle, as the structured form of Binaryen IR does not allow
+  // removing such a drop, and so we can only do it here in StackIR.
+  //
+  // TODO: We can look even further back, say if there is another drop of
+  //       something before, then we can remove that drop as well. To do that
+  //       we'd need to inspect the stack going backwards.
+  for (Index i = 1; i < insts.size(); i++) {
+    auto* inst = insts[i];
+    if (!inst || inst->op != StackInst::Basic ||
+        !inst->origin->is<Unreachable>()) {
+      continue;
+    }
+
+    // Look back past nulls.
+    Index j = i - 1;
+    while (j > 0 && !insts[j]) {
+      j--;
+    }
+
+    auto*& prev = insts[j];
+    if (prev && prev->op == StackInst::Basic && prev->origin->is<Drop>()) {
+      prev = nullptr;
+    }
+  }
 }
 
 // Remove obviously-unneeded code.
diff --git a/test/lit/passes/stack-ir-dce.wast b/test/lit/passes/stack-ir-dce.wast
new file mode 100644
index 000000000..24afb4840
--- /dev/null
+++ b/test/lit/passes/stack-ir-dce.wast
@@ -0,0 +1,118 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+
+;; RUN: wasm-opt %s --generate-stack-ir --optimize-stack-ir -all --print-stack-ir | filecheck %s
+;; Also verify we roundtrip the output here properly.
+;; RUN: wasm-opt %s --generate-stack-ir --optimize-stack-ir -all --roundtrip --print | filecheck %s --check-prefix=ROUNDTRIP
+
+(module
+  ;; CHECK:      (func $drop-unreachable (type $0) (result i32)
+  ;; CHECK-NEXT:  call $drop-unreachable
+  ;; CHECK-NEXT:  unreachable
+  ;; CHECK-NEXT: )
+  ;; ROUNDTRIP:      (func $drop-unreachable (type $0) (result i32)
+  ;; ROUNDTRIP-NEXT:  (drop
+  ;; ROUNDTRIP-NEXT:   (call $drop-unreachable)
+  ;; ROUNDTRIP-NEXT:  )
+  ;; ROUNDTRIP-NEXT:  (unreachable)
+  ;; ROUNDTRIP-NEXT: )
+  (func $drop-unreachable (result i32)
+    ;; This drop can be removed.
+    (drop
+      (call $drop-unreachable)
+    )
+    (unreachable)
+  )
+
+  ;; CHECK:      (func $unreachable (type $0) (result i32)
+  ;; CHECK-NEXT:  unreachable
+  ;; CHECK-NEXT: )
+  ;; ROUNDTRIP:      (func $unreachable (type $0) (result i32)
+  ;; ROUNDTRIP-NEXT:  (unreachable)
+  ;; ROUNDTRIP-NEXT: )
+  (func $unreachable (result i32)
+    ;; An unreachable with nothing before it. Check we do not error here.
+    (unreachable)
+  )
+
+  ;; CHECK:      (func $unreachable-non-drop (type $1)
+  ;; CHECK-NEXT:  call $unreachable-non-drop
+  ;; CHECK-NEXT:  unreachable
+  ;; CHECK-NEXT: )
+  ;; ROUNDTRIP:      (func $unreachable-non-drop (type $1)
+  ;; ROUNDTRIP-NEXT:  (call $unreachable-non-drop)
+  ;; ROUNDTRIP-NEXT:  (unreachable)
+  ;; ROUNDTRIP-NEXT: )
+  (func $unreachable-non-drop
+    ;; An unreachable with something other than a drop before it. Check we do
+    ;; not error here.
+    (call $unreachable-non-drop)
+    (unreachable)
+  )
+
+  ;; CHECK:      (func $many-drop-unreachable (type $0) (result i32)
+  ;; CHECK-NEXT:  i32.const 1
+  ;; CHECK-NEXT:  if (result i32)
+  ;; CHECK-NEXT:   call $drop-unreachable
+  ;; CHECK-NEXT:   unreachable
+  ;; CHECK-NEXT:  else
+  ;; CHECK-NEXT:   call $drop-unreachable
+  ;; CHECK-NEXT:   unreachable
+  ;; CHECK-NEXT:  end
+  ;; CHECK-NEXT:  drop
+  ;; CHECK-NEXT:  call $drop-unreachable
+  ;; CHECK-NEXT:  unreachable
+  ;; CHECK-NEXT: )
+  ;; ROUNDTRIP:      (func $many-drop-unreachable (type $0) (result i32)
+  ;; ROUNDTRIP-NEXT:  (drop
+  ;; ROUNDTRIP-NEXT:   (if (result i32)
+  ;; ROUNDTRIP-NEXT:    (i32.const 1)
+  ;; ROUNDTRIP-NEXT:    (then
+  ;; ROUNDTRIP-NEXT:     (drop
+  ;; ROUNDTRIP-NEXT:      (call $drop-unreachable)
+  ;; ROUNDTRIP-NEXT:     )
+  ;; ROUNDTRIP-NEXT:     (unreachable)
+  ;; ROUNDTRIP-NEXT:    )
+  ;; ROUNDTRIP-NEXT:    (else
+  ;; ROUNDTRIP-NEXT:     (drop
+  ;; ROUNDTRIP-NEXT:      (call $drop-unreachable)
+  ;; ROUNDTRIP-NEXT:     )
+  ;; ROUNDTRIP-NEXT:     (unreachable)
+  ;; ROUNDTRIP-NEXT:    )
+  ;; ROUNDTRIP-NEXT:   )
+  ;; ROUNDTRIP-NEXT:  )
+  ;; ROUNDTRIP-NEXT:  (drop
+  ;; ROUNDTRIP-NEXT:   (call $drop-unreachable)
+  ;; ROUNDTRIP-NEXT:  )
+  ;; ROUNDTRIP-NEXT:  (unreachable)
+  ;; ROUNDTRIP-NEXT: )
+  (func $many-drop-unreachable (result i32)
+    ;; Two drop-unreachables in an if. The drop on the if can remain, but all
+    ;; others are removable.
+    (drop
+      (if (result i32)
+        (i32.const 1)
+        (then
+          (drop
+            (call $drop-unreachable)
+          )
+          (unreachable)
+        )
+        (else
+          (drop
+            (call $drop-unreachable)
+          )
+          (unreachable)
+        )
+      )
+    )
+    ;;  Two more outside the if.
+    (drop
+      (call $drop-unreachable)
+    )
+    (unreachable)
+    (drop
+      (call $drop-unreachable)
+    )
+    (unreachable)
+  )
+)