[NFC] Refactor ChildLocalizer to handle unreachable code better (WebA…

…ssembly#6394)

This is NFC in the current users, but is necessary functionality for a later
PR.

ChildLocalizer moves children into locals as needed. It used to stop when it
saw the first unreachable. After this change we move such unreachable
children out of the parent as well, making this more uniform: all interacting
effects are moved out, and all that is left nested in the parent can be
moved around and removed as desired.

Also add a getReplacement helper that makes using this easier.

This cannot be tested comprehensively with the current user as that user
will not call this code path on an unreachable parent at all, so this just
adds what can be tested. The later PR will have tests for all corner cases.

src/ir/localize.h

@@ -46,28 +46,50 @@ struct Localizer {
 
 // Replaces all children with gets of locals, if they have any effects that
 // interact with any of the others, or if they have side effects which cannot be
-// removed.
+// removed. Also replace unreachable things with an unreachable, leaving in
+// place only things without interacting effects. For example:
 //
-// After this, the original input has only local.gets as inputs, or other things
-// that have no interacting effects, and so those children can be reordered
-// and/or removed as needed.
+//  (parent
+//    (call $foo)
+//    (br $out)
+//    (i32.const)
+//  )
 //
-// The sets of the locals are emitted on a |sets| property on the class. Those
-// must be emitted right before the input.
+// =>
 //
-// This stops at the first unreachable child, as there is no code executing
-// after that point anyhow.
+//  (local.set $temp.foo
+//    (call $foo)            ;; moved out
+//  )
+//  (br $out)                ;; moved out
+//  (parent
+//    (local.get $temp.foo)  ;; value saved to a local
+//    (unreachable)          ;; complex effect replaced by unreachable
+//    (i32.const)
+//  )
+//
+// After this it is safe to reorder and remove things from the parent: all
+// interesting interactions happen before the parent.
+//
+// Typical usage is to call getReplacement() will produces the entire output
+// just shown (i.e., possible initial local.sets and other stuff that was pulled
+// out, followed by the parent, as relevant). Note that getReplacement() may
+// omit the parent, if it had an unreachable child. That is useful behavior in
+// that it removes unneeded code (& otherwise some users of this code would need
+// to write their own removal logic). However, that does imply that it is valid
+// to remove the parent in such cases, which is not so for e.g. br when it is
+// the last thing keeping a block reachable. Calling this with something like a
+// struct.new or a call (the current intended users) is valid; if we want to
+// generalize this fully then we need to make changes here.
 //
 // TODO: use in more places
 struct ChildLocalizer {
-  std::vector<LocalSet*> sets;
-
-  ChildLocalizer(Expression* input,
+  ChildLocalizer(Expression* parent,
                  Function* func,
-                 Module* wasm,
-                 const PassOptions& options) {
-    Builder builder(*wasm);
-    ChildIterator iterator(input);
+                 Module& wasm,
+                 const PassOptions& options)
+    : parent(parent), wasm(wasm) {
+    Builder builder(wasm);
+    ChildIterator iterator(parent);
     auto& children = iterator.children;
     auto num = children.size();
 
@@ -77,15 +99,31 @@ struct ChildLocalizer {
       // The children are in reverse order in ChildIterator, but we want to
       // process them in the normal order.
       auto* child = *children[num - 1 - i];
-      effects.emplace_back(options, *wasm, child);
+      effects.emplace_back(options, wasm, child);
     }
 
     // Go through the children and move to locals those that we need to.
     for (Index i = 0; i < num; i++) {
       auto** childp = children[num - 1 - i];
       auto* child = *childp;
       if (child->type == Type::unreachable) {
-        break;
+        // Move the child out, and put an unreachable in its place (note that we
+        // don't need an actual set here, as there is no value to set to a
+        // local).
+        sets.push_back(child);
+        *childp = builder.makeUnreachable();
+        hasUnreachableChild = true;
+        continue;
+      }
+
+      if (hasUnreachableChild) {
+        // Once we pass one unreachable, we only need to copy the children over.
+        // (The only reason we still need them is that they may be needed for
+        // validation, e.g. if one contains a break to a block that is the only
+        // reason the block has type none.)
+        sets.push_back(builder.makeDrop(child));
+        *childp = builder.makeUnreachable();
+        continue;
       }
 
       // Use a local if we need to. That is the case either if this has side
@@ -106,6 +144,36 @@ struct ChildLocalizer {
       }
     }
   }
+
+  // Helper that gets a replacement for the parent: a block containing the
+  // sets + the parent. This will not contain the parent if we don't need it
+  // (if it was never reached).
+  Expression* getReplacement() {
+    if (sets.empty()) {
+      // Nothing to add.
+      return parent;
+    }
+
+    auto* block = Builder(wasm).makeBlock();
+    block->list.set(sets);
+    if (hasUnreachableChild) {
+      // If there is an unreachable child then we do not need the parent at all,
+      // and we know the type is unreachable.
+      block->type = Type::unreachable;
+    } else {
+      // Otherwise, add the parent and finalize.
+      block->list.push_back(parent);
+      block->finalize();
+    }
+    return block;
+  }
+
+private:
+  Expression* parent;
+  Module& wasm;
+
+  std::vector<Expression*> sets;
+  bool hasUnreachableChild = false;
 };
 
 } // namespace wasm

src/passes/GlobalTypeOptimization.cpp

@@ -365,13 +365,8 @@ struct GlobalTypeOptimization : public Pass {
           if (!func) {
             Fatal() << "TODO: side effects in removed fields in globals\n";
           }
-          auto* block = Builder(*getModule()).makeBlock();
-          auto sets =
-            ChildLocalizer(curr, func, getModule(), getPassOptions()).sets;
-          block->list.set(sets);
-          block->list.push_back(curr);
-          block->finalize(curr->type);
-          replaceCurrent(block);
+          ChildLocalizer localizer(curr, func, *getModule(), getPassOptions());
+          replaceCurrent(localizer.getReplacement());
         }
 
         // Remove the unneeded operands.

test/lit/passes/gto-removals.wast

@@ -925,3 +925,49 @@
     (unreachable)
   )
 )
+
+(module
+  ;; CHECK:      (rec
+  ;; CHECK-NEXT:  (type $struct (sub (struct )))
+  (type $struct (sub (struct (field anyref) (field i32) (field f32) (field f64))))
+
+  ;; CHECK:       (type $1 (func (result (ref $struct))))
+
+  ;; CHECK:      (func $func (type $1) (result (ref $struct))
+  ;; CHECK-NEXT:  (local $0 (ref $struct))
+  ;; CHECK-NEXT:  (local $1 f64)
+  ;; CHECK-NEXT:  (local.set $0
+  ;; CHECK-NEXT:   (call $func)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (local.set $1
+  ;; CHECK-NEXT:   (block (result f64)
+  ;; CHECK-NEXT:    (if
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:     (then
+  ;; CHECK-NEXT:      (unreachable)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (f64.const 30)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.new_default $struct)
+  ;; CHECK-NEXT: )
+  (func $func (result (ref $struct))
+    ;; The fields can be removed here, but the effects must be preserved before
+    ;; the struct.new. The consts in the middle can vanish entirely.
+    (struct.new $struct
+      (call $func)
+      (i32.const 10)
+      (f32.const 20)
+      (block (result f64)
+        (if
+          (i32.const 0)
+          (then
+            (unreachable)
+          )
+        )
+        (f64.const 30)
+      )
+    )
+  )
+)