OptimizeInstructions: Push StructNew down to help it fold away StructSets (#6584)

Heap stores (struct.set) are optimized into the struct.new when they are adjacent
in a statement list.

Pushing struct.new down past irrelevant instructions increases the likelihood that
it ends up adjacent to sets.

2 files changed, 109 insertions, 29 deletions

diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index 81d0aef76..10c1aac9e 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -1864,11 +1864,9 @@ struct OptimizeInstructions
   // =>
   //  (local.set $x (struct.new X' Y Z))
   //
-  // We also handle other struct.sets immediately after this one, but we only
-  // handle the case where they are all in sequence and right after the
-  // local.set (anything in the middle of this pattern will stop us from
-  // optimizing later struct.sets, which might be improved later but would
-  // require an analysis of effects TODO).
+  // We also handle other struct.sets immediately after this one. If the
+  // instruction following the new is not a struct.set we push the new down if
+  // possible.
   void optimizeHeapStores(ExpressionList& list) {
     for (Index i = 0; i < list.size(); i++) {
       auto* localSet = list[i]->dynCast<LocalSet>();
@@ -1880,30 +1878,70 @@ struct OptimizeInstructions
         continue;
       }
 
-      // This local.set of a struct.new looks good. Find struct.sets after it
-      // to optimize.
-      for (Index j = i + 1; j < list.size(); j++) {
+      // This local.set of a struct.new looks good. Find struct.sets after it to
+      // optimize.
+      Index localSetIndex = i;
+      for (Index j = localSetIndex + 1; j < list.size(); j++) {
+
+        // Check that the next instruction is a struct.set on the same local as
+        // the struct.new.
         auto* structSet = list[j]->dynCast<StructSet>();
-        if (!structSet) {
-          // Any time the pattern no longer matches, stop optimizing possible
-          // struct.sets for this struct.new.
-          break;
-        }
-        auto* localGet = structSet->ref->dynCast<LocalGet>();
-        if (!localGet || localGet->index != localSet->index) {
+        auto* localGet =
+          structSet ? structSet->ref->dynCast<LocalGet>() : nullptr;
+        if (!structSet || !localGet || localGet->index != localSet->index) {
+          // Any time the pattern no longer matches, we try to push the
+          // struct.new further down but if it is not possible we stop
+          // optimizing possible struct.sets for this struct.new.
+          if (trySwap(list, localSetIndex, j)) {
+            // Update the index and continue to try again.
+            localSetIndex = j;
+            continue;
+          }
           break;
         }
+
+        // The pattern matches, try to optimize.
         if (!optimizeSubsequentStructSet(new_, structSet, localGet->index)) {
           break;
         } else {
-          // Success. Replace the set with a nop, and continue to
-          // perhaps optimize more.
+          // Success. Replace the set with a nop, and continue to perhaps
+          // optimize more.
           ExpressionManipulator::nop(structSet);
         }
       }
     }
   }
 
+  // Helper function for optimizeHeapStores. Tries pushing the struct.new at
+  // index i down to index j, swapping it with the instruction already at j, so
+  // that it is closer to (potential) later struct.sets.
+  bool trySwap(ExpressionList& list, Index i, Index j) {
+    if (j == list.size() - 1) {
+      // There is no reason to swap with the last element of the list as it
+      // won't match the pattern because there wont be anything after. This also
+      // avoids swapping an instruction that does not leave anything in the
+      // stack by one that could leave something, and that which would be
+      // incorrect.
+      return false;
+    }
+
+    if (list[j]->is<LocalSet>() &&
+        list[j]->dynCast<LocalSet>()->value->is<StructNew>()) {
+      // Don't swap two struct.new instructions to avoid going back and forth.
+      return false;
+    }
+    // Check if the two expressions can be swapped safely considering their
+    // effects.
+    auto firstEffects = effects(list[i]);
+    auto secondEffects = effects(list[j]);
+    if (secondEffects.invalidates(firstEffects)) {
+      return false;
+    }
+
+    std::swap(list[i], list[j]);
+    return true;
+  }
+
   // Given a struct.new and a struct.set that occurs right after it, and that
   // applies to the same data, try to apply the set during the new. This can be
   // either with a nested tee:
diff --git a/test/lit/passes/optimize-instructions-gc-heap.wast b/test/lit/passes/optimize-instructions-gc-heap.wast
index 6db63fae3..2988b1bd8 100644
--- a/test/lit/passes/optimize-instructions-gc-heap.wast
+++ b/test/lit/passes/optimize-instructions-gc-heap.wast
@@ -9,11 +9,10 @@
   ;; CHECK:      (type $struct (struct (field (mut i32))))
   (type $struct (struct (field (mut i32))))
 
-  ;; CHECK:      (type $struct3 (struct (field (mut i32)) (field (mut i32)) (field (mut i32))))
-
   ;; CHECK:      (type $struct2 (struct (field (mut i32)) (field (mut i32))))
   (type $struct2 (struct (field (mut i32)) (field (mut i32))))
 
+  ;; CHECK:      (type $struct3 (struct (field (mut i32)) (field (mut i32)) (field (mut i32))))
   (type $struct3 (struct (field (mut i32)) (field (mut i32)) (field (mut i32))))
 
   ;; CHECK:      (func $tee (type $1)
@@ -272,12 +271,15 @@
 
   ;; CHECK:      (func $pattern-breaker (type $1)
   ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (local $ref2 (ref null $struct))
   ;; CHECK-NEXT:  (local.set $ref
   ;; CHECK-NEXT:   (struct.new $struct
   ;; CHECK-NEXT:    (i32.const 10)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
-  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT:  (local.set $ref2
+  ;; CHECK-NEXT:   (local.get $ref)
+  ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (struct.set $struct 0
   ;; CHECK-NEXT:   (local.get $ref)
   ;; CHECK-NEXT:   (i32.const 20)
@@ -285,13 +287,56 @@
   ;; CHECK-NEXT: )
   (func $pattern-breaker
     (local $ref (ref null $struct))
+    (local $ref2 (ref null $struct))
+    (local.set $ref
+      (struct.new $struct
+        (i32.const 10)
+      )
+    )
+    ;; Any instruction that can not be swapped and is not
+    ;; the expected struct.set breaks the pattern.
+    (local.set $ref2 (local.get $ref))
+    (struct.set $struct 0
+      (local.get $ref)
+      (i32.const 20)
+    )
+  )
+
+  ;; CHECK:      (func $dont-swap-subsequent-struct-new (type $1)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (local $ref2 (ref null $struct))
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT:  (local.set $ref
+  ;; CHECK-NEXT:   (struct.new $struct
+  ;; CHECK-NEXT:    (i32.const 10)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (local.set $ref2
+  ;; CHECK-NEXT:   (struct.new $struct
+  ;; CHECK-NEXT:    (i32.const 20)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.set $struct 0
+  ;; CHECK-NEXT:   (local.get $ref)
+  ;; CHECK-NEXT:   (i32.const 20)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $dont-swap-subsequent-struct-new
+    (local $ref (ref null $struct))
+    (local $ref2 (ref null $struct))
     (local.set $ref
       (struct.new $struct
         (i32.const 10)
       )
     )
-    ;; Anything that we don't recognize breaks the pattern.
     (nop)
+    ;; We do not swap with another local.set of struct.new.
+    (local.set $ref2
+      (struct.new $struct
+        (i32.const 20)
+      )
+    )
+    ;; last instruction in the block won't be swapped.
     (struct.set $struct 0
       (local.get $ref)
       (i32.const 20)
@@ -611,20 +656,17 @@
   ;; CHECK:      (func $many-news (type $1)
   ;; CHECK-NEXT:  (local $ref (ref null $struct3))
   ;; CHECK-NEXT:  (local $ref2 (ref null $struct3))
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT:  (nop)
   ;; CHECK-NEXT:  (local.set $ref
   ;; CHECK-NEXT:   (struct.new $struct3
   ;; CHECK-NEXT:    (i32.const 40)
   ;; CHECK-NEXT:    (i32.const 50)
-  ;; CHECK-NEXT:    (i32.const 30)
+  ;; CHECK-NEXT:    (i32.const 60)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
-  ;; CHECK-NEXT:  (nop)
-  ;; CHECK-NEXT:  (nop)
-  ;; CHECK-NEXT:  (struct.set $struct3 2
-  ;; CHECK-NEXT:   (local.get $ref)
-  ;; CHECK-NEXT:   (i32.const 60)
-  ;; CHECK-NEXT:  )
-  ;; CHECK-NEXT:  (nop)
   ;; CHECK-NEXT:  (local.set $ref
   ;; CHECK-NEXT:   (struct.new $struct3
   ;; CHECK-NEXT:    (i32.const 400)