Add EH support for CodeFolding (#2665)

This does two things:
- Treats the target branch of `br_on_exn` as unoptimizables, because it
  is a conditional branch.
- Makes sure we don't move expressions that contain `exnref.pop`, which
  should follow right after `catch`.
- Adds `containsChild` utility function, which can search all children,
  optionally with limited depth. This was actually added to be used in
  CodeFolding but ended up not being used, but wasn't removed in case
  there will be uses later.

4 files changed, 160 insertions, 0 deletions

diff --git a/src/ir/iteration.h b/src/ir/iteration.h
index 289e5fa01..fd275f749 100644
--- a/src/ir/iteration.h
+++ b/src/ir/iteration.h
@@ -82,6 +82,30 @@ public:
   Iterator end() const { return Iterator(*this, children.size()); }
 };
 
+// Returns true if the current expression contains a certain kind of expression,
+// within the given depth of BFS. If depth is -1, this searches all children.
+template<typename T> bool containsChild(Expression* parent, int depth = -1) {
+  std::vector<Expression*> exprs;
+  std::vector<Expression*> nextExprs;
+  exprs.push_back(parent);
+  while (!exprs.empty() && depth > 0) {
+    for (auto* expr : exprs) {
+      for (auto* child : ChildIterator(expr)) {
+        if (child->is<T>()) {
+          return true;
+        }
+        nextExprs.push_back(child);
+      }
+    }
+    exprs.swap(nextExprs);
+    nextExprs.clear();
+    if (depth > 0) {
+      depth--;
+    }
+  }
+  return false;
+}
+
 } // namespace wasm
 
 #endif // wasm_ir_iteration_h
diff --git a/src/passes/CodeFolding.cpp b/src/passes/CodeFolding.cpp
index e4d24e33d..70a357bbf 100644
--- a/src/passes/CodeFolding.cpp
+++ b/src/passes/CodeFolding.cpp
@@ -59,6 +59,7 @@
 
 #include "ir/branch-utils.h"
 #include "ir/effects.h"
+#include "ir/find_all.h"
 #include "ir/label-utils.h"
 #include "ir/utils.h"
 #include "pass.h"
@@ -154,6 +155,8 @@ struct CodeFolding : public WalkerPass<ControlFlowWalker<CodeFolding>> {
     unoptimizables.insert(curr->default_);
   }
 
+  void visitBrOnExn(BrOnExn* curr) { unoptimizables.insert(curr->name); }
+
   void visitUnreachable(Unreachable* curr) {
     // we can only optimize if we are at the end of the parent block
     if (!controlFlowStack.empty()) {
@@ -301,6 +304,14 @@ private:
         // anything exiting that is in all targets is something bad
         return false;
       }
+      // Currently pop instructions are only used for exnref.pop, which is a
+      // pseudo instruction following a catch. We check if the current
+      // expression has a pop child. This can be overly conservative, because
+      // this can also exclude whole try-catches that contain a pop within them.
+      if (getModule()->features.hasExceptionHandling() &&
+          !FindAll<Pop>(item).list.empty()) {
+        return false;
+      }
     }
     return true;
   }
diff --git a/test/passes/remove-unused-names_code-folding.txt b/test/passes/remove-unused-names_code-folding_all-features.txt
index 93f503579..2aa40da24 100644
--- a/test/passes/remove-unused-names_code-folding.txt
+++ b/test/passes/remove-unused-names_code-folding_all-features.txt
@@ -3,6 +3,7 @@
  (type $none_=>_i32 (func (result i32)))
  (type $i32_=>_i32 (func (param i32) (result i32)))
  (type $i32_i32_=>_i32 (func (param i32 i32) (result i32)))
+ (event $e (attr 0) (param))
  (func $ifs (; 0 ;)
   (if
    (i32.const 0)
@@ -1708,4 +1709,78 @@
    (i32.const 2)
   )
  )
+ (func $exnref_pop_test (; 36 ;)
+  (local $exn exnref)
+  (block $folding-inner0
+   (try
+    (try
+     (nop)
+     (catch
+      (local.set $exn
+       (exnref.pop)
+      )
+      (br $folding-inner0)
+     )
+    )
+    (catch
+     (local.set $exn
+      (exnref.pop)
+     )
+     (br $folding-inner0)
+    )
+   )
+   (return)
+  )
+  (drop
+   (i32.const 111)
+  )
+  (drop
+   (i32.const 222)
+  )
+  (drop
+   (i32.const 333)
+  )
+  (unreachable)
+ )
+ (func $br_on_exn_target_block (; 37 ;)
+  (local $exn exnref)
+  (block $x
+   (if
+    (i32.const 0)
+    (block
+     (drop
+      (i32.const 1)
+     )
+     (drop
+      (i32.const 2)
+     )
+     (br $x)
+    )
+   )
+   (if
+    (i32.const 0)
+    (block
+     (drop
+      (i32.const 1)
+     )
+     (drop
+      (i32.const 2)
+     )
+     (br $x)
+    )
+   )
+   (drop
+    (br_on_exn $x $e
+     (local.get $exn)
+    )
+   )
+   (drop
+    (i32.const 1)
+   )
+   (drop
+    (i32.const 2)
+   )
+   (br $x)
+  )
+ )
 )
diff --git a/test/passes/remove-unused-names_code-folding.wast b/test/passes/remove-unused-names_code-folding_all-features.wast
index 9979a0a3c..ad803864e 100644
--- a/test/passes/remove-unused-names_code-folding.wast
+++ b/test/passes/remove-unused-names_code-folding_all-features.wast
@@ -1191,4 +1191,54 @@
       )
     )
   )
+
+  (func $exnref_pop_test (local $exn exnref)
+    (try
+      (try
+        (catch
+          ;; Expressions containing exnref.pop should NOT be taken out and
+          ;; folded.
+          (local.set $exn (exnref.pop))
+          (drop (i32.const 111))
+          (drop (i32.const 222))
+          (drop (i32.const 333))
+          (unreachable)
+        )
+      )
+      (catch
+        (local.set $exn (exnref.pop))
+        (drop (i32.const 111))
+        (drop (i32.const 222))
+        (drop (i32.const 333))
+        (unreachable)
+      )
+    )
+  )
+
+  (event $e (attr 0)) ;; exception with no param
+  (func $br_on_exn_target_block (local $exn exnref)
+    ;; Here this block $x is targeted by br_on_exn, so code folding out of this
+    ;; block should NOT happen.
+    (block $x
+      (if (i32.const 0)
+        (block
+          (drop (i32.const 1))
+          (drop (i32.const 2))
+          (br $x)
+        )
+      )
+      (if (i32.const 0)
+        (block
+          (drop (i32.const 1))
+          (drop (i32.const 2))
+          (br $x)
+        )
+      )
+      (drop (br_on_exn $x $e (local.get $exn)))
+      ;; no fallthrough, another thing to merge
+      (drop (i32.const 1))
+      (drop (i32.const 2))
+      (br $x)
+    )
+  )
 )