Fix Heap2Local on pops inside of newly-created blocks (#6938)

1 files changed, 17 insertions, 0 deletions

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 09a50eea8..c31bc8436 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -152,6 +152,7 @@
 
 #include "ir/bits.h"
 #include "ir/branch-utils.h"
+#include "ir/eh-utils.h"
 #include "ir/find_all.h"
 #include "ir/local-graph.h"
 #include "ir/parents.h"
@@ -1191,9 +1192,16 @@ struct Heap2Local {
         // some cases, so to be careful here use a fairly small limit.
         return size < 20;
       }
+
+      // Also note if a pop exists here, as they may require fixups.
+      bool hasPop = false;
+
+      void visitPop(Pop* curr) { hasPop = true; }
     } finder;
     finder.walk(func->body);
 
+    bool optimized = false;
+
     // First, lower non-escaping arrays into structs. That allows us to handle
     // arrays in a single place, and let all the rest of this pass assume we are
     // working on structs. We are in fact only optimizing struct-like arrays
@@ -1215,6 +1223,7 @@ struct Heap2Local {
         auto* structNew =
           Array2Struct(allocation, analyzer, func, wasm).structNew;
         Struct2Local(structNew, analyzer, func, wasm);
+        optimized = true;
       }
     }
 
@@ -1231,8 +1240,16 @@ struct Heap2Local {
         localGraph, parents, branchTargets, passOptions, wasm);
       if (!analyzer.escapes(allocation)) {
         Struct2Local(allocation, analyzer, func, wasm);
+        optimized = true;
       }
     }
+
+    // We conservatively run the EH pop fixup if this function has a 'pop' and
+    // if we have ever optimized, as all of the things we do here involve
+    // creating blocks, so we might have moved pops into the blocks.
+    if (finder.hasPop && optimized) {
+      EHUtils::handleBlockNestedPops(func, wasm);
+    }
   }
 
   bool canHandleAsLocal(const Field& field) {