[Wasm GC] Add TypeSSA pass (#5299)

This creates new nominal types for each (interesting) struct.new. That then allows
type-based optimizations to be more precise, as those optimizations will track
separate info for each struct.new, in effect. That is kind of like SSA, however, we
do not handle merges. For example:

x = struct.new $A (5);
print(x.value);
y = struct.new $A (11);
print(y.value);

// => //

x = struct.new $A.x (5);
print(x.value);
y = struct.new $A.y (11);
print(y.value);

After the pass runs each of those struct.new creates a unique type, and type-based
analysis can see that 5 or 11 are the only values written in that type (if nothing else
writes there).

This bloats the type section with the new subtypes, so it is best used with a pass
to merge unneeded duplicate types, which a later PR will add. That later PR will
exactly merge back in the types created here, which are nominally different but
indistinguishable otherwise.

This pass is not enabled by default. It's not clear yet where is the best place to do it,
as it must be balanced by type merging, but it might be better to do multiple
rounds of optimization between the two. Needs more investigation.

5 files changed, 286 insertions, 0 deletions

diff --git a/src/ir/names.h b/src/ir/names.h
index b2165c2c2..b908ad0e7 100644
--- a/src/ir/names.h
+++ b/src/ir/names.h
@@ -92,6 +92,12 @@ inline Name getValidLocalName(Function& func, Name root) {
                       [&](Name test) { return !func.hasLocalIndex(test); });
 }
 
