1 files changed, 142 insertions, 0 deletions
diff --git a/src/passes/LocalSubtyping.cpp b/src/passes/LocalSubtyping.cpp
new file mode 100644
index 000000000..f435e53d4
--- /dev/null
+++ b/src/passes/LocalSubtyping.cpp
@@ -0,0 +1,142 @@
+/*
+ * Copyright 2021 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//
+// Refines the types of locals where possible. That is, if a local is assigned
+// types that are more specific than the local's declared type, refine the
+// declared type. This can then potentially unlock optimizations later when the
+// local is used, as we have more type info. (However, it may also increase code
+// size in theory, if we end up declaring more types - TODO investigate.)
+//
+
+#include <ir/find_all.h>
+#include <ir/linear-execution.h>
+#include <ir/utils.h>
+#include <pass.h>
+#include <wasm.h>
+
+namespace wasm {
+
+struct LocalSubtyping : public WalkerPass<PostWalker<LocalSubtyping>> {
+  bool isFunctionParallel() override { return true; }
+
+  Pass* create() override { return new LocalSubtyping(); }
+
+  // Shared code to find all sets or gets for each local index. Returns a vector
+  // that maps local indexes => vector of LocalGet*|LocalSet* expressions.
+  template<typename T>
+  std::vector<std::vector<T*>> getLocalOperations(Function* func) {
+    std::vector<std::vector<T*>> ret;
+    ret.resize(func->getNumLocals());
+    FindAll<T> operations(func->body);
+    for (auto* operation : operations.list) {
+      ret[operation->index].push_back(operation);
+    }
+    return ret;
+  }
+
+  void doWalkFunction(Function* func) {
+    if (!getModule()->features.hasGC()) {
+      return;
+    }
+
+    auto varBase = func->getVarIndexBase();
+    auto numLocals = func->getNumLocals();
+
+    auto setsForLocal = getLocalOperations<LocalSet>(func);
+    auto getsForLocal = getLocalOperations<LocalGet>(func);
+
+    // Keep iterating while we find things to change. There can be chains like
+    // X -> Y -> Z where one change enables more. Note that we are O(N^2) on
+    // that atm, but it is a rare pattern as general optimizations
+    // (SimplifyLocals and CoalesceLocals) break up such things; also, even if
+    // we tracked changes more carefully we'd have the case of nested tees
+    // where we could still be O(N^2), so we'd need something more complex here
+    // involving topological sorting. Leave that for if the need arises.
+
+    // TODO: handle cycles of X -> Y -> X etc.
+
+    bool more;
+
+    do {
+      more = false;
+
+      // First, refinalize which will recompute least upper bounds on ifs and
+      // blocks, etc., potentially finding a more specific type. Note that
+      // that utility does not tell us if it changed anything, so we depend on
+      // the next step for knowing if there is more work to do.
+      ReFinalize().walkFunctionInModule(func, getModule());
+
+      // Second, find vars whose actual applied values allow a more specific
+      // type.
+
+      for (Index i = varBase; i < numLocals; i++) {
+        // Find all the types assigned to the var, and compute the optimal LUB.
+        // Note that we do not need to take into account the initial value of
+        // zero or null that locals have: that value has the type of the local,
+        // which is a supertype of all the assigned values anyhow. It will never
+        // be able to tell us of a more specific subtype that is possible.
+        std::unordered_set<Type> types;
+        for (auto* set : setsForLocal[i]) {
+          types.insert(set->value->type);
+        }
+        if (types.empty()) {
+          // Nothing is assigned to this local (other opts will remove it).
+          continue;
+        }
+
+        auto oldType = func->getLocalType(i);
+        auto newType = Type::getLeastUpperBound(types);
+        assert(newType != Type::none); // in valid wasm there must be a LUB
+
+        // Remove non-nullability if we disallow that in locals.
+        if (!getModule()->features.hasGCNNLocals() && newType.isNonNullable()) {
+          newType = Type(newType.getHeapType(), Nullable);
+          // Note that the old type must have been nullable as well, as non-
+          // nullable types cannot be locals without that feature being enabled,
+          // which means that we will not have to do any extra work to handle
+          // non-nullability if we update the type: we are just updating the
+          // heap type, and leaving the type nullable as it was.
+          assert(oldType.isNullable());
+        }
+
+        if (newType != oldType) {
+          // We found a more specific type!
+          assert(Type::isSubType(newType, oldType));
+          func->vars[i - varBase] = newType;
+          more = true;
+
+          // Update gets and tees.
+          for (auto* get : getsForLocal[i]) {
+            get->type = newType;
+          }
+
+          // NB: These tee updates will not be needed if the type of tees
+          //     becomes that of their value, in the spec.
+          for (auto* set : setsForLocal[i]) {
+            if (set->isTee() && set->type != Type::unreachable) {
+              set->type = newType;
+            }
+          }
+        }
+      }
+    } while (more);
+  }
+};
+
+Pass* createLocalSubtypingPass() { return new LocalSubtyping(); }
+
+} // namespace wasm