[Wasm EH] Optimize throws caught by TryTable into breaks (#6980)

E.g.

(try_table (catch_all $catch)
  (throw $e)
)

=>

(try_table (catch_all $catch)
  (br $catch)
)

This can then allow other passes to remove the TryTable, if no throwing
things remain.

3 files changed, 407 insertions, 18 deletions

diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index 7fe169bdb..0c316e7e7 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -18,16 +18,17 @@
 // Removes branches for which we go to where they go anyhow
 //
 
-#include <ir/branch-utils.h>
-#include <ir/cost.h>
-#include <ir/effects.h>
-#include <ir/gc-type-utils.h>
-#include <ir/literal-utils.h>
-#include <ir/utils.h>
-#include <parsing.h>
-#include <pass.h>
-#include <wasm-builder.h>
-#include <wasm.h>
+#include "ir/branch-utils.h"
+#include "ir/cost.h"
+#include "ir/drop.h"
+#include "ir/effects.h"
+#include "ir/gc-type-utils.h"
+#include "ir/literal-utils.h"
+#include "ir/utils.h"
+#include "parsing.h"
+#include "pass.h"
+#include "wasm-builder.h"
+#include "wasm.h"
 
 namespace wasm {
 
@@ -260,7 +261,7 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
     } else if (curr->is<Nop>()) {
       // ignore (could be result of a previous cycle)
       self->stopValueFlow();
-    } else if (curr->is<Loop>()) {
+    } else if (curr->is<Loop>()) { // TODO: eh
       // do nothing - it's ok for values to flow out
     } else if (auto* sw = curr->dynCast<Switch>()) {
       self->stopFlow();
@@ -462,13 +463,69 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
     //       later down, see visitLocalSet.
   }
 
+  // A stack of try_tables that are parents of the current expression.
+  std::vector<TryTable*> tryTables;
+
+  static void popTryTable(RemoveUnusedBrs* self, Expression** currp) {
+    assert(!self->tryTables.empty() && self->tryTables.back() == *currp);
+    self->tryTables.pop_back();
+  }
+
+  void visitThrow(Throw* curr) {
+    // If a throw will definitely be caught, and it is not a catch with a
+    // reference, then it is just a branch (i.e. the code is using exceptions as
+    // control flow). Turn it into a branch here so that the rest of the pass
+    // can optimize it with all other branches.
+    //
+    // To do so, look at the closest try and see if it will catch us, and
+    // proceed outwards if not.
+    auto thrownTag = curr->tag;
+    for (int i = tryTables.size() - 1; i >= 0; i--) {
+      auto* tryy = tryTables[i];
+      for (Index j = 0; j < tryy->catchTags.size(); j++) {
+        auto tag = tryy->catchTags[j];
+        // The tag must match, or be a catch_all.
+        if (tag == thrownTag || tag.isNull()) {
+          // This must not be a catch with exnref.
+          if (!tryy->catchRefs[j]) {
+            // Success! Create a break to replace the throw.
+            auto dest = tryy->catchDests[j];
+            auto& wasm = *getModule();
+            Builder builder(wasm);
+            if (!tag.isNull()) {
+              // We are catching a specific tag, so values might be sent.
+              Expression* value = nullptr;
+              if (curr->operands.size() == 1) {
+                value = curr->operands[0];
+              } else if (curr->operands.size() > 1) {
+                value = builder.makeTupleMake(curr->operands);
+              }
+              auto* br = builder.makeBreak(dest, value);
+              replaceCurrent(br);
+              return;
+            }
+
+            // catch_all: no values are sent. Drop the throw's children (while
+            // ignoring parent effects: the parent is a throw, but we have
+            // proven we can remove that effect).
+            auto* br = builder.makeBreak(dest);
+            auto* rep = getDroppedChildrenAndAppend(
+              curr, wasm, getPassOptions(), br, DropMode::IgnoreParentEffects);
+            replaceCurrent(rep);
+          }
+
+          // Return even if we did not optimize: we found our tag was caught.
+          return;
+        }
+      }
+    }
+  }
+
   // override scan to add a pre and a post check task to all nodes
   static void scan(RemoveUnusedBrs* self, Expression** currp) {
     self->pushTask(visitAny, currp);
 
-    auto* iff = (*currp)->dynCast<If>();
-
-    if (iff) {
+    if (auto* iff = (*currp)->dynCast<If>()) {
       if (iff->condition->type == Type::unreachable) {
         // avoid trying to optimize this, we never reach it anyhow
         return;
@@ -484,9 +541,15 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
       self->pushTask(scan, &iff->ifTrue);
       self->pushTask(clear, currp); // clear all flow after the condition
       self->pushTask(scan, &iff->condition);
-    } else {
-      Super::scan(self, currp);
+      return;
+    }
+    if (auto* tryy = (*currp)->dynCast<TryTable>()) {
+      // Push the try we are reaching, and add a task to pop it, after all the
+      // tasks that Super::scan will push for its children.
+      self->tryTables.push_back(tryy);
+      self->pushTask(popTryTable, currp);
     }
+    Super::scan(self, currp);
   }
 
   // optimizes a loop. returns true if we made changes
diff --git a/test/lit/passes/remove-unused-brs-eh.wast b/test/lit/passes/remove-unused-brs-eh.wast
new file mode 100644
index 000000000..a4c6591ab
--- /dev/null
+++ b/test/lit/passes/remove-unused-brs-eh.wast
@@ -0,0 +1,324 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: foreach %s %t wasm-opt -all --remove-unused-brs -S -o - | filecheck %s
+
+(module
+  ;; CHECK:      (tag $e)
+  (tag $e)
+  ;; CHECK:      (tag $f)
+  (tag $f)
+  ;; CHECK:      (tag $g)
+  (tag $g)
+
+  ;; CHECK:      (func $throw-caught-all (type $0)
+  ;; CHECK-NEXT:  (block $catch
+  ;; CHECK-NEXT:   (try_table (catch_all $catch)
+  ;; CHECK-NEXT:    (br $catch)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-caught-all
+    (block $catch
+      (try_table (catch_all $catch)
+        ;; This throw can be a br.
+        (throw $e)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $throw-caught-all-more (type $2) (param $x i32)
+  ;; CHECK-NEXT:  (block $catch
+  ;; CHECK-NEXT:   (try_table (catch_all $catch)
+  ;; CHECK-NEXT:    (br_if $catch
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-caught-all-more (param $x i32)
+    (block $catch
+      (try_table (catch_all $catch)
+        ;; Look into nested children. After we turn the throw into a br, it also
+        ;; fuses with the if into a br_if.
+        (if
+          (local.get $x)
+          (then
+            (throw $e)
+          )
+        )
+      )
+    )
+  )
+
+  ;; CHECK:      (func $throw-caught-precise (type $0)
+  ;; CHECK-NEXT:  (block $catch
+  ;; CHECK-NEXT:   (try_table (catch $e $catch)
+  ;; CHECK-NEXT:    (br $catch)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-caught-precise
+    (block $catch
+      ;; We can still optimize here, even though we replaced the catch_all with
+      ;; a precise tag, because the tag matches.
+      (try_table (catch $e $catch)
+        (throw $e)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $throw-caught-precise-later (type $0)
+  ;; CHECK-NEXT:  (block $fail
+  ;; CHECK-NEXT:   (block $catch
+  ;; CHECK-NEXT:    (try_table (catch $f $fail) (catch $e $catch)
+  ;; CHECK-NEXT:     (br $catch)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (call $throw-caught-precise-later)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-caught-precise-later
+    (block $fail
+      (block $catch
+        ;; We can still optimize here, by looking through the tags til we find
+        ;; ours.
+        (try_table (catch $f $fail) (catch $e $catch)
+          (throw $e)
+        )
+      )
+      ;; Add an effect here, so the two blocks are not mergeable.
+      (call $throw-caught-precise-later)
+    )
+  )
+
+  ;; CHECK:      (func $throw-caught-all-later (type $0)
+  ;; CHECK-NEXT:  (block $fail
+  ;; CHECK-NEXT:   (block $catch
+  ;; CHECK-NEXT:    (try_table (catch $f $fail) (catch_all $catch)
+  ;; CHECK-NEXT:     (br $catch)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (call $throw-caught-precise-later)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-caught-all-later
+    (block $fail
+      (block $catch
+        ;; We can still optimize here, by looking through the tags til we find
+        ;; the catch_all
+        (try_table (catch $f $fail) (catch_all $catch)
+          (throw $e)
+        )
+      )
+      ;; Add an effect here, so the two blocks are not mergeable.
+      (call $throw-caught-precise-later)
+    )
+  )
+
+  ;; CHECK:      (func $throw-caught-fail (type $0)
+  ;; CHECK-NEXT:  (block $catch
+  ;; CHECK-NEXT:   (try_table (catch $f $catch)
+  ;; CHECK-NEXT:    (throw $e)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-caught-fail
+    (block $catch
+      ;; The tag does *not* match.
+      (try_table (catch $f $catch)
+        (throw $e)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $throw-caught-outer (type $0)
+  ;; CHECK-NEXT:  (block $fail
+  ;; CHECK-NEXT:   (block $catch
+  ;; CHECK-NEXT:    (try_table (catch $e $catch)
+  ;; CHECK-NEXT:     (try_table (catch $f $fail)
+  ;; CHECK-NEXT:      (br $catch)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (call $throw-caught-precise-later)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-caught-outer
+    (block $fail
+      (block $catch
+        (try_table (catch $e $catch)
+          (try_table (catch $f $fail)
+            ;; This throw can be a br, thanks to the outer try.
+            (throw $e)
+          )
+        )
+      )
+      ;; Add an effect here, so the two blocks are not mergeable.
+      (call $throw-caught-precise-later)
+    )
+  )
+
+  ;; CHECK:      (func $throw-catch-all-ref (type $1) (result exnref)
+  ;; CHECK-NEXT:  (block $catch (result exnref)
+  ;; CHECK-NEXT:   (try_table (catch_all_ref $catch)
+  ;; CHECK-NEXT:    (throw $e)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-catch-all-ref (result exnref)
+    (block $catch (result exnref)
+      (try_table (catch_all_ref $catch)
+        ;; This is caught with a ref, so we do not optimize.
+        (throw $e)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $throw-caught-ref (type $1) (result exnref)
+  ;; CHECK-NEXT:  (block $catch (result exnref)
+  ;; CHECK-NEXT:   (try_table (catch_ref $e $catch)
+  ;; CHECK-NEXT:    (throw $e)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-caught-ref (result exnref)
+    (block $catch (result exnref)
+      (try_table (catch_ref $e $catch)
+        ;; As above, but without catch_all.
+        (throw $e)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $throw-caught-ref-later-all (type $1) (result exnref)
+  ;; CHECK-NEXT:  (block $outer (result exnref)
+  ;; CHECK-NEXT:   (block $catch
+  ;; CHECK-NEXT:    (try_table (catch_ref $e $outer) (catch_all $catch)
+  ;; CHECK-NEXT:     (throw $e)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (unreachable)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-caught-ref-later-all (result exnref)
+    (block $outer (result exnref)
+      (block $catch
+        (try_table (catch_ref $e $outer) (catch_all $catch)
+          ;; This is caught with a ref, before we reach the catch all, so we do
+          ;; not optimize.
+          (throw $e)
+        )
+      )
+      (unreachable)
+    )
+  )
+
+  ;; CHECK:      (func $throw-multi (type $0)
+  ;; CHECK-NEXT:  (block $outer
+  ;; CHECK-NEXT:   (block $middle
+  ;; CHECK-NEXT:    (block $inner
+  ;; CHECK-NEXT:     (try_table (catch $e $outer) (catch $f $middle) (catch_all $inner)
+  ;; CHECK-NEXT:      (br $outer)
+  ;; CHECK-NEXT:      (br $middle)
+  ;; CHECK-NEXT:      (br $inner)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (call $throw-caught-precise-later)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (call $throw-caught-precise-later)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-multi
+    (block $outer
+      (block $middle
+        (block $inner
+          (try_table (catch $e $outer) (catch $f $middle) (catch_all $inner)
+            ;; Multiple throws, optimizable in different ways.
+            (throw $e)
+            (throw $f)
+            (throw $g)
+          )
+        )
+        ;; Add an effect here, so the two blocks are not mergeable.
+        (call $throw-caught-precise-later)
+      )
+      (call $throw-caught-precise-later)
+    )
+  )
+)
+
+(module
+  ;; CHECK:      (import "a" "b" (func $effect (type $1) (result i32)))
+  (import "a" "b" (func $effect (result i32)))
+
+  ;; CHECK:      (tag $e (param i32))
+  (tag $e (param i32))
+
+  ;; CHECK:      (tag $multi (param i32 f64))
+  (tag $multi (param i32 f64))
+
+  ;; CHECK:      (func $throw-caught-all (type $0) (param $x i32)
+  ;; CHECK-NEXT:  (block $catch
+  ;; CHECK-NEXT:   (try_table (catch_all $catch)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (call $effect)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (br $catch)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-caught-all (param $x i32)
+    (block $catch
+      (try_table (catch_all $catch)
+        ;; This throw can be a br. The call must be kept in a drop.
+        (throw $e
+          (call $effect)
+        )
+      )
+    )
+  )
+
+  ;; CHECK:      (func $throw-br-contents (type $1) (result i32)
+  ;; CHECK-NEXT:  (block $catch (result i32)
+  ;; CHECK-NEXT:   (try_table (catch $e $catch)
+  ;; CHECK-NEXT:    (br $catch
+  ;; CHECK-NEXT:     (i32.const 42)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-br-contents (result i32)
+    (block $catch (result i32)
+      (try_table (catch $e $catch)
+        ;; This throw is not caught by catch_all as above, so the value must be
+        ;; sent as a value on the br we optimize it to.
+        (throw $e
+          (i32.const 42)
+        )
+      )
+    )
+  )
+
+  ;; CHECK:      (func $throw-br-contents-multi (type $2) (result i32 f64)
+  ;; CHECK-NEXT:  (block $catch (type $2) (result i32 f64)
+  ;; CHECK-NEXT:   (try_table (catch $multi $catch)
+  ;; CHECK-NEXT:    (br $catch
+  ;; CHECK-NEXT:     (tuple.make 2
+  ;; CHECK-NEXT:      (i32.const 42)
+  ;; CHECK-NEXT:      (f64.const 3.14159)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $throw-br-contents-multi (result i32 f64)
+    ;; As above, but now with a multivalue tag.
+    (block $catch (result i32 f64)
+      (try_table (catch $multi $catch)
+        (throw $multi
+          (i32.const 42)
+          (f64.const 3.14159)
+        )
+      )
+    )
+  )
+)
diff --git a/test/lit/passes/vacuum-eh.wast b/test/lit/passes/vacuum-eh.wast
index ec62f3c58..d42fed56b 100644
--- a/test/lit/passes/vacuum-eh.wast
+++ b/test/lit/passes/vacuum-eh.wast
@@ -151,8 +151,10 @@
   ;; CHECK-NEXT: )
   (func $trivial-catch-all-of-throw (local $0 i32)
     ;; This try_table's body throws, but the catch_all catches it, so the entire
-    ;; try_table could be optimized out. We do this for `try` but not yet for
-    ;; `try_table`.
+    ;; try_table could be optimized out. We do this for `try` but not for
+    ;; `try_table` - we leave such optimizations to --remove-unused-brs (that
+    ;; pass can see that the throw can be converted to a br).
+    
     (block $catch
       (try_table (catch_all $catch)
         (throw $e (i32.const 0))