Vacuum code leading up to a trap in TrapsNeverHappen mode (#5228)

This adds two rules to vacuum in TNH mode:

if (..) trap()  =>  if (..) {}

{ stuff, trap() }  =>  {}

That is, we assume traps never happen so an if will not branch to one, and
code right before a trap can be assumed to not execute. Together, we should
be removing practically all possible code in TNH mode (though we could also
add support for br_if etc.).

3 files changed, 515 insertions, 3 deletions

diff --git a/src/passes/Vacuum.cpp b/src/passes/Vacuum.cpp
index 8d89553b2..41a8e02da 100644
--- a/src/passes/Vacuum.cpp
+++ b/src/passes/Vacuum.cpp
@@ -130,9 +130,55 @@ struct Vacuum : public WalkerPass<ExpressionStackWalker<Vacuum>> {
   }
 
   void visitBlock(Block* curr) {
+    auto& list = curr->list;
+
+    // If traps are assumed to never happen, we can remove code on paths that
+    // must reach a trap:
+    //
+    //  (block
+    //    (i32.store ..)
+    //    (br_if ..)      ;; execution branches here, so the first store remains
+    //    (i32.store ..)  ;; this store can be removed
+    //    (unreachable);
+    //  )
+    //
+    // For this to be useful we need to have at least 2 elements: something to
+    // remove, and an unreachable.
+    if (getPassOptions().trapsNeverHappen && list.size() >= 2) {
+      // Go backwards. When we find a trap, mark the things before it as heading
+      // to a trap.
+      auto headingToTrap = false;
+      for (int i = list.size() - 1; i >= 0; i--) {
+        if (list[i]->is<Unreachable>()) {
+          headingToTrap = true;
+          continue;
+        }
+
+        if (!headingToTrap) {
+          continue;
+        }
+
+        // Check if we may no longer be heading to a trap. Two situations count
+        // here: Control flow might branch, or we might call (since a call might
+        // reach an import; see notes on that in pass.h:trapsNeverHappen).
+        //
+        // We also cannot remove a pop as it is necessary for structural
+        // reasons.
+        EffectAnalyzer effects(getPassOptions(), *getModule(), list[i]);
+        if (effects.transfersControlFlow() || effects.calls ||
+            effects.danglingPop) {
+          headingToTrap = false;
+          continue;
+        }
+
+        // This code can be removed! Turn it into a nop, and leave it for the
+        // code lower down to finish cleaning up.
+        ExpressionManipulator::nop(list[i]);
+      }
+    }
+
     // compress out nops and other dead code
     int skip = 0;
-    auto& list = curr->list;
     size_t size = list.size();
     for (size_t z = 0; z < size; z++) {
       auto* child = list[z];
@@ -210,6 +256,38 @@ struct Vacuum : public WalkerPass<ExpressionStackWalker<Vacuum>> {
       return;
     }
     // from here on, we can assume the condition executed
+
+    // In trapsNeverHappen mode, a definitely-trapping arm can be assumed to not
+    // happen. Such conditional code can be assumed to never be reached in this
+    // mode.
+    //
+    // Ignore the case of an unreachable if, such as having both arms be
+    // unreachable. In that case we'd need to fix up the IR to avoid changing
+    // the type; leave that for DCE to simplify first. After checking that
+    // curr->type != unreachable, we can assume that only one of the arms is
+    // unreachable (at most).
+    if (getPassOptions().trapsNeverHappen && curr->type != Type::unreachable) {
+      auto optimizeArm = [&](Expression* arm, Expression* otherArm) {
+        if (!arm->is<Unreachable>()) {
+          return false;
+        }
+        Builder builder(*getModule());
+        Expression* rep = builder.makeDrop(curr->condition);
+        if (otherArm) {
+          rep = builder.makeSequence(rep, otherArm);
+        }
+        replaceCurrent(rep);
+        return true;
+      };
+
+      // As mentioned above, do not try to optimize both arms; leave that case
+      // for DCE.
+      if (optimizeArm(curr->ifTrue, curr->ifFalse) ||
+          (curr->ifFalse && optimizeArm(curr->ifFalse, curr->ifTrue))) {
+        return;
+      }
+    }
+
     if (curr->ifFalse) {
       if (curr->ifFalse->is<Nop>()) {
         curr->ifFalse = nullptr;
diff --git a/test/lit/passes/vacuum-tnh-mvp.wast b/test/lit/passes/vacuum-tnh-mvp.wast
new file mode 100644
index 000000000..93d9bb203
--- /dev/null
+++ b/test/lit/passes/vacuum-tnh-mvp.wast
@@ -0,0 +1,33 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt %s --vacuum -tnh -S -o - | filecheck %s
+
+(module
+  (memory 1 1)
+
+  ;; CHECK:      (func $block-unreachable-but-call
+  ;; CHECK-NEXT:  (i32.store
+  ;; CHECK-NEXT:   (i32.const 0)
+  ;; CHECK-NEXT:   (i32.const 1)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (call $block-unreachable-but-call)
+  ;; CHECK-NEXT:  (unreachable)
+  ;; CHECK-NEXT: )
+  (func $block-unreachable-but-call
+    ;; A call cannot be removed, even if it leads to a trap, since it might have
+    ;; non-trap effects (like mayNotReturn). We can remove the store after it,
+    ;; though.
+    ;;
+    ;; This duplicates a test in vacuum-tnh but in MVP mode (to check for a
+    ;; possible bug with the throws effect which varies based on features).
+    (i32.store
+      (i32.const 0)
+      (i32.const 1)
+    )
+    (call $block-unreachable-but-call)
+    (i32.store
+      (i32.const 2)
+      (i32.const 3)
+    )
+    (unreachable)
+  )
+)
diff --git a/test/lit/passes/vacuum-tnh.wast b/test/lit/passes/vacuum-tnh.wast
index e4ce9765d..bac1c3eb6 100644
--- a/test/lit/passes/vacuum-tnh.wast
+++ b/test/lit/passes/vacuum-tnh.wast
@@ -6,10 +6,16 @@
 ;; RUN: wasm-opt %s --vacuum                      -all -S -o - | filecheck %s --check-prefix=NO_TNH
 
 (module
-  (memory 1 1)
-
   ;; YESTNH:      (type $struct (struct (field (mut i32))))
+
+  ;; YESTNH:      (tag $tag (param i32))
   ;; NO_TNH:      (type $struct (struct (field (mut i32))))
+
+  ;; NO_TNH:      (tag $tag (param i32))
+  (tag $tag (param i32))
+
+  (memory 1 1)
+
   (type $struct (struct (field (mut i32))))
 
   ;; YESTNH:      (func $drop (type $i32_anyref_=>_none) (param $x i32) (param $y anyref)
@@ -265,4 +271,399 @@
       )
     )
   )
+
+  ;; YESTNH:      (func $if-unreachable (type $i32_=>_none) (param $p i32)
+  ;; YESTNH-NEXT:  (drop
+  ;; YESTNH-NEXT:   (local.get $p)
+  ;; YESTNH-NEXT:  )
+  ;; YESTNH-NEXT:  (block
+  ;; YESTNH-NEXT:   (drop
+  ;; YESTNH-NEXT:    (local.get $p)
+  ;; YESTNH-NEXT:   )
+  ;; YESTNH-NEXT:   (call $if-unreachable
+  ;; YESTNH-NEXT:    (i32.const 0)
+  ;; YESTNH-NEXT:   )
+  ;; YESTNH-NEXT:  )
+  ;; YESTNH-NEXT:  (if
+  ;; YESTNH-NEXT:   (local.get $p)
+  ;; YESTNH-NEXT:   (unreachable)
+  ;; YESTNH-NEXT:   (unreachable)
+  ;; YESTNH-NEXT:  )
+  ;; YESTNH-NEXT: )
+  ;; NO_TNH:      (func $if-unreachable (type $i32_=>_none) (param $p i32)
+  ;; NO_TNH-NEXT:  (if
+  ;; NO_TNH-NEXT:   (local.get $p)
+  ;; NO_TNH-NEXT:   (unreachable)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT:  (if
+  ;; NO_TNH-NEXT:   (local.get $p)
+  ;; NO_TNH-NEXT:   (call $if-unreachable
+  ;; NO_TNH-NEXT:    (i32.const 0)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:   (unreachable)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT:  (if
+  ;; NO_TNH-NEXT:   (local.get $p)
+  ;; NO_TNH-NEXT:   (unreachable)
+  ;; NO_TNH-NEXT:   (unreachable)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT: )
+  (func $if-unreachable (param $p i32)
+    ;; The if arm can be nopped, as in tnh we assume we never reach it.
+    (if
+      (local.get $p)
+      (unreachable)
+    )
+    ;; This else arm can be removed.
+    (if
+      (local.get $p)
+      (call $if-unreachable
+        (i32.const 0)
+      )
+      (unreachable)
+    )
+    ;; Both of these can be removed, but we leave this for DCE to handle.
+    (if
+      (local.get $p)
+      (unreachable)
+      (unreachable)
+    )
+  )
+
+  ;; YESTNH:      (func $if-unreachable-value (type $i32_=>_i32) (param $p i32) (result i32)
+  ;; YESTNH-NEXT:  (drop
+  ;; YESTNH-NEXT:   (local.get $p)
+  ;; YESTNH-NEXT:  )
+  ;; YESTNH-NEXT:  (i32.const 1)
+  ;; YESTNH-NEXT: )
+  ;; NO_TNH:      (func $if-unreachable-value (type $i32_=>_i32) (param $p i32) (result i32)
+  ;; NO_TNH-NEXT:  (if (result i32)
+  ;; NO_TNH-NEXT:   (local.get $p)
+  ;; NO_TNH-NEXT:   (unreachable)
+  ;; NO_TNH-NEXT:   (i32.const 1)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT: )
+  (func $if-unreachable-value (param $p i32) (result i32)
+    ;; When removing the unreachable arm we must update the IR properly, as it
+    ;; cannot have a nop there.
+    (if (result i32)
+      (local.get $p)
+      (unreachable)
+      (i32.const 1)
+    )
+  )
+
+  ;; YESTNH:      (func $if-unreachable-value-2 (type $i32_=>_i32) (param $p i32) (result i32)
+  ;; YESTNH-NEXT:  (drop
+  ;; YESTNH-NEXT:   (local.get $p)
+  ;; YESTNH-NEXT:  )
+  ;; YESTNH-NEXT:  (i32.const 1)
+  ;; YESTNH-NEXT: )
+  ;; NO_TNH:      (func $if-unreachable-value-2 (type $i32_=>_i32) (param $p i32) (result i32)
+  ;; NO_TNH-NEXT:  (if (result i32)
+  ;; NO_TNH-NEXT:   (local.get $p)
+  ;; NO_TNH-NEXT:   (i32.const 1)
+  ;; NO_TNH-NEXT:   (unreachable)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT: )
+  (func $if-unreachable-value-2 (param $p i32) (result i32)
+    ;; As above but in the other arm.
+    (if (result i32)
+      (local.get $p)
+      (i32.const 1)
+      (unreachable)
+    )
+  )
+
+  ;; YESTNH:      (func $block-unreachable (type $i32_=>_none) (param $p i32)
+  ;; YESTNH-NEXT:  (if
+  ;; YESTNH-NEXT:   (local.get $p)
+  ;; YESTNH-NEXT:   (block
+  ;; YESTNH-NEXT:    (i32.store
+  ;; YESTNH-NEXT:     (i32.const 0)
+  ;; YESTNH-NEXT:     (i32.const 1)
+  ;; YESTNH-NEXT:    )
+  ;; YESTNH-NEXT:    (if
+  ;; YESTNH-NEXT:     (local.get $p)
+  ;; YESTNH-NEXT:     (return)
+  ;; YESTNH-NEXT:    )
+  ;; YESTNH-NEXT:    (unreachable)
+  ;; YESTNH-NEXT:   )
+  ;; YESTNH-NEXT:  )
+  ;; YESTNH-NEXT: )
+  ;; NO_TNH:      (func $block-unreachable (type $i32_=>_none) (param $p i32)
+  ;; NO_TNH-NEXT:  (if
+  ;; NO_TNH-NEXT:   (local.get $p)
+  ;; NO_TNH-NEXT:   (block
+  ;; NO_TNH-NEXT:    (i32.store
+  ;; NO_TNH-NEXT:     (i32.const 0)
+  ;; NO_TNH-NEXT:     (i32.const 1)
+  ;; NO_TNH-NEXT:    )
+  ;; NO_TNH-NEXT:    (if
+  ;; NO_TNH-NEXT:     (local.get $p)
+  ;; NO_TNH-NEXT:     (return)
+  ;; NO_TNH-NEXT:    )
+  ;; NO_TNH-NEXT:    (i32.store
+  ;; NO_TNH-NEXT:     (i32.const 2)
+  ;; NO_TNH-NEXT:     (i32.const 3)
+  ;; NO_TNH-NEXT:    )
+  ;; NO_TNH-NEXT:    (unreachable)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT: )
+  (func $block-unreachable (param $p i32)
+    (if
+      (local.get $p)
+      (block
+        (i32.store
+          (i32.const 0)
+          (i32.const 1)
+        )
+        (if
+          (local.get $p)
+          (return)
+        )
+        ;; This store can be removed as it leads up to an unreachable which we
+        ;; assume is never reached.
+        (i32.store
+          (i32.const 2)
+          (i32.const 3)
+        )
+        (unreachable)
+      )
+    )
+  )
+
+  ;; YESTNH:      (func $block-unreachable-named (type $i32_=>_none) (param $p i32)
+  ;; YESTNH-NEXT:  (if
+  ;; YESTNH-NEXT:   (local.get $p)
+  ;; YESTNH-NEXT:   (block $named
+  ;; YESTNH-NEXT:    (i32.store
+  ;; YESTNH-NEXT:     (i32.const 0)
+  ;; YESTNH-NEXT:     (i32.const 1)
+  ;; YESTNH-NEXT:    )
+  ;; YESTNH-NEXT:    (br_if $named
+  ;; YESTNH-NEXT:     (local.get $p)
+  ;; YESTNH-NEXT:    )
+  ;; YESTNH-NEXT:    (unreachable)
+  ;; YESTNH-NEXT:   )
+  ;; YESTNH-NEXT:  )
+  ;; YESTNH-NEXT: )
+  ;; NO_TNH:      (func $block-unreachable-named (type $i32_=>_none) (param $p i32)
+  ;; NO_TNH-NEXT:  (if
+  ;; NO_TNH-NEXT:   (local.get $p)
+  ;; NO_TNH-NEXT:   (block $named
+  ;; NO_TNH-NEXT:    (i32.store
+  ;; NO_TNH-NEXT:     (i32.const 0)
+  ;; NO_TNH-NEXT:     (i32.const 1)
+  ;; NO_TNH-NEXT:    )
+  ;; NO_TNH-NEXT:    (br_if $named
+  ;; NO_TNH-NEXT:     (local.get $p)
+  ;; NO_TNH-NEXT:    )
+  ;; NO_TNH-NEXT:    (i32.store
+  ;; NO_TNH-NEXT:     (i32.const 2)
+  ;; NO_TNH-NEXT:     (i32.const 3)
+  ;; NO_TNH-NEXT:    )
+  ;; NO_TNH-NEXT:    (unreachable)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT: )
+  (func $block-unreachable-named (param $p i32)
+    (if
+      (local.get $p)
+      (block $named
+        (i32.store
+          (i32.const 0)
+          (i32.const 1)
+        )
+        ;; As above, but now the block is named and we use a br_if. We should
+        ;; again only remove the last store.
+        (br_if $named
+          (local.get $p)
+        )
+        (i32.store
+          (i32.const 2)
+          (i32.const 3)
+        )
+        (unreachable)
+      )
+    )
+  )
+
+  ;; YESTNH:      (func $block-unreachable-all (type $i32_=>_none) (param $p i32)
+  ;; YESTNH-NEXT:  (nop)
+  ;; YESTNH-NEXT: )
+  ;; NO_TNH:      (func $block-unreachable-all (type $i32_=>_none) (param $p i32)
+  ;; NO_TNH-NEXT:  (if
+  ;; NO_TNH-NEXT:   (local.get $p)
+  ;; NO_TNH-NEXT:   (block
+  ;; NO_TNH-NEXT:    (i32.store
+  ;; NO_TNH-NEXT:     (i32.const 0)
+  ;; NO_TNH-NEXT:     (i32.const 1)
+  ;; NO_TNH-NEXT:    )
+  ;; NO_TNH-NEXT:    (i32.store
+  ;; NO_TNH-NEXT:     (i32.const 2)
+  ;; NO_TNH-NEXT:     (i32.const 3)
+  ;; NO_TNH-NEXT:    )
+  ;; NO_TNH-NEXT:    (unreachable)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT: )
+  (func $block-unreachable-all (param $p i32)
+    (if
+      (local.get $p)
+      (block
+        ;; Both stores can be removed, and even the entire if arm and then the
+        ;; entire if.
+        (i32.store
+          (i32.const 0)
+          (i32.const 1)
+        )
+        (i32.store
+          (i32.const 2)
+          (i32.const 3)
+        )
+        (unreachable)
+      )
+    )
+  )
+
+  ;; YESTNH:      (func $block-unreachable-but-call (type $none_=>_none)
+  ;; YESTNH-NEXT:  (i32.store
+  ;; YESTNH-NEXT:   (i32.const 0)
+  ;; YESTNH-NEXT:   (i32.const 1)
+  ;; YESTNH-NEXT:  )
+  ;; YESTNH-NEXT:  (call $block-unreachable-but-call)
+  ;; YESTNH-NEXT:  (unreachable)
+  ;; YESTNH-NEXT: )
+  ;; NO_TNH:      (func $block-unreachable-but-call (type $none_=>_none)
+  ;; NO_TNH-NEXT:  (i32.store
+  ;; NO_TNH-NEXT:   (i32.const 0)
+  ;; NO_TNH-NEXT:   (i32.const 1)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT:  (call $block-unreachable-but-call)
+  ;; NO_TNH-NEXT:  (i32.store
+  ;; NO_TNH-NEXT:   (i32.const 2)
+  ;; NO_TNH-NEXT:   (i32.const 3)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT:  (unreachable)
+  ;; NO_TNH-NEXT: )
+  (func $block-unreachable-but-call
+    ;; A call cannot be removed, even if it leads to a trap, since it might have
+    ;; non-trap effects (like mayNotReturn). We can remove the store after it,
+    ;; though.
+    (i32.store
+      (i32.const 0)
+      (i32.const 1)
+    )
+    (call $block-unreachable-but-call)
+    (i32.store
+      (i32.const 2)
+      (i32.const 3)
+    )
+    (unreachable)
+  )
+
+  ;; YESTNH:      (func $catch-pop (type $none_=>_none)
+  ;; YESTNH-NEXT:  (try $try
+  ;; YESTNH-NEXT:   (do
+  ;; YESTNH-NEXT:    (call $catch-pop)
+  ;; YESTNH-NEXT:   )
+  ;; YESTNH-NEXT:   (catch $tag
+  ;; YESTNH-NEXT:    (drop
+  ;; YESTNH-NEXT:     (pop i32)
+  ;; YESTNH-NEXT:    )
+  ;; YESTNH-NEXT:    (unreachable)
+  ;; YESTNH-NEXT:   )
+  ;; YESTNH-NEXT:  )
+  ;; YESTNH-NEXT: )
+  ;; NO_TNH:      (func $catch-pop (type $none_=>_none)
+  ;; NO_TNH-NEXT:  (try $try
+  ;; NO_TNH-NEXT:   (do
+  ;; NO_TNH-NEXT:    (call $catch-pop)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:   (catch $tag
+  ;; NO_TNH-NEXT:    (drop
+  ;; NO_TNH-NEXT:     (pop i32)
+  ;; NO_TNH-NEXT:    )
+  ;; NO_TNH-NEXT:    (i32.store
+  ;; NO_TNH-NEXT:     (i32.const 0)
+  ;; NO_TNH-NEXT:     (i32.const 1)
+  ;; NO_TNH-NEXT:    )
+  ;; NO_TNH-NEXT:    (unreachable)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT: )
+  (func $catch-pop
+    (try $try
+      (do
+        ;; Put a call here so the entire try-catch is not removed as trivial.
+        (call $catch-pop)
+      )
+      (catch $tag
+        ;; A pop on the way to a trap cannot be removed. But the store can.
+        ;; TODO: The pop can actually be removed since it is valid per the spec
+        ;;       because of the unreachable afterwards. We need to fix our
+        ;;       validation rules to handle that though.
+        (drop
+          (pop i32)
+        )
+        (i32.store
+          (i32.const 0)
+          (i32.const 1)
+        )
+        (unreachable)
+      )
+    )
+  )
+
+  ;; YESTNH:      (func $loop-unreachable (type $i32_=>_none) (param $p i32)
+  ;; YESTNH-NEXT:  (loop $loop
+  ;; YESTNH-NEXT:   (i32.store
+  ;; YESTNH-NEXT:    (i32.const 0)
+  ;; YESTNH-NEXT:    (i32.const 1)
+  ;; YESTNH-NEXT:   )
+  ;; YESTNH-NEXT:   (if
+  ;; YESTNH-NEXT:    (local.get $p)
+  ;; YESTNH-NEXT:    (br $loop)
+  ;; YESTNH-NEXT:   )
+  ;; YESTNH-NEXT:   (unreachable)
+  ;; YESTNH-NEXT:  )
+  ;; YESTNH-NEXT: )
+  ;; NO_TNH:      (func $loop-unreachable (type $i32_=>_none) (param $p i32)
+  ;; NO_TNH-NEXT:  (loop $loop
+  ;; NO_TNH-NEXT:   (i32.store
+  ;; NO_TNH-NEXT:    (i32.const 0)
+  ;; NO_TNH-NEXT:    (i32.const 1)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:   (if
+  ;; NO_TNH-NEXT:    (local.get $p)
+  ;; NO_TNH-NEXT:    (br $loop)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:   (i32.store
+  ;; NO_TNH-NEXT:    (i32.const 2)
+  ;; NO_TNH-NEXT:    (i32.const 3)
+  ;; NO_TNH-NEXT:   )
+  ;; NO_TNH-NEXT:   (unreachable)
+  ;; NO_TNH-NEXT:  )
+  ;; NO_TNH-NEXT: )
+  (func $loop-unreachable (param $p i32)
+    (loop $loop
+      (i32.store
+        (i32.const 0)
+        (i32.const 1)
+      )
+      (if
+        (local.get $p)
+        (br $loop)
+      )
+      ;; This store can be removed as it leads up to an unreachable which we
+      ;; assume is never reached.
+      (i32.store
+        (i32.const 2)
+        (i32.const 3)
+      )
+      (unreachable)
+    )
+  )
 )