Generalize moving of identical code from if/select arms (#3833)

Effects are fine in the moved code, if we are doing so on an if
(which runs just one arm anyhow).

Allow unreachable, which lets us hoist returns for example.

Allow none, which lets us hoist drop and call for example. For
this we also need to be careful with subtyping, as at least drop
is polymorphic, so the child types may not have an LUB (see
example in code).

Adds a small ShallowEffectAnalyzer child of EffectAnalyzer that
calls visit to just do a shallow analysis (instead of walk which
walks the children).

4 files changed, 210 insertions, 44 deletions

diff --git a/src/ir/effects.h b/src/ir/effects.h
index 3d71a8433..f9b6ca6fc 100644
--- a/src/ir/effects.h
+++ b/src/ir/effects.h
@@ -716,6 +716,20 @@ private:
   }
 };
 
+// Calculate effects only on the node itself (shallowly), and not on
+// children.
+class ShallowEffectAnalyzer : public EffectAnalyzer {
+public:
+  ShallowEffectAnalyzer(const PassOptions& passOptions,
+                        FeatureSet features,
+                        Expression* ast = nullptr)
+    : EffectAnalyzer(passOptions, features) {
+    if (ast) {
+      visit(ast);
+    }
+  }
+};
+
 } // namespace wasm
 
 #endif // wasm_ir_effects_h
diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index 5df057dcc..5eba1e3b9 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -2786,10 +2786,6 @@ private:
   // Optimize an if-else or a select, something with a condition and two
   // arms with outputs.
   template<typename T> void optimizeTernary(T* curr) {
-    if (curr->type == Type::unreachable) {
-      return;
-    }
-
     using namespace Match;
     Builder builder(*getModule());
 
@@ -2809,7 +2805,7 @@ private:
     //      (Y)
     //    )
     //  )
-    {
+    if (curr->type != Type::unreachable) {
       Unary* un;
       Expression* x;
       Const* c;
@@ -2860,27 +2856,42 @@ private:
       // as we go, then do a single replaceCurrent() at the end.
       SmallVector<Expression*, 1> chain;
       while (1) {
-        // TODO: handle none as well as concrete types (can pull a drop out, for
-        //       example)
-        // TODO: consider the case with more children than 1
-        if (curr->ifTrue->type.isConcrete() &&
-            !Properties::isControlFlowStructure(curr->ifTrue) &&
+        // Ignore control flow structures (which are handled in MergeBlocks).
+        if (!Properties::isControlFlowStructure(curr->ifTrue) &&
             ExpressionAnalyzer::shallowEqual(curr->ifTrue, curr->ifFalse)) {
+          // TODO: consider the case with more than one child.
           ChildIterator ifTrueChildren(curr->ifTrue);
           if (ifTrueChildren.children.size() == 1) {
+            ChildIterator ifFalseChildren(curr->ifFalse);
+            // ifTrue and ifFalse's children will become the direct children of
+            // curr, and so there must be an LUB for curr to have a proper type.
+            // An example where that does not happen is this:
+            //
+            //  (if
+            //    (condition)
+            //    (drop (i32.const 1))
+            //    (drop (f64.const 2.0))
+            //  )
+            auto* ifTrueChild = *ifTrueChildren.begin();
+            auto* ifFalseChild = *ifFalseChildren.begin();
+            bool validTypes =
+              Type::hasLeastUpperBound(ifTrueChild->type, ifFalseChild->type);
             // If the expression we are about to move outside has side effects,
-            // we cannot replace two instances with one (and there might be
-            // ordering issues as well, for a select's condition).
-            // TODO: handle certain side effects when possible
-            EffectAnalyzer shallowEffects(getPassOptions(),
-                                          getModule()->features);
-            shallowEffects.visit(curr->ifTrue);
-            if (!shallowEffects.hasSideEffects()) {
+            // then we cannot do so in general with a select: we'd be reducing
+            // the amount of the effects as well as moving them. For an if,
+            // the side effects execute once, so there is no problem.
+            // TODO: handle certain side effects when possible in select
+            bool validEffects = std::is_same<T, If>::value ||
+                                !ShallowEffectAnalyzer(getPassOptions(),
+                                                       getModule()->features,
+                                                       curr->ifTrue)
+                                   .hasSideEffects();
+            if (validTypes && validEffects) {
               // Replace ifTrue with its child.
-              curr->ifTrue = *ifTrueChildren.begin();
+              curr->ifTrue = ifTrueChild;
               // Relace ifFalse with its child, and reuse that node outside.
               auto* reuse = curr->ifFalse;
-              curr->ifFalse = *ChildIterator(curr->ifFalse).begin();
+              curr->ifFalse = ifFalseChild;
               // curr's type may have changed, if the instructions we moved out
               // had different input types than output types.
               curr->finalize();
diff --git a/test/lit/passes/optimize-instructions-gc.wast b/test/lit/passes/optimize-instructions-gc.wast
index 3ae0e1d91..c0474a823 100644
--- a/test/lit/passes/optimize-instructions-gc.wast
+++ b/test/lit/passes/optimize-instructions-gc.wast
@@ -849,7 +849,7 @@
       )
     )
   )
-  ;; CHECK:      (func $ternary-identical-arms-but-side-effect (param $x (ref null $struct)) (param $y (ref null $struct)) (param $z i32)
+  ;; CHECK:      (func $select-identical-arms-but-side-effect (param $x (ref null $struct)) (param $y (ref null $struct)) (param $z i32)
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (select
   ;; CHECK-NEXT:    (struct.get_u $struct $i8
@@ -862,7 +862,7 @@
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
-  (func $ternary-identical-arms-but-side-effect (param $x (ref null $struct)) (param $y (ref null $struct)) (param $z i32)
+  (func $select-identical-arms-but-side-effect (param $x (ref null $struct)) (param $y (ref null $struct)) (param $z i32)
     (drop
       (select
         ;; the arms are equal but have side effects
@@ -901,4 +901,30 @@
       )
     )
   )
+  ;; CHECK:      (func $if-identical-arms-with-side-effect (param $x (ref null $struct)) (param $y (ref null $struct)) (param $z i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (struct.get_u $struct $i8
+  ;; CHECK-NEXT:    (if (result (ref null $struct))
+  ;; CHECK-NEXT:     (local.get $z)
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:     (local.get $y)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $if-identical-arms-with-side-effect (param $x (ref null $struct)) (param $y (ref null $struct)) (param $z i32)
+    (drop
+      (if (result i32)
+        (local.get $z)
+        ;; the arms are equal and have side effects, but that is ok with an if
+        ;; which only executes one side anyhow
+        (struct.get_u $struct 0
+          (local.get $x)
+        )
+        (struct.get_u $struct 0
+          (local.get $y)
+        )
+      )
+    )
+  )
 )
diff --git a/test/lit/passes/optimize-instructions.wast b/test/lit/passes/optimize-instructions.wast
index 7150fc839..3ed0b3115 100644
--- a/test/lit/passes/optimize-instructions.wast
+++ b/test/lit/passes/optimize-instructions.wast
@@ -101,12 +101,10 @@
     )
   )
   ;; CHECK:      (func $if-eqz-two-arms (param $i1 i32)
-  ;; CHECK-NEXT:  (if
-  ;; CHECK-NEXT:   (local.get $i1)
-  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (if (result i32)
+  ;; CHECK-NEXT:    (local.get $i1)
   ;; CHECK-NEXT:    (i32.const 12)
-  ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (drop
   ;; CHECK-NEXT:    (i32.const 11)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
@@ -125,14 +123,12 @@
     )
   )
   ;; CHECK:      (func $if-eqz-two-arms-i64 (param $i2 i64)
-  ;; CHECK-NEXT:  (if
-  ;; CHECK-NEXT:   (i64.eqz
-  ;; CHECK-NEXT:    (local.get $i2)
-  ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (if (result i32)
+  ;; CHECK-NEXT:    (i64.eqz
+  ;; CHECK-NEXT:     (local.get $i2)
+  ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (i32.const 11)
-  ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (drop
   ;; CHECK-NEXT:    (i32.const 12)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
@@ -11928,29 +11924,44 @@
       )
     )
   )
-  ;; CHECK:      (func $ternary-identical-arms-but-type-is-none (param $x i32) (param $y i32) (param $z i32)
-  ;; CHECK-NEXT:  (if
-  ;; CHECK-NEXT:   (local.get $z)
-  ;; CHECK-NEXT:   (drop
-  ;; CHECK-NEXT:    (i32.eqz
+  ;; CHECK:      (func $ternary-identical-arms-and-type-is-none (param $x i32) (param $y i32) (param $z i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.eqz
+  ;; CHECK-NEXT:    (if (result i32)
+  ;; CHECK-NEXT:     (local.get $z)
   ;; CHECK-NEXT:     (local.get $x)
-  ;; CHECK-NEXT:    )
-  ;; CHECK-NEXT:   )
-  ;; CHECK-NEXT:   (drop
-  ;; CHECK-NEXT:    (i32.eqz
   ;; CHECK-NEXT:     (local.get $y)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
-  (func $ternary-identical-arms-but-type-is-none (param $x i32) (param $y i32) (param $z i32)
+  (func $ternary-identical-arms-and-type-is-none (param $x i32) (param $y i32) (param $z i32)
     (if
       (local.get $z)
-      ;; identical arms, but type is none
       (drop (i32.eqz (local.get $x)))
       (drop (i32.eqz (local.get $y)))
     )
   )
+  ;; CHECK:      (func $ternary-identical-arms-and-type-is-none-child-types-mismatch (param $x i32) (param $y i32) (param $z i32)
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (local.get $z)
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (f64.const 2.34)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $ternary-identical-arms-and-type-is-none-child-types-mismatch (param $x i32) (param $y i32) (param $z i32)
+    (if
+      (local.get $z)
+      ;; the drop cannot be hoisted out, since the children's type mismatch
+      ;; would not allow us to give a proper type to the if.
+      (drop (i32.const 1))
+      (drop (f64.const 2.34))
+    )
+  )
   ;; CHECK:      (func $ternary-identical-arms-but-block (param $x i32) (param $y i32) (param $z i32)
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (select
@@ -12013,4 +12024,108 @@
       )
     )
   )
+  ;; CHECK:      (func $ternary-identical-arms-br_if-same (param $x i32) (param $y i32) (param $z i32)
+  ;; CHECK-NEXT:  (block $block
+  ;; CHECK-NEXT:   (br_if $block
+  ;; CHECK-NEXT:    (if (result i32)
+  ;; CHECK-NEXT:     (local.get $z)
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:     (local.get $y)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $ternary-identical-arms-br_if-same (param $x i32) (param $y i32) (param $z i32)
+    (block $block
+      (if
+        (local.get $z)
+        ;; two br_ifs with the same target are shallowly identical
+        (br_if $block
+          (local.get $x)
+        )
+        (br_if $block
+          (local.get $y)
+        )
+      )
+    )
+  )
+  ;; CHECK:      (func $ternary-identical-arms-br_if-different (param $x i32) (param $y i32) (param $z i32)
+  ;; CHECK-NEXT:  (block $block1
+  ;; CHECK-NEXT:   (block $block2
+  ;; CHECK-NEXT:    (if
+  ;; CHECK-NEXT:     (local.get $z)
+  ;; CHECK-NEXT:     (br_if $block1
+  ;; CHECK-NEXT:      (local.get $x)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (br_if $block2
+  ;; CHECK-NEXT:      (local.get $y)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $ternary-identical-arms-br_if-different (param $x i32) (param $y i32) (param $z i32)
+    (block $block1
+      (block $block2
+        (if
+          (local.get $z)
+          ;; two br_ifs with different targets are not shallowly identical
+          (br_if $block1
+            (local.get $x)
+          )
+          (br_if $block2
+            (local.get $y)
+          )
+        )
+      )
+    )
+  )
+  ;; CHECK:      (func $ternary-identical-arms-return (param $x i32) (param $y i32) (param $z i32) (result i32)
+  ;; CHECK-NEXT:  (block $block
+  ;; CHECK-NEXT:   (return
+  ;; CHECK-NEXT:    (if (result i32)
+  ;; CHECK-NEXT:     (local.get $z)
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:     (local.get $y)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $ternary-identical-arms-return (param $x i32) (param $y i32) (param $z i32) (result i32)
+    (block $block
+      (if
+        (local.get $z)
+        (return
+          (local.get $x)
+        )
+        (return
+          (local.get $y)
+        )
+      )
+    )
+  )
+  ;; CHECK:      (func $send-i32 (param $0 i32)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $send-i32 (param i32))
+  ;; CHECK:      (func $ternary-identical-arms-call (param $x i32) (param $y i32) (param $z i32)
+  ;; CHECK-NEXT:  (call $send-i32
+  ;; CHECK-NEXT:   (if (result i32)
+  ;; CHECK-NEXT:    (local.get $z)
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:    (local.get $y)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $ternary-identical-arms-call (param $x i32) (param $y i32) (param $z i32)
+    (if
+      (local.get $z)
+      (call $send-i32
+        (local.get $x)
+      )
+      (call $send-i32
+        (local.get $y)
+      )
+    )
+  )
 )