From d7d17c925af6ee1fd78f39d3fbf2cb22de9e74eb Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Wed, 17 May 2023 16:10:57 -0700
Subject: 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.).
---
 src/passes/Vacuum.cpp | 80 ++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 79 insertions(+), 1 deletion(-)

(limited to 'src')

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;
-- 
cgit v1.2.3