+template<typename T>
+inline Name getValidNameGivenExisting(Name root, const T& existingNames) {
+  return getValidName(root,
+                      [&](Name test) { return !existingNames.count(test); });
+}
+
 class MinifiedNameGenerator {
   size_t state = 0;
 
diff --git a/src/passes/CMakeLists.txt b/src/passes/CMakeLists.txt
index c19f19473..ea06bebb1 100644
--- a/src/passes/CMakeLists.txt
+++ b/src/passes/CMakeLists.txt
@@ -100,6 +100,7 @@ set(passes_SOURCES
   ReReloop.cpp
   TrapMode.cpp
   TypeRefining.cpp
+  TypeSSA.cpp
   SafeHeap.cpp
   SimplifyGlobals.cpp
   SimplifyLocals.cpp
diff --git a/src/passes/TypeSSA.cpp b/src/passes/TypeSSA.cpp
new file mode 100644
index 000000000..75c2d4651
--- /dev/null
+++ b/src/passes/TypeSSA.cpp
@@ -0,0 +1,275 @@
+/*
+ * Copyright 2022 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.
+ */
+
+//
+// Adds a new type for each struct/array.new. This is similar to SSA, which
+// defines each value in its own SSA register as opposed to reusing a local;
+// likewise, this defines a new type at each location that creates an instance.
+// As with SSA, this then allows other passes to be more efficient:
+//
+//  x = struct.new $A (5);
+//  print(x.value);
+//
+//  y = struct.new $A (11);
+//  print(y.value);
+//
+// A type-based optimization on that will not be able to do much, as it will see
+// that we write either 5 or 11 into the value of type $A. But after TypeSSA we
+// have this:
+//
+//  x = struct.new $A.x (5);
+//  print(x.value);
+//
+//  y = struct.new $A.y (11);
+//  print(y.value);
+//
+// Now each place has its own type, and a single value is possible in each type,
+// allowing us to infer what is printed.
+//
+// Note that the metaphor of SSA is not perfect here as we do not handle merges,
+// that is, when there is a value that can read from more than one struct.new
+// then we do nothing atm. We could create a phi there, but in general that
+// would require multiple inheritance. TODO think more on that
+//
+
+#include "ir/find_all.h"
+#include "ir/module-utils.h"
+#include "ir/names.h"
+#include "ir/possible-constant.h"
+#include "pass.h"
+#include "wasm-builder.h"
+#include "wasm.h"
+
+namespace wasm {
+
+namespace {
+
+struct News {
+  std::vector<StructNew*> structNews;
+  // TODO: arrayNews of all kinds
+};
+
+struct NewFinder : public PostWalker<NewFinder> {
+  News news;
+
+  void visitStructNew(StructNew* curr) { news.structNews.push_back(curr); }
+
+  // TODO arrays
+};
+
+struct TypeSSA : public Pass {
+  // Only modifies struct/array.new types.
+  bool requiresNonNullableLocalFixups() override { return false; }
+
+  Module* module;
+
+  void run(Module* module_) override {
+    module = module_;
+
+    if (!module->features.hasGC()) {
+      return;
+    }
+
+    // First, find all the struct/array.news.
+
+    ModuleUtils::ParallelFunctionAnalysis<News> analysis(
+      *module, [&](Function* func, News& news) {
+        if (func->imported()) {
+          return;
+        }
+
+        NewFinder finder;
+        finder.walk(func->body);
+        news = std::move(finder.news);
+      });
+
+    // Also find news in the module scope.
+    NewFinder moduleFinder;
+    moduleFinder.walkModuleCode(module);
+
+    // Process all the news to find the ones we want to modify, adding them to
+    // newsToModify. Note that we must do so in a deterministic order.
+    ModuleUtils::iterDefinedFunctions(
+      *module, [&](Function* func) { processNews(analysis.map[func]); });
+    processNews(moduleFinder.news);
+
+    // Modify the ones we found are relevant. We must modify them all at once as
+    // in the isorecursive type system we want to create a single new rec group
+    // for them all (see below).
+    modifyNews();
+  }
+
+  News newsToModify;
+
+  // As we generate new names, use a consistent index.
+  Index nameCounter = 0;
+
+  void processNews(const News& news) {
+    for (auto* curr : news.structNews) {
+      if (isInteresting(curr)) {
+        newsToModify.structNews.push_back(curr);
+      }
+    }
+  }
+
+  void modifyNews() {
+    auto& structNews = newsToModify.structNews;
+    auto num = structNews.size();
+    if (num == 0) {
+      return;
+    }
+
+    // We collected all the instructions we want to create new types for. Now we
+    // can create those new types, which we do all at once. It is important to
+    // do so in the isorecursive type system as if we create a new singleton rec
+    // group for each then all those will end up identical to each other, so
+    // instead we'll make a single big rec group for them all at once.
+    //
+    // Our goal is to create a completely new/fresh/private type for each
+    // instruction that we found. In the isorecursive type system there isn't an
+    // explicit way to do so, but at least if the new rec group is very large
+    // the risk of collision with another rec group in the program is small.
+    // Note that the risk of collision with things outside of this module is
+    // also a possibility, and so for that reason this pass is likely mostly
+    // useful in the closed-world scenario.
+
+    TypeBuilder builder(num);
+    for (Index i = 0; i < num; i++) {
+      auto* curr = structNews[i];
+      auto oldType = curr->type.getHeapType();
+      builder[i] = oldType.getStruct();
+      builder[i].subTypeOf(oldType);
+    }
+    builder.createRecGroup(0, num);
+    auto result = builder.build();
+    assert(!result.getError());
+    auto newTypes = *result;
+    assert(newTypes.size() == num);
+
+    // The new types must not overlap with any existing ones. If they do, then
+    // it would be unsafe to apply this optimization (if casts exist to the
+    // existing types, the new types merged with them would now succeed on those
+    // casts). We could try to create a "weird" rec group to avoid this (e.g. we
+    // could make a rec group larger than any existing one, or with an initial
+    // member that is "random"), but hopefully this is rare, so just error for
+    // now.
+    //
+    // Note that it is enough to check one of the types: either the entire rec
+    // group gets merged, so they are all merged, or not.
+    std::vector<HeapType> typesVec = ModuleUtils::collectHeapTypes(*module);
+    std::unordered_set<HeapType> typesSet(typesVec.begin(), typesVec.end());
+    if (typesSet.count(newTypes[0])) {
+      Fatal() << "Rec group collision in TypeSSA! Please file a bug";
+    }
+#ifndef NDEBUG
+    // Verify the above assumption, just to be safe.
+    for (auto newType : newTypes) {
+      assert(!typesSet.count(newType));
+    }
+#endif
+
+    // Success: we can apply the new types.
+
+    // We'll generate nice names as we go, if we can, so first scan existing
+    // ones to avoid collisions.
+    std::unordered_set<Name> existingTypeNames;
+    auto& typeNames = module->typeNames;
+    for (auto& [type, info] : typeNames) {
+      existingTypeNames.insert(info.name);
+    }
+
+    for (Index i = 0; i < num; i++) {
+      auto* curr = structNews[i];
+      auto oldType = curr->type.getHeapType();
+      auto newType = newTypes[i];
+      curr->type = Type(newType, NonNullable);
+
+      // If the old type has a nice name, make a nice name for the new one.
+      if (typeNames.count(oldType)) {
+        auto intendedName = typeNames[oldType].name.toString() + '$' +
+                            std::to_string(++nameCounter);
+        auto newName =
+          Names::getValidNameGivenExisting(intendedName, existingTypeNames);
+        // Copy the old field names; only change the type's name itself.
+        auto info = typeNames[oldType];
+        info.name = newName;
+        typeNames[newType] = info;
+        existingTypeNames.insert(newName);
+      }
+    }
+
+    // TODO: arrays
+  }
+
+  // An interesting StructNew, which we think is worth creating a new type for,
+  // is one that can be optimized better with a new type. That means it must
+  // have something interesting for optimizations to work with.
+  //
+  // TODO: We may add new optimizations in the future that can benefit from more
+  //       things, so it may be interesting to experiment with considering all
+  //       news as "interesting" when we add major new type-based optimization
+  //       passes.
+  bool isInteresting(StructNew* curr) {
+    if (curr->type == Type::unreachable) {
+      // This is dead code anyhow.
+      return false;
+    }
+
+    if (curr->isWithDefault()) {
+      // This starts with all default values - zeros and nulls - and that might
+      // be useful.
+      //
+      // (A struct whose fields are all bottom types only has a single possible
+      // value in each field anyhow, so that is not interesting, but also
+      // unreasonable to occur in practice as other optimizations should handle
+      // it.)
+      return true;
+    }
+
+    // Look for at least one interesting operand.
+    auto& fields = curr->type.getHeapType().getStruct().fields;
+    for (Index i = 0; i < fields.size(); i++) {
+      assert(i <= curr->operands.size());
+      auto* operand = curr->operands[i];
+      if (operand->type != fields[i].type) {
+        // Excellent, this field has an interesting type - more refined than the
+        // declared type, and which optimizations might benefit from.
+        //
+        // TODO: If we add GUFA integration, we could check for an exact type
+        //       here - even if the type is not more refined, but it is more
+        //       precise, that is interesting.
+        return true;
+      }
+
+      // TODO: fallthrough
+      PossibleConstantValues value;
+      value.note(operand, *module);
+      if (value.isConstantLiteral() || value.isConstantGlobal()) {
+        // This is a constant that passes may benefit from.
+        return true;
+      }
+    }
+
+    // Nothing interesting.
+    return false;
+  }
+};
+
+} // anonymous namespace
+
+Pass* createTypeSSAPass() { return new TypeSSA(); }
+
+} // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index ee4e2ed21..b97d182c6 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -461,6 +461,9 @@ void PassRegistry::registerPasses() {
   registerPass("trap-mode-js",
                "replace trapping operations with js semantics",
                createTrapModeJS);
+  registerPass("type-ssa",
+               "create new nominal types to help other optimizations",
+               createTypeSSAPass);
   registerPass("untee",
                "removes local.tees, replacing them with sets and gets",
                createUnteePass);
diff --git a/src/passes/passes.h b/src/passes/passes.h
index ae0cc62e2..126dc5f12 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -148,6 +148,7 @@ Pass* createSSAifyNoMergePass();
 Pass* createTrapModeClamp();
 Pass* createTrapModeJS();
 Pass* createTypeRefiningPass();
+Pass* createTypeSSAPass();
 Pass* createUnteePass();
 Pass* createVacuumPass();