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.

1 files changed, 55 insertions, 17 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: