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

WebAssembly · Sep 16, 2024 · 4913080 · 4913080
1 parent 29d30fa
commit 4913080
Show file tree

Hide file tree

Showing 2 changed files with 141 additions and 0 deletions.
diff --git 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) {

diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
@@ -4485,3 +4485,127 @@
     )
   )
 )
+
+;; Pops require fixups.
+(module
+  (type $struct (struct (field (mut i32))))
+
+  (type $array (array (mut i32)))
+
+  ;; CHECK:      (type $0 (func))
+
+  ;; CHECK:      (type $1 (func (param i32)))
+
+  ;; CHECK:      (type $2 (struct (field (mut i32)) (field (mut i32)) (field (mut i32))))
+
+  ;; CHECK:      (tag $tag (param i32))
+  (tag $tag (param i32))
+
+  ;; CHECK:      (func $struct-with-pop (type $0)
+  ;; CHECK-NEXT:  (local $0 i32)
+  ;; CHECK-NEXT:  (local $1 i32)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (try
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:    (nop)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch $tag
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (pop i32)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (block (result nullref)
+  ;; CHECK-NEXT:      (local.set $1
+  ;; CHECK-NEXT:       (local.get $2)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $0
+  ;; CHECK-NEXT:       (local.get $1)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (ref.null none)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $struct-with-pop
+    (try
+      (do
+        (nop)
+      )
+      (catch $tag
+        (drop
+          ;; We create a block when we replace the struct with locals, which the
+          ;; pop must be moved out of.
+          (struct.new $struct
+            (pop i32)
+          )
+        )
+      )
+    )
+  )
+
+  ;; CHECK:      (func $array-with-pop (type $0)
+  ;; CHECK-NEXT:  (local $0 i32)
+  ;; CHECK-NEXT:  (local $1 i32)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (local $3 i32)
+  ;; CHECK-NEXT:  (local $4 i32)
+  ;; CHECK-NEXT:  (local $5 i32)
+  ;; CHECK-NEXT:  (local $6 i32)
+  ;; CHECK-NEXT:  (local $7 i32)
+  ;; CHECK-NEXT:  (try
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:    (nop)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch $tag
+  ;; CHECK-NEXT:    (local.set $7
+  ;; CHECK-NEXT:     (pop i32)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (block (result (ref null $2))
+  ;; CHECK-NEXT:      (local.set $0
+  ;; CHECK-NEXT:       (local.get $7)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (block (result nullref)
+  ;; CHECK-NEXT:       (local.set $4
+  ;; CHECK-NEXT:        (local.get $0)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $5
+  ;; CHECK-NEXT:        (local.get $0)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $6
+  ;; CHECK-NEXT:        (local.get $0)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $1
+  ;; CHECK-NEXT:        (local.get $4)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $2
+  ;; CHECK-NEXT:        (local.get $5)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $3
+  ;; CHECK-NEXT:        (local.get $6)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $array-with-pop
+    (try
+      (do
+        (nop)
+      )
+      (catch $tag
+        (drop
+          ;; As above, but with an array
+          (array.new $array
+            (pop i32)
+            (i32.const 3)
+          )
+        )
+      )
+    )
+  )
+)