[LICM] Defer replacing hoisted code with nop (#9326)

LoopInvariantCodeMotion uses a `LazyLocalGraph` that delays walking the
function and building its CFG until the first `getSets` call.
Previously, `visitLoop` immediately replaced hoisted expressions with
nops before later inserting them above the loop. When an expression that
wrote to a local without reading any locals was hoisted first (bypassing
`hasGetDependingOnLoopSet`), and a subsequent expression in the same
loop read a local and triggered the first `LazyLocalGraph::getSets`
call, `LocalGraphFlower` walked the function while the hoisted
expression was detached from the AST. As a result, `LocalGraph` missed
the `local.set` inside the detached expression and later misclassified
expressions in enclosing loops that read that local as loop-invariant.

Record pointers to the hoisted expression slots in `movedCode` and only
replace them with nops when constructing the replacement block at the
end of `visitLoop`, keeping the AST intact whenever `LazyLocalGraph`
initializes.

Author: tlively

src/passes/LoopInvariantCodeMotion.cpp

@@ -60,7 +60,7 @@ struct LoopInvariantCodeMotion
   void visitLoop(Loop* loop) {
     // We accumulate all the code we can move out, and will place it
     // in a block just preceding the loop.
-    std::vector<Expression*> movedCode;
+    std::vector<Expression**> movedCode;
     // Accumulate effects of things we can't move out - things
     // we move out later must cross them, so we must verify it
     // is ok to do so.
@@ -172,10 +172,11 @@ struct LoopInvariantCodeMotion
                 numSetsForIndex[set->index]++;
               }
             } else {
-              // We can move it! Leave the changes, move the code, and update
-              // loopSets.
-              movedCode.push_back(curr);
-              *currp = Builder(*getModule()).makeNop();
+              // We can move it! Leave the changes, record the code to move, and
+              // update loopSets. Do not replace *currp with a nop yet, as that
+              // would detach the expression from the function body before a
+              // potential lazy initialization of localGraph.
+              movedCode.push_back(currp);
               for (auto* set : currSets.list) {
                 loopSets.erase(set);
               }
@@ -195,10 +196,14 @@ struct LoopInvariantCodeMotion
     if (!movedCode.empty()) {
       // Finish the moving by emitting the code outside.
       Builder builder(*getModule());
-      auto* ret = builder.makeBlock(movedCode);
-      ret->list.push_back(loop);
-      ret->finalize(loop->type);
-      replaceCurrent(ret);
+      std::vector<Expression*> exprs;
+      exprs.reserve(movedCode.size() + 1);
+      for (auto** currp : movedCode) {
+        exprs.push_back(*currp);
+        *currp = builder.makeNop();
+      }
+      exprs.push_back(loop);
+      replaceCurrent(builder.makeBlock(exprs, loop->type));
       // Note that we do not need to modify the localGraph - we keep
       // each get in a position to be influenced by exactly the same
       // sets as before.

test/lit/passes/licm.wast

@@ -186,6 +186,79 @@
     (local.get $y)
   )
 
+  ;; CHECK:      (func $nested-loop-lazy-local-graph (type $1) (param $cond i32) (result i32)
+  ;; CHECK-NEXT:  (local $x i32)
+  ;; CHECK-NEXT:  (local $y i32)
+  ;; CHECK-NEXT:  (local $z i32)
+  ;; CHECK-NEXT:  (loop $outer
+  ;; CHECK-NEXT:   (local.set $y
+  ;; CHECK-NEXT:    (i32.add
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:     (i32.const 10)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (block
+  ;; CHECK-NEXT:    (if
+  ;; CHECK-NEXT:     (i32.const 1)
+  ;; CHECK-NEXT:     (then
+  ;; CHECK-NEXT:      (local.set $x
+  ;; CHECK-NEXT:       (i32.const 1)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $z
+  ;; CHECK-NEXT:     (i32.eqz
+  ;; CHECK-NEXT:      (local.get $x)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (loop $inner
+  ;; CHECK-NEXT:     (nop)
+  ;; CHECK-NEXT:     (nop)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (br_if $outer
+  ;; CHECK-NEXT:    (local.get $cond)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (local.get $y)
+  ;; CHECK-NEXT: )
+  (func $nested-loop-lazy-local-graph (param $cond i32) (result i32)
+    (local $x i32)
+    (local $y i32)
+    (local $z i32)
+    ;; Hoisting the (if ...) out of $inner must not detach its (local.set $x)
+    ;; from the function before LazyLocalGraph initializes when checking the
+    ;; subsequent (local.get $x) in $inner. Otherwise, $outer would miss the
+    ;; write to $x and invalidly hoist (local.set $y ...).
+    (loop $outer
+      (local.set $y
+        (i32.add
+          (local.get $x)
+          (i32.const 10)
+        )
+      )
+      (loop $inner
+        (if
+          (i32.const 1)
+          (then
+            (local.set $x
+              (i32.const 1)
+            )
+          )
+        )
+        (local.set $z
+          (i32.eqz
+            (local.get $x)
+          )
+        )
+      )
+      (br_if $outer
+        (local.get $cond)
+      )
+    )
+    (local.get $y)
+  )
+
   ;; CHECK:      (func $hoist-trap-before-backedge (type $1) (param $x i32) (result i32)
   ;; CHECK-NEXT:  (local $y i32)
   ;; CHECK-NEXT:  (block