RemoveUnusedBrs: Make branches to traps just trap themselves

src/passes/RemoveUnusedBrs.cpp

@@ -1214,6 +1214,8 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
         for (auto target : relevantTargets) {
           labelToBranches[target].push_back(curr);
         }
+        // TODO: If all but one target trap, in TNH we can unconditionally jump
+        //       to the non-trapping place.
       }
 
       void visitBlock(Block* curr) {
@@ -1231,14 +1233,39 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
             }
           }
         } else if (list.size() == 2) {
-          // if this block has two children, a child-block and a simple jump,
-          // then jumps to child-block can be replaced with jumps to the new
-          // target
-          auto* child = list[0]->dynCast<Block>();
-          auto* jump = list[1]->dynCast<Break>();
-          if (child && child->name.is() && jump &&
-              ExpressionAnalyzer::isSimple(jump)) {
-            redirectBranches(child, jump->name);
+          // With two items, we look for this form:
+          //
+          //  (block ;; curr
+          //    (block $child ;; child
+          //    )
+          //    after
+          //  )
+          //
+          // Anything branching to the child will end up in that second
+          // instruction |after|. That lets us optimize the cases where |after|
+          // is a br or an unreachable.
+          if (auto* child = list[0]->dynCast<Block>()) {
+            if (child->name) {
+              if (auto* jump = list[1]->dynCast<Break>()) {
+                if (ExpressionAnalyzer::isSimple(jump)) {
+                  // Unconditional jumps to the child can skip ahead to where
+                  // the child jumps.
+                  redirectBranches(child, jump->name);
+                }
+              } else if (list[1]->dynCast<Unreachable>()) {
+                // Unconditional jumps to a trap can just trap.
+                for (auto* branch : labelToBranches[child->name]) {
+                  if (auto* br = branch->dynCast<Break>()) {
+                    if (ExpressionAnalyzer::isSimple(br)) {
+                      // It is safe to just modify the br in-place because it is
+                      // only jumping here, so this the last place that will
+                      // read it.
+                      ExpressionManipulator::unreachable(br);
+                    }
+                  }
+                }
+              }
+            }
           }
         }
       }

test/lit/passes/remove-unused-brs_trap.wast

@@ -0,0 +1,92 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt %s --remove-unused-brs -all -S -o - | filecheck %s
+
+(module
+  ;; CHECK:      (import "a" "b" (func $import (type $0)))
+  (import "a" "b" (func $import))
+
+  ;; CHECK:      (func $jump-to-trap (type $1) (param $x i32) (result i32)
+  ;; CHECK-NEXT:  (block $trap
+  ;; CHECK-NEXT:   (call $import)
+  ;; CHECK-NEXT:   (br_if $trap
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (call $import)
+  ;; CHECK-NEXT:   (if
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:    (then
+  ;; CHECK-NEXT:     (call $import)
+  ;; CHECK-NEXT:     (return
+  ;; CHECK-NEXT:      (i32.const 0)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (else
+  ;; CHECK-NEXT:     (call $import)
+  ;; CHECK-NEXT:     (call $import)
+  ;; CHECK-NEXT:     (unreachable)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (return
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (unreachable)
+  ;; CHECK-NEXT: )
+  (func $jump-to-trap (param $x i32) (result i32)
+    (block $trap
+      (call $import) ;; add effects to avoid other opts
+      ;; This might go to the trap outside, but might not, so we do nothing.
+      (br_if $trap
+        (local.get $x)
+      )
+      (call $import)
+      (if
+        (local.get $x)
+        (then
+          (call $import)
+          (return (i32.const 0))
+        )
+        (else
+          (call $import)
+          (call $import)
+          ;; This goes to a trap, and we can just turn it into a trap.
+          (br $trap)
+        )
+      )
+      (return (i32.const 1))
+    )
+    (unreachable)
+  )
+
+  ;; CHECK:      (func $dispatch (type $2) (param $x i32)
+  ;; CHECK-NEXT:  (block $trap
+  ;; CHECK-NEXT:   (block $mid
+  ;; CHECK-NEXT:    (block $top
+  ;; CHECK-NEXT:     (br_table $top $mid $trap
+  ;; CHECK-NEXT:      (local.get $x)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (call $import)
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (call $import)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (unreachable)
+  ;; CHECK-NEXT: )
+  (func $dispatch (export "a") (param $x i32)
+    ;; A realistic example with br_table (which also avoids other opts from
+    ;; getting in the way).
+    (block $trap
+      (block $mid
+        (block $top
+          (br_table $top $mid $trap (local.get $x))
+        )
+        (call $import) ;; add effects to avoid other opts
+        ;; This goes to a trap, and we can just turn it into a trap.
+        (br $trap)
+      )
+      (call $import)
+    )
+    (unreachable)
+  )
+)

test/passes/remove-unused-brs_enable-multivalue.txt

@@ -2396,7 +2396,7 @@
         (unreachable)
        )
        (then
-        (br $label$3)
+        (unreachable)
        )
       )
      )