GlobalStructInference: Handle >2 globals if values coincide (#4714)

In GSI we look for a read of a global in a situation like this:

$global1: value1
$global2: value2

(struct.get $Type (ref))

If global inference shows this get must be of either $global1 or $global2, then we
can optimize to this:

(ref) == $global1 ? value1 : value2

We focus on the case of two values because 1 is handled by other passes, and >2
makes the tradeoffs less clear.

However, a simple extension is the case where there are more than 2 globals, but
there are only two values, and one value is unique to one global:

$global1: valueA
$global2: valueB
$global3: valueA

=>

(ref) == $global2 ? valueB : valueA

We can still use a single comparison here, on the global that has the
unique value. Then the else will handle all the other globals.

This increases the cases that GSI can optimize J2Wasm output by over 50%.

1 files changed, 80 insertions, 14 deletions

diff --git a/src/passes/GlobalStructInference.cpp b/src/passes/GlobalStructInference.cpp
index 42fadf295..c7eec2b94 100644
--- a/src/passes/GlobalStructInference.cpp
+++ b/src/passes/GlobalStructInference.cpp
@@ -187,44 +187,110 @@ struct GlobalStructInference : public Pass {
           return;
         }
 
-        auto& globals = iter->second;
-
-        // TODO: more sizes
-        if (globals.size() != 2) {
-          return;
-        }
-
-        // Check if the relevant fields contain constants, and are immutable.
-        auto& wasm = *getModule();
+        // The field must be immutable.
         auto fieldIndex = curr->index;
         auto& field = type.getHeapType().getStruct().fields[fieldIndex];
         if (field.mutable_ == Mutable) {
           return;
         }
-        auto fieldType = field.type;
+
+        // We are looking for the case where we can pick between two values
+        // using a single comparison. More than two values, or more than a
+        // single comparison, add tradeoffs that may not be worth it, and a
+        // single value (or no value) is already handled by other passes.
+        //
+        // That situation may involve more than two globals. For example we may
+        // have three relevant globals, but two may have the same value. In that
+        // case we can compare against the third:
+        //
+        //  $global0: (struct.new $Type (i32.const 42))
+        //  $global1: (struct.new $Type (i32.const 42))
+        //  $global2: (struct.new $Type (i32.const 1337))
+        //
+        // (struct.get $Type (ref))
+        //   =>
+        // (select
+        //   (i32.const 1337)
+        //   (i32.const 42)
+        //   (ref.eq (ref) $global2))
+        auto& globals = iter->second;
+        if (globals.size() < 2) {
+          return;
+        }
+
+        // Find the constant values and which globals correspond to them.
+        // TODO: SmallVectors?
         std::vector<Literal> values;
+        std::vector<std::vector<Name>> globalsForValue;
+
+        // Check if the relevant fields contain constants.
+        auto& wasm = *getModule();
+        auto fieldType = field.type;
         for (Index i = 0; i < globals.size(); i++) {
-          auto* structNew = wasm.getGlobal(globals[i])->init->cast<StructNew>();
+          Name global = globals[i];
+          auto* structNew = wasm.getGlobal(global)->init->cast<StructNew>();
+          Literal value;
           if (structNew->isWithDefault()) {
-            values.push_back(Literal::makeZero(fieldType));
+            value = Literal::makeZero(fieldType);
           } else {
             auto* init = structNew->operands[fieldIndex];
             if (!Properties::isConstantExpression(init)) {
               // Non-constant; give up entirely.
               return;
             }
-            values.push_back(Properties::getLiteral(init));
+            value = Properties::getLiteral(init);
+          }
+
+          // Process the current value, comparing it against the previous.
+          auto found = std::find(values.begin(), values.end(), value);
+          if (found == values.end()) {
+            // This is a new value.
+            assert(values.size() <= 2);
+            if (values.size() == 2) {
+              // Adding this value would mean we have too many, so give up.
+              return;
+            }
+            values.push_back(value);
+            globalsForValue.push_back({global});
+          } else {
+            // This is an existing value.
+            Index index = found - values.begin();
+            globalsForValue[index].push_back(global);
           }
         }
 
+        // We have some globals (at least 2), and so must have at least one
+        // value. And we have already exited if we have more than 2, so that
+        // only leaves 1 and 2. We are looking for the case of 2 here, since
+        // other passes (ConstantFieldPropagation) can handle 1.
+        if (values.size() == 1) {
+          return;
+        }
+        assert(values.size() == 2);
+
+        // We have two values. Check that we can pick between them using a
+        // single comparison. While doing so, ensure that the index we can check
+        // on is 0, that is, the first value has a single global.
+        if (globalsForValue[0].size() == 1) {
+          // The checked global is already in index 0.
+        } else if (globalsForValue[1].size() == 1) {
+          std::swap(values[0], values[1]);
+          std::swap(globalsForValue[0], globalsForValue[1]);
+        } else {
+          // Both indexes have more than one option, so we'd need more than one
+          // comparison. Give up.
+          return;
+        }
+
         // Excellent, we can optimize here! Emit a select.
         //
         // Note that we must trap on null, so add a ref.as_non_null here.
+        auto checkGlobal = globalsForValue[0][0];
         Builder builder(wasm);
         replaceCurrent(builder.makeSelect(
           builder.makeRefEq(builder.makeRefAs(RefAsNonNull, curr->ref),
                             builder.makeGlobalGet(
-                              globals[0], wasm.getGlobal(globals[0])->type)),
+                              checkGlobal, wasm.getGlobal(checkGlobal)->type)),
           builder.makeConstantExpression(values[0]),
           builder.makeConstantExpression(values[1])));
       }