HeapStoreOptimization: Fix a bug with jumping from the later value (v2) (#7070)

This PR fixes this situation:

(block $out
  (local.set $x (struct.new X Y Z))
  (struct.set $X 0 (local.get $x) (..br $out..))  ;; X' here has a br
)
(local.get $x)

=>

(block $out
  (local.set $x (struct.new (..br $out..) Y Z))
)
(local.get $x)

We want to fold the struct.set into the struct.new, but the br is
a problem: if it executes then we skip the struct.set, and the last
local.get in fact reads the struct before the write. And, if we did this
optimization, we'd end up with the br on the struct.new, so it
would skip that instruction and even the local.set.

To fix this, we use the new API from #7039, which lets us query,
"is it ok to move the local.set to where the struct.set is?"

2 files changed, 873 insertions, 19 deletions

diff --git a/src/passes/HeapStoreOptimization.cpp b/src/passes/HeapStoreOptimization.cpp
index f227dc536..605caba41 100644
--- a/src/passes/HeapStoreOptimization.cpp
+++ b/src/passes/HeapStoreOptimization.cpp
@@ -22,6 +22,7 @@
 
 #include "cfg/cfg-traversal.h"
 #include "ir/effects.h"
+#include "ir/local-graph.h"
 #include "pass.h"
 #include "wasm-builder.h"
 #include "wasm.h"
@@ -90,7 +91,7 @@ struct HeapStoreOptimization
     //
     if (auto* tee = curr->ref->dynCast<LocalSet>()) {
       if (auto* new_ = tee->value->dynCast<StructNew>()) {
-        if (optimizeSubsequentStructSet(new_, curr, tee->index)) {
+        if (optimizeSubsequentStructSet(new_, curr, tee)) {
           // Success, so we do not need the struct.set any more, and the tee
           // can just be a set instead of us.
           tee->makeSet();
@@ -147,7 +148,7 @@ struct HeapStoreOptimization
         }
 
         // The pattern matches, try to optimize.
-        if (!optimizeSubsequentStructSet(new_, structSet, localGet->index)) {
+        if (!optimizeSubsequentStructSet(new_, structSet, localSet)) {
           break;
         } else {
           // Success. Replace the set with a nop, and continue to perhaps
@@ -209,7 +210,7 @@ struct HeapStoreOptimization
   // Returns true if we succeeded.
   bool optimizeSubsequentStructSet(StructNew* new_,
                                    StructSet* set,
-                                   Index refLocalIndex) {
+                                   LocalSet* localSet) {
     // Leave unreachable code for DCE, to avoid updating types here.
     if (new_->type == Type::unreachable || set->type == Type::unreachable) {
       return false;
@@ -217,6 +218,7 @@ struct HeapStoreOptimization
 
     auto index = set->index;
     auto& operands = new_->operands;
+    auto refLocalIndex = localSet->index;
 
     // Check for effects that prevent us moving the struct.set's value (X' in
     // the function comment) into its new position in the struct.new. First, it
@@ -246,6 +248,12 @@ struct HeapStoreOptimization
       }
     }
 
+    // We must also be careful of branches out from the value that skip the
+    // local.set, see below.
+    if (canSkipLocalSet(set, setValueEffects, localSet)) {
+      return false;
+    }
+
     // We can optimize here!
     Builder builder(*getModule());
 
@@ -271,9 +279,74 @@ struct HeapStoreOptimization
     return true;
   }
 
+  // We must be careful of branches that skip the local.set. For example:
+  //
+  //  (block $out
+  //    (local.set $x (struct.new X Y Z))
+  //    (struct.set (local.get $x) (..br $out..))  ;; X' here has a br
+  //  )
+  //  ..local.get $x..
+  //
+  // Note how we do the local.set first. Imagine we optimized to this:
+  //
+  //  (block $out
+  //    (local.set $x (struct.new (..br $out..) Y Z))
+  //  )
+  //  ..local.get $x..
+  //
+  // Now the br happens first, skipping the local.set entirely, and the use
+  // later down will not get the proper value. This is the problem we must
+  // detect here.
+  //
+  // We are given a struct.set, the computed effects of its value (the caller
+  // already has those, so this is an optimization to avoid recomputation), and
+  // the local.set.
+  bool canSkipLocalSet(StructSet* set,
+                       const EffectAnalyzer& setValueEffects,
+                       LocalSet* localSet) {
+    // First, if the set's value cannot branch at all, then we have no problem.
+    if (!setValueEffects.transfersControlFlow()) {
+      return false;
+    }
+
+    // We may branch, so do an analysis using a LocalGraph. We will check
+    // whether it is valid to move the local.set to where the struct.set is, so
+    // we provide StructSet as the query class.
+    //
+    // It is valid to reuse the LocalGraph many times because the optimization
+    // that we do in this pass does not generate new, dangerous control flow. We
+    // only optimize if moving a LocalSet is valid, and that does not invalidate
+    // any other one.
+    if (!localGraph) {
+      localGraph.emplace(getFunction(), getModule(), StructSet::SpecificId);
+    }
+    auto badGets = localGraph->canMoveSet(localSet, set);
+    if (badGets.size() == 0) {
+      // No problems at all.
+      return false;
+    }
+    // It is ok to have a local.get in the struct.set itself, as if we optimize
+    // then that local.get goes away anyhow, that is,
+    //
+    //   (local.set $x (struct.new X Y Z))
+    //   (struct.set (local.get $x) (X'))
+    //  =>
+    //   (local.set $x (struct.new X' Y Z))
+    //
+    // Both the local.get and the struct.set are removed.
+    if (badGets.size() >= 2) {
+      return true;
+    }
+    return *badGets.begin() != set->ref;
+  }
+
   EffectAnalyzer effects(Expression* expr) {
     return EffectAnalyzer(getPassOptions(), *getModule(), expr);
   }
+
+private:
+  // A local graph that is constructed the first time we need it.
+  std::optional<LazyLocalGraph> localGraph;
 };
 
 } // anonymous namespace
diff --git a/test/lit/passes/heap-store-optimization.wast b/test/lit/passes/heap-store-optimization.wast
index 635963523..1c1abba4a 100644
--- a/test/lit/passes/heap-store-optimization.wast
+++ b/test/lit/passes/heap-store-optimization.wast
@@ -7,12 +7,18 @@
 
 (module
   ;; CHECK:      (type $struct (struct (field (mut i32))))
-  (type $struct (struct (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))))
+
+  ;; CHECK:      (tag $tag)
+  (tag $tag)
+
+  (type $struct (struct (field (mut i32))))
+
+  (type $struct2 (struct (field (mut i32)) (field (mut i32))))
+
   (type $struct3 (struct (field (mut i32)) (field (mut i32)) (field (mut i32))))
 
   ;; CHECK:      (func $tee (type $1)
@@ -816,26 +822,29 @@
     )
   )
 
-  ;; CHECK:      (func $helper-i32 (type $4) (param $x i32) (result i32)
+  ;; CHECK:      (func $helper-i32 (type $7) (param $x i32) (result i32)
   ;; CHECK-NEXT:  (i32.const 42)
   ;; CHECK-NEXT: )
   (func $helper-i32 (param $x i32) (result i32)
     (i32.const 42)
   )
 
-  ;; CHECK:      (func $control-flow-in-set-value (type $5) (result i32)
+  ;; CHECK:      (func $control-flow-in-set-value (type $4) (result i32)
   ;; CHECK-NEXT:  (local $ref (ref null $struct))
   ;; CHECK-NEXT:  (block $label
-  ;; CHECK-NEXT:   (local.set $ref
-  ;; CHECK-NEXT:    (struct.new $struct
-  ;; CHECK-NEXT:     (if (result i32)
+  ;; CHECK-NEXT:   (struct.set $struct 0
+  ;; CHECK-NEXT:    (local.tee $ref
+  ;; CHECK-NEXT:     (struct.new $struct
   ;; CHECK-NEXT:      (i32.const 1)
-  ;; CHECK-NEXT:      (then
-  ;; CHECK-NEXT:       (br $label)
-  ;; CHECK-NEXT:      )
-  ;; CHECK-NEXT:      (else
-  ;; CHECK-NEXT:       (i32.const 42)
-  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (if (result i32)
+  ;; CHECK-NEXT:     (i32.const 1)
+  ;; CHECK-NEXT:     (then
+  ;; CHECK-NEXT:      (br $label)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (else
+  ;; CHECK-NEXT:      (i32.const 42)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
@@ -846,8 +855,7 @@
   ;; CHECK-NEXT: )
   (func $control-flow-in-set-value (result i32)
     ;; Test we handle control flow in the struct.set's value when we combine a
-    ;; struct.set with a struct.new.
-    ;; XXX The output above is wrong atm, and the pass needs fixing!
+    ;; struct.set with a struct.new. We should not optimize here.
     (local $ref (ref null $struct))
     (block $label
       (struct.set $struct 0
@@ -876,5 +884,778 @@
       (local.get $ref)
     )
   )
-)
 
+  ;; CHECK:      (func $control-flow-in-set-value-safe (type $4) (result i32)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (block
+  ;; CHECK-NEXT:   (local.set $ref
+  ;; CHECK-NEXT:    (struct.new $struct
+  ;; CHECK-NEXT:     (if (result i32)
+  ;; CHECK-NEXT:      (i32.const 1)
+  ;; CHECK-NEXT:      (then
+  ;; CHECK-NEXT:       (i32.const 1337)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (else
+  ;; CHECK-NEXT:       (i32.const 42)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.get $struct 0
+  ;; CHECK-NEXT:   (local.get $ref)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $control-flow-in-set-value-safe (result i32)
+    ;; As above, but now the control flow in the value is safe: it does not
+    ;; escape out. We should optimize here.
+    (local $ref (ref null $struct))
+    (block $label
+      (struct.set $struct 0
+        (local.tee $ref
+          (struct.new $struct
+            (i32.const 1)
+          )
+        )
+        (if (result i32)
+          (i32.const 1)
+          (then
+            (i32.const 1337)  ;; the break to $out was replaced
+          )
+          (else
+            (i32.const 42)
+          )
+        )
+      )
+    )
+    (struct.get $struct 0
+      (local.get $ref)
+    )
+  )
+
+  ;; CHECK:      (func $control-flow-in-set-value-safe-call (type $4) (result i32)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (block
+  ;; CHECK-NEXT:   (local.set $ref
+  ;; CHECK-NEXT:    (struct.new $struct
+  ;; CHECK-NEXT:     (call $helper-i32
+  ;; CHECK-NEXT:      (i32.const 42)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.get $struct 0
+  ;; CHECK-NEXT:   (local.get $ref)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $control-flow-in-set-value-safe-call (result i32)
+    ;; As above, but now the possible control flow is a call. It may throw, but
+    ;; that would go outside of the function, which is fine.
+    (local $ref (ref null $struct))
+    (block $label
+      (struct.set $struct 0
+        (local.tee $ref
+          (struct.new $struct
+            (i32.const 1)
+          )
+        )
+        (call $helper-i32 (i32.const 42))  ;; the if was replaced by this call
+      )
+    )
+    (struct.get $struct 0
+      (local.get $ref)
+    )
+  )
+
+  (func $control-flow-in-set-value-safe-return (result i32)
+    ;; As above, but replace the call with a return in an if. We can still
+    ;; optimize (if the return is taken, we go outside of the function anyhow).
+    (local $ref (ref null $struct))
+    (block $label
+      (struct.set $struct 0
+        (local.tee $ref
+          (struct.new $struct
+            (i32.const 1)
+          )
+        )
+        (if (result i32)
+          (i32.const 1)
+          (then
+            (return (i32.const 42))
+          )
+          (else
+            (i32.const 42)
+          )
+        )
+      )
+    )
+    (struct.get $struct 0
+      (local.get $ref)
+    )
+  )
+
+  ;; CHECK:      (func $control-flow-in-set-value-unsafe-call (type $4) (result i32)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (block $label
+  ;; CHECK-NEXT:   (try_table (catch $tag $label)
+  ;; CHECK-NEXT:    (struct.set $struct 0
+  ;; CHECK-NEXT:     (local.tee $ref
+  ;; CHECK-NEXT:      (struct.new $struct
+  ;; CHECK-NEXT:       (i32.const 1)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (call $helper-i32
+  ;; CHECK-NEXT:      (i32.const 42)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.get $struct 0
+  ;; CHECK-NEXT:   (local.get $ref)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $control-flow-in-set-value-unsafe-call (result i32)
+    ;; As above, but now the call's possible throw could be caught *inside* the
+    ;; function, which means it is dangerous, and we do not optimize.
+    (local $ref (ref null $struct))
+    (block $label
+      (try_table (catch $tag $label)  ;; this try was added
+        (struct.set $struct 0
+          (local.tee $ref
+            (struct.new $struct
+              (i32.const 1)
+            )
+          )
+          (call $helper-i32 (i32.const 42))
+        )
+      )
+    )
+    (struct.get $struct 0
+      (local.get $ref)
+    )
+  )
+
+  ;; CHECK:      (func $control-flow-later (type $5) (param $x i32)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (block $out
+  ;; CHECK-NEXT:   (br_if $out
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (local.set $ref
+  ;; CHECK-NEXT:   (struct.new $struct
+  ;; CHECK-NEXT:    (i32.const 42)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (local.get $x)
+  ;; CHECK-NEXT:   (then
+  ;; CHECK-NEXT:    (nop)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $control-flow-later (param $x i32)
+    (local $ref (ref null $struct))
+    (struct.set $struct 0
+      (local.tee $ref
+        (struct.new $struct
+          (i32.const 1)
+        )
+      )
+      (i32.const 42)
+    )
+    ;; This later control flow should not prevent optimizing the struct.set
+    ;; before it.
+    (block $out
+      (br_if $out
+        (local.get $x)
+      )
+    )
+    (if
+      (local.get $x)
+      (then
+        (nop)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $loop (type $1)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (loop $loop
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (struct.get $struct 0
+  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (struct.set $struct 0
+  ;; CHECK-NEXT:    (local.tee $ref
+  ;; CHECK-NEXT:     (struct.new $struct
+  ;; CHECK-NEXT:      (i32.const 1)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (if (result i32)
+  ;; CHECK-NEXT:     (i32.const 1)
+  ;; CHECK-NEXT:     (then
+  ;; CHECK-NEXT:      (br $loop)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (else
+  ;; CHECK-NEXT:      (i32.const 42)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $loop
+    (local $ref (ref null $struct))
+    (loop $loop
+      ;; There is a use of the reference at the top of the loop, and the br_if
+      ;; may get here, so this is a basic block before the struct.set that we
+      ;; need to be careful of reaching. We should not optimize here.
+      (drop
+        (struct.get $struct 0
+          (local.get $ref)
+        )
+      )
+      (struct.set $struct 0
+        (local.tee $ref
+          (struct.new $struct
+            (i32.const 1)
+          )
+        )
+        (if (result i32)
+          (i32.const 1)
+          (then
+            (br $loop)
+          )
+          (else
+            (i32.const 42)
+          )
+        )
+      )
+    )
+  )
+
+  ;; CHECK:      (func $loop-more-flow (type $1)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (loop $loop
+  ;; CHECK-NEXT:   (if
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:    (then
+  ;; CHECK-NEXT:     (br $loop)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (struct.get $struct 0
+  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (struct.set $struct 0
+  ;; CHECK-NEXT:    (local.tee $ref
+  ;; CHECK-NEXT:     (struct.new $struct
+  ;; CHECK-NEXT:      (i32.const 1)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (if (result i32)
+  ;; CHECK-NEXT:     (i32.const 1)
+  ;; CHECK-NEXT:     (then
+  ;; CHECK-NEXT:      (br $loop)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (else
+  ;; CHECK-NEXT:      (i32.const 42)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $loop-more-flow
+    (local $ref (ref null $struct))
+    (loop $loop
+      ;; As above, but add this if which adds more control flow at the loop top.
+      ;; We should still not optimize here.
+      (if
+        (i32.const 1)
+        (then
+          (br $loop)
+        )
+      )
+      (drop
+        (struct.get $struct 0
+          (local.get $ref)
+        )
+      )
+      (struct.set $struct 0
+        (local.tee $ref
+          (struct.new $struct
+            (i32.const 1)
+          )
+        )
+        (if (result i32)
+          (i32.const 1)
+          (then
+            (br $loop)
+          )
+          (else
+            (i32.const 42)
+          )
+        )
+      )
+    )
+  )
+
+  ;; CHECK:      (func $loop-in-value (type $5) (param $x i32)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (local.set $ref
+  ;; CHECK-NEXT:   (struct.new $struct
+  ;; CHECK-NEXT:    (loop $loop (result i32)
+  ;; CHECK-NEXT:     (br_if $loop
+  ;; CHECK-NEXT:      (local.get $x)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (i32.const 42)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $loop-in-value (param $x i32)
+    (local $ref (ref null $struct))
+    ;; The struct.set's value has a loop in it, but that is fine, as while there
+    ;; are backedges there, they are still contained in the value: we can't skip
+    ;; the struct.set. We can optimize here.
+    (struct.set $struct 0
+      (local.tee $ref
+        (struct.new $struct
+          (i32.const 1)
+        )
+      )
+      (loop $loop (result i32)
+        (br_if $loop
+          (local.get $x)
+        )
+        (i32.const 42)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $in-if-arm (type $6) (param $x i32) (param $y i32) (result i32)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (local.get $x)
+  ;; CHECK-NEXT:   (then
+  ;; CHECK-NEXT:    (block $out
+  ;; CHECK-NEXT:     (struct.set $struct 0
+  ;; CHECK-NEXT:      (local.tee $ref
+  ;; CHECK-NEXT:       (struct.new $struct
+  ;; CHECK-NEXT:        (i32.const 1)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (block (result i32)
+  ;; CHECK-NEXT:       (br_if $out
+  ;; CHECK-NEXT:        (local.get $x)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (i32.const 42)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.get $struct 0
+  ;; CHECK-NEXT:   (local.get $ref)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $in-if-arm (param $x i32) (param $y i32) (result i32)
+    (local $ref (ref null $struct))
+    (if
+      (local.get $x)
+      (then
+        (block $out
+          ;; We cannot optimize here, as the struct.get outside of the if can
+          ;; read different state if the br_if happens.
+          (struct.set $struct 0
+            (local.tee $ref
+              (struct.new $struct
+                (i32.const 1)
+              )
+            )
+            (block (result i32)
+              (br_if $out
+                (local.get $x)
+              )
+              (i32.const 42)
+            )
+          )
+        )
+      )
+    )
+    (struct.get $struct 0
+      (local.get $ref)
+    )
+  )
+
+  ;; CHECK:      (func $in-if-arm-yes (type $6) (param $x i32) (param $y i32) (result i32)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (local.get $x)
+  ;; CHECK-NEXT:   (then
+  ;; CHECK-NEXT:    (block $out
+  ;; CHECK-NEXT:     (local.set $ref
+  ;; CHECK-NEXT:      (struct.new $struct
+  ;; CHECK-NEXT:       (block (result i32)
+  ;; CHECK-NEXT:        (br_if $out
+  ;; CHECK-NEXT:         (local.get $x)
+  ;; CHECK-NEXT:        )
+  ;; CHECK-NEXT:        (i32.const 42)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (i32.const 1337)
+  ;; CHECK-NEXT: )
+  (func $in-if-arm-yes (param $x i32) (param $y i32) (result i32)
+    (local $ref (ref null $struct))
+    ;; As before, but the struct.get at the end is removed, so we can optimize.
+    (if
+      (local.get $x)
+      (then
+        (block $out
+          (struct.set $struct 0
+            (local.tee $ref
+              (struct.new $struct
+                (i32.const 1)
+              )
+            )
+            (block (result i32)
+              (br_if $out
+                (local.get $x)
+              )
+              (i32.const 42)
+            )
+          )
+        )
+      )
+    )
+    (i32.const 1337) ;; this changed
+  )
+
+  ;; CHECK:      (func $control-flow-in-set-value-sequence (type $4) (result i32)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block $out (result i32)
+  ;; CHECK-NEXT:    (local.set $ref
+  ;; CHECK-NEXT:     (struct.new_default $struct)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (struct.set $struct 0
+  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:     (br_if $out
+  ;; CHECK-NEXT:      (i32.const 1)
+  ;; CHECK-NEXT:      (i32.const 2)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (struct.set $struct 0
+  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:     (i32.const 3)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (struct.get $struct 0
+  ;; CHECK-NEXT:      (local.get $ref)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 42)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.get $struct 0
+  ;; CHECK-NEXT:   (local.get $ref)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $control-flow-in-set-value-sequence (result i32)
+    (local $ref (ref null $struct))
+    (drop
+      (block $out (result i32)
+        (local.set $ref
+          ;; Also test struct.new_default here, with control flow.
+          (struct.new_default $struct)
+        )
+        ;; The struct.get outside is a problem, so we do not optimize here,
+        ;; nor the set after us.
+        (struct.set $struct 0
+          (local.get $ref)
+          (br_if $out
+            (i32.const 1)
+            (i32.const 2)
+          )
+        )
+        (struct.set $struct 0
+          (local.get $ref)
+          (i32.const 3)
+        )
+        ;; This struct.set is not a problem: if we branch, we don't reach it
+        ;; anyhow.
+        (drop
+          (struct.get $struct 0
+            (local.get $ref)
+          )
+        )
+        (i32.const 42)
+      )
+    )
+    (struct.get $struct 0
+      (local.get $ref)
+    )
+  )
+
+  ;; CHECK:      (func $control-flow-in-set-value-sequence-2 (type $4) (result i32)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block $out (result i32)
+  ;; CHECK-NEXT:    (local.set $ref
+  ;; CHECK-NEXT:     (struct.new $struct
+  ;; CHECK-NEXT:      (i32.const 3)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (nop)
+  ;; CHECK-NEXT:    (struct.set $struct 0
+  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:     (br_if $out
+  ;; CHECK-NEXT:      (i32.const 1)
+  ;; CHECK-NEXT:      (i32.const 2)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (struct.get $struct 0
+  ;; CHECK-NEXT:      (local.get $ref)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 42)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.get $struct 0
+  ;; CHECK-NEXT:   (local.get $ref)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $control-flow-in-set-value-sequence-2 (result i32)
+    (local $ref (ref null $struct))
+    (drop
+      (block $out (result i32)
+        (local.set $ref
+          (struct.new_default $struct)
+        )
+        ;; As above, but the order of struct.sets is flipped. We can at least
+        ;; optimize the first here.
+        (struct.set $struct 0
+          (local.get $ref)
+          (i32.const 3)
+        )
+        (struct.set $struct 0
+          (local.get $ref)
+          (br_if $out
+            (i32.const 1)
+            (i32.const 2)
+          )
+        )
+        (drop
+          (struct.get $struct 0
+            (local.get $ref)
+          )
+        )
+        (i32.const 42)
+      )
+    )
+    (struct.get $struct 0
+      (local.get $ref)
+    )
+  )
+
+  ;; CHECK:      (func $control-flow-in-set-value-sequence-yes (type $4) (result i32)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block $out (result i32)
+  ;; CHECK-NEXT:    (local.set $ref
+  ;; CHECK-NEXT:     (struct.new $struct
+  ;; CHECK-NEXT:      (br_if $out
+  ;; CHECK-NEXT:       (i32.const 1)
+  ;; CHECK-NEXT:       (i32.const 2)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (nop)
+  ;; CHECK-NEXT:    (nop)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (struct.get $struct 0
+  ;; CHECK-NEXT:      (local.get $ref)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 42)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (i32.const 1337)
+  ;; CHECK-NEXT: )
+  (func $control-flow-in-set-value-sequence-yes (result i32)
+    (local $ref (ref null $struct))
+    ;; As above, but the struct.get at the end is removed, allowing us to
+    ;; optimize it all.
+    (drop
+      (block $out (result i32)
+        (local.set $ref
+          (struct.new_default $struct)
+        )
+        (struct.set $struct 0
+          (local.get $ref)
+          (i32.const 3)
+        )
+        (struct.set $struct 0
+          (local.get $ref)
+          (br_if $out
+            (i32.const 1)
+            (i32.const 2)
+          )
+        )
+        ;; Note how this struct.get remains and does not stop us.
+        (drop
+          (struct.get $struct 0
+            (local.get $ref)
+          )
+        )
+        (i32.const 42)
+      )
+    )
+    (i32.const 1337)
+  )
+
+  ;; CHECK:      (func $multi-control-flow-in-set-value-sequence-yes (type $4) (result i32)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block $out (result i32)
+  ;; CHECK-NEXT:    (local.set $ref
+  ;; CHECK-NEXT:     (struct.new $struct
+  ;; CHECK-NEXT:      (br_if $out
+  ;; CHECK-NEXT:       (i32.const 1)
+  ;; CHECK-NEXT:       (i32.const 2)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (nop)
+  ;; CHECK-NEXT:    (struct.set $struct 0
+  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:     (br_if $out
+  ;; CHECK-NEXT:      (i32.const 3)
+  ;; CHECK-NEXT:      (i32.const 4)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (struct.set $struct 0
+  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:     (br_if $out
+  ;; CHECK-NEXT:      (i32.const 5)
+  ;; CHECK-NEXT:      (i32.const 6)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 42)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (i32.const 1337)
+  ;; CHECK-NEXT: )
+  (func $multi-control-flow-in-set-value-sequence-yes (result i32)
+    (local $ref (ref null $struct))
+    ;; As above, but now we have multiple br_ifs. We optimize one, but then
+    ;; stop because of the control flow that is now in the struct.new. TODO we
+    ;; could be more precise here.
+    (drop
+      (block $out (result i32)
+        (local.set $ref
+          (struct.new_default $struct)
+        )
+        (struct.set $struct 0
+          (local.get $ref)
+          (br_if $out
+            (i32.const 1)
+            (i32.const 2)
+          )
+        )
+        (struct.set $struct 0
+          (local.get $ref)
+          (br_if $out
+            (i32.const 3)
+            (i32.const 4)
+          )
+        )
+        (struct.set $struct 0
+          (local.get $ref)
+          (br_if $out
+            (i32.const 5)
+            (i32.const 6)
+          )
+        )
+        (i32.const 42)
+      )
+    )
+    (i32.const 1337)
+  )
+
+  ;; CHECK:      (func $multi-control-flow-in-set-value-sequence-no (type $8) (result anyref)
+  ;; CHECK-NEXT:  (local $ref (ref null $struct))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block $out (result i32)
+  ;; CHECK-NEXT:    (local.set $ref
+  ;; CHECK-NEXT:     (struct.new_default $struct)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (struct.set $struct 0
+  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:     (br_if $out
+  ;; CHECK-NEXT:      (i32.const 1)
+  ;; CHECK-NEXT:      (i32.const 2)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (struct.set $struct 0
+  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:     (br_if $out
+  ;; CHECK-NEXT:      (i32.const 3)
+  ;; CHECK-NEXT:      (i32.const 4)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (struct.set $struct 0
+  ;; CHECK-NEXT:     (local.get $ref)
+  ;; CHECK-NEXT:     (br_if $out
+  ;; CHECK-NEXT:      (i32.const 5)
+  ;; CHECK-NEXT:      (i32.const 6)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (i32.const 42)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (local.get $ref)
+  ;; CHECK-NEXT: )
+  (func $multi-control-flow-in-set-value-sequence-no (result anyref)
+    (local $ref (ref null $struct))
+    ;; As above, but now we have a dangerous local.get at the end, stopping us
+    ;; from optimizing.
+    (drop
+      (block $out (result i32)
+        (local.set $ref
+          (struct.new_default $struct)
+        )
+        (struct.set $struct 0
+          (local.get $ref)
+          (br_if $out
+            (i32.const 1)
+            (i32.const 2)
+          )
+        )
+        (struct.set $struct 0
+          (local.get $ref)
+          (br_if $out
+            (i32.const 3)
+            (i32.const 4)
+          )
+        )
+        (struct.set $struct 0
+          (local.get $ref)
+          (br_if $out
+            (i32.const 5)
+            (i32.const 6)
+          )
+        )
+        (i32.const 42)
+      )
+    )
+    (local.get $ref)
+  )
+)