From 1166d3a6727c9b4b7092545b465b23170a29fbe5 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Thu, 21 Mar 2024 12:24:47 -0700
Subject: [PATCH 1/2] fix

---
 src/passes/CodeFolding.cpp               | 12 +++++
 test/lit/passes/code-folding-eh-old.wast | 58 +++++++++++++++++++++++-
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/src/passes/CodeFolding.cpp b/src/passes/CodeFolding.cpp
index 0e57a79a21c..57e70b8e6da 100644
--- a/src/passes/CodeFolding.cpp
+++ b/src/passes/CodeFolding.cpp
@@ -58,6 +58,7 @@
 #include <iterator>
 
 #include "ir/branch-utils.h"
+#include "ir/eh-utils.h"
 #include "ir/effects.h"
 #include "ir/find_all.h"
 #include "ir/label-utils.h"
@@ -119,7 +120,11 @@ struct CodeFolding : public WalkerPass<ControlFlowWalker<CodeFolding>> {
 
   // state
 
+  // Set when we optimized and believe another pass is warranted.
   bool anotherPass;
+  // Set when we optimized in a manner that requires EH fixups specifically,
+  // which is generally the case when we wrap things in a block.
+  bool needEHFixups;
 
   // pass state
 
@@ -229,6 +234,7 @@ struct CodeFolding : public WalkerPass<ControlFlowWalker<CodeFolding>> {
       // we must ensure we present the same type as the if had
       ret->finalize(curr->type);
       replaceCurrent(ret);
+      needEHFixups = true;
     } else {
       // if both are blocks, look for a tail we can merge
       auto* left = curr->ifTrue->dynCast<Block>();
@@ -266,6 +272,7 @@ struct CodeFolding : public WalkerPass<ControlFlowWalker<CodeFolding>> {
     anotherPass = true;
     while (anotherPass) {
       anotherPass = false;
+      needEHFixups = false;
       super::doWalkFunction(func);
       optimizeTerminatingTails(unreachableTails);
       // optimize returns at the end, so we can benefit from a fallthrough if
@@ -279,6 +286,9 @@ struct CodeFolding : public WalkerPass<ControlFlowWalker<CodeFolding>> {
       returnTails.clear();
       unoptimizables.clear();
       modifieds.clear();
+      if (needEHFixups) {
+        EHUtils::handleBlockNestedPops(func, *getModule());
+      }
       // if we did any work, types may need to be propagated
       if (anotherPass) {
         ReFinalize().walkFunctionInModule(func, getModule());
@@ -488,6 +498,7 @@ struct CodeFolding : public WalkerPass<ControlFlowWalker<CodeFolding>> {
     // ensure the replacement has the same type, so the outside is not surprised
     block->finalize(oldType);
     replaceCurrent(block);
+    needEHFixups = true;
   }
 
   // optimize tails that terminate control flow in this function, so we
@@ -745,6 +756,7 @@ struct CodeFolding : public WalkerPass<ControlFlowWalker<CodeFolding>> {
     // ensure the replacement has the same type, so the outside is not surprised
     outer->finalize(getFunction()->getResults());
     getFunction()->body = outer;
+    needEHFixups = true;
     return true;
   }
 
diff --git a/test/lit/passes/code-folding-eh-old.wast b/test/lit/passes/code-folding-eh-old.wast
index e83559c2be6..836df867c44 100644
--- a/test/lit/passes/code-folding-eh-old.wast
+++ b/test/lit/passes/code-folding-eh-old.wast
@@ -72,7 +72,7 @@
   ;; CHECK-NEXT: )
   (func $foo)
 
-  ;; CHECK:      (func $try-call-optimize-terminating-tails (type $2) (result i32)
+  ;; CHECK:      (func $try-call-optimize-terminating-tails (type $1) (result i32)
   ;; CHECK-NEXT:  (try
   ;; CHECK-NEXT:   (do
   ;; CHECK-NEXT:    (call $foo)
@@ -155,4 +155,60 @@
       (unreachable)
     )
   )
+
+  ;; CHECK:      (func $if-arms-in-catch (type $1) (result i32)
+  ;; CHECK-NEXT:  (local $0 i32)
+  ;; CHECK-NEXT:  (try
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch $e-i32
+  ;; CHECK-NEXT:    (local.set $0
+  ;; CHECK-NEXT:     (pop i32)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (block
+  ;; CHECK-NEXT:     (block
+  ;; CHECK-NEXT:      (drop
+  ;; CHECK-NEXT:       (local.get $0)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (drop
+  ;; CHECK-NEXT:       (i32.eqz
+  ;; CHECK-NEXT:        (i32.const 1)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (unreachable)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $if-arms-in-catch (result i32)
+    (try
+      (do
+        (unreachable)
+      )
+      (catch $e-i32
+        ;; These if arms can be folded, after which the if is replaced by a
+        ;; block, so we need a fixup for the pop.
+        (if
+          (pop i32)
+          (then
+            (drop
+              (i32.eqz
+                (i32.const 1)
+              )
+            )
+          )
+          (else
+            (drop
+              (i32.eqz
+                (i32.const 1)
+              )
+            )
+          )
+        )
+        (unreachable)
+      )
+    )
+  )
 )

From a6f1eab3198d4306b3a70529d9dc93ce9d9c6d33 Mon Sep 17 00:00:00 2001
From: Alon Zakai <azakai@google.com>
Date: Fri, 22 Mar 2024 10:30:49 -0700
Subject: [PATCH 2/2] format

---
 src/passes/CodeFolding.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/passes/CodeFolding.cpp b/src/passes/CodeFolding.cpp
index 57e70b8e6da..5a18937bbea 100644
--- a/src/passes/CodeFolding.cpp
+++ b/src/passes/CodeFolding.cpp
@@ -58,8 +58,8 @@
 #include <iterator>
 
 #include "ir/branch-utils.h"
-#include "ir/eh-utils.h"
 #include "ir/effects.h"
+#include "ir/eh-utils.h"
 #include "ir/find_all.h"
 #include "ir/label-utils.h"
 #include "ir/utils.h"