[Wasm GC] GlobalTypeOptimization: Turn fields immutable when possible (#4213)

Add a new pass to perform global type optimization. So far this just
does one thing, to find fields with no struct.set and to turn them
immutable (where possible - sub and supertypes must agree).

To do that, this adds a GlobalTypeRewriter utility which rewrites
all the heap types in the module, allowing changes while doing so.
In this PR, the change is to flip the mutable field. Otherwise, the
utility handles all the boilerplate of creating temp heap types using
a TypeBuilder, and it handles replacing the types in every place
they are used in the module.

This is not enabled by default yet as I don't see enough of a benefit
on j2cl. This PR is basically the simplest thing to do in the space of
global type optimization, and the simplest way I can think of to
fully test the GlobalTypeRewriter (which can't be done as a unit
test, really, since we want to emit a full module and validate it etc.).

This PR builds the foundation for more complicated things like
removing unused fields, subtyping fields, and more.

4 files changed, 187 insertions, 0 deletions

diff --git a/src/passes/CMakeLists.txt b/src/passes/CMakeLists.txt
index a49ae8980..4ee7fe9d1 100644
--- a/src/passes/CMakeLists.txt
+++ b/src/passes/CMakeLists.txt
@@ -87,6 +87,7 @@ set(passes_SOURCES
   SSAify.cpp
   Untee.cpp
   Vacuum.cpp
+  GlobalTypeOptimization.cpp
   ${CMAKE_CURRENT_BINARY_DIR}/WasmIntrinsics.cpp
   ${passes_HEADERS}
 )
diff --git a/src/passes/GlobalTypeOptimization.cpp b/src/passes/GlobalTypeOptimization.cpp
new file mode 100644
index 000000000..89d5b8c0a
--- /dev/null
+++ b/src/passes/GlobalTypeOptimization.cpp
@@ -0,0 +1,182 @@
+/*
+ * 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.
+ */
+
+//
+// Optimize types at the global level, altering fields etc. on the set of heap
+// types defined in the module.
+//
+//  * Immutability: If a field has no struct.set, it can become immutable.
+//
+// TODO: Specialize field types.
+// TODO: Remove unused fields.
+//
+
+#include "ir/struct-utils.h"
+#include "ir/subtypes.h"
+#include "ir/type-updating.h"
+#include "ir/utils.h"
+#include "pass.h"
+#include "support/small_set.h"
+#include "wasm-builder.h"
+#include "wasm-type.h"
+#include "wasm.h"
+
+using namespace std;
+
+namespace wasm {
+
+namespace {
+
+// Information about usage of a field.
+struct FieldInfo {
+  bool hasWrite = false;
+
+  void noteWrite() { hasWrite = true; }
+
+  bool combine(const FieldInfo& other) {
+    if (!hasWrite && other.hasWrite) {
+      hasWrite = true;
+      return true;
+    }
+    return false;
+  }
+};
+
+struct FieldInfoScanner : public Scanner<FieldInfo, FieldInfoScanner> {
+  Pass* create() override {
+    return new FieldInfoScanner(functionNewInfos, functionSetInfos);
+  }
+
+  FieldInfoScanner(FunctionStructValuesMap<FieldInfo>& functionNewInfos,
+                   FunctionStructValuesMap<FieldInfo>& functionSetInfos)
+    : Scanner<FieldInfo, FieldInfoScanner>(functionNewInfos, functionSetInfos) {
+  }
+
+  void noteExpression(Expression* expr,
+                      HeapType type,
+                      Index index,
+                      FieldInfo& info) {
+    info.noteWrite();
+  }
+
+  void
+  noteDefault(Type fieldType, HeapType type, Index index, FieldInfo& info) {
+    info.noteWrite();
+  }
+
+  void noteCopy(HeapType type, Index index, FieldInfo& info) {
+    info.noteWrite();
+  }
+};
+
+struct GlobalTypeOptimization : public Pass {
+  void run(PassRunner* runner, Module* module) override {
+    if (getTypeSystem() != TypeSystem::Nominal) {
+      Fatal() << "GlobalTypeOptimization requires nominal typing";
+    }
+
+    // Find and analyze struct operations inside each function.
+    FunctionStructValuesMap<FieldInfo> functionNewInfos(*module),
+      functionSetInfos(*module);
+    FieldInfoScanner scanner(functionNewInfos, functionSetInfos);
+    scanner.run(runner, module);
+    scanner.walkModuleCode(module);
+
+    // Combine the data from the functions.
+    StructValuesMap<FieldInfo> combinedNewInfos, combinedSetInfos;
+    functionSetInfos.combineInto(combinedSetInfos);
+    // TODO: combine newInfos as well, once we have a need for that (we will
+    //       when we do things like subtyping).
+
+    // Find which fields are immutable in all super- and sub-classes. To see
+    // that, propagate sets in both directions. This is necessary because we
+    // cannot have a supertype's field be immutable while a subtype's is not -
+    // they must match for us to preserve subtyping.
+    //
+    // Note that we do not need to care about types here: If the fields were
+    // mutable before, then they must have had identical types for them to be
+    // subtypes (as wasm only allows the type to differ if the fields are
+    // immutable). Note that by making more things immutable we therefore make
+    // it possible to apply more specific subtypes in subtype fields.
+    TypeHierarchyPropagator<FieldInfo> propagator(*module);
+    propagator.propagateToSuperAndSubTypes(combinedSetInfos);
+
+    // Maps types to a vector of booleans that indicate if we can turn the
+    // field immutable. To avoid eager allocation of memory, the vectors are
+    // only resized when we actually have a true to place in them (which is
+    // rare).
+    using CanBecomeImmutable = std::unordered_map<HeapType, std::vector<bool>>;
+    CanBecomeImmutable canBecomeImmutable;
+
+    for (auto type : propagator.subTypes.types) {
+      if (!type.isStruct()) {
+        continue;
+      }
+
+      auto& fields = type.getStruct().fields;
+      for (Index i = 0; i < fields.size(); i++) {
+        if (fields[i].mutable_ == Immutable) {
+          // Already immutable; nothing to do.
+          continue;
+        }
+
+        if (combinedSetInfos[type][i].hasWrite) {
+          // A set exists.
+          continue;
+        }
+
+        // No set exists. Mark it as something we can make immutable.
+        auto& vec = canBecomeImmutable[type];
+        vec.resize(i + 1);
+        vec[i] = true;
+      }
+    }
+
+    // The types are now generally correct, except for their internals, which we
+    // rewrite now.
+    class TypeRewriter : public GlobalTypeRewriter {
+      CanBecomeImmutable& canBecomeImmutable;
+
+    public:
+      TypeRewriter(Module& wasm, CanBecomeImmutable& canBecomeImmutable)
+        : GlobalTypeRewriter(wasm), canBecomeImmutable(canBecomeImmutable) {}
+
+      virtual void modifyStruct(HeapType oldStructType, Struct& struct_) {
+        if (!canBecomeImmutable.count(oldStructType)) {
+          return;
+        }
+
+        auto& newFields = struct_.fields;
+        auto& immutableVec = canBecomeImmutable[oldStructType];
+        for (Index i = 0; i < immutableVec.size(); i++) {
+          if (immutableVec[i]) {
+            newFields[i].mutable_ = Immutable;
+          }
+        }
+      }
+    };
+
+    TypeRewriter(*module, canBecomeImmutable).update();
+  }
+};
+
+} // anonymous namespace
+
+Pass* createGlobalTypeOptimizationPass() {
+  return new GlobalTypeOptimization();
+}
+
+} // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index 5f53a1d78..67087c328 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -155,6 +155,8 @@ void PassRegistry::registerPasses() {
   registerPass(
     "generate-stack-ir", "generate Stack IR", createGenerateStackIRPass);
   registerPass(
+    "gto", "globally optimize GC types", createGlobalTypeOptimizationPass);
+  registerPass(
     "heap2local", "replace GC allocations with locals", createHeap2LocalPass);
   registerPass(
     "inline-main", "inline __original_main into main", createInlineMainPass);
@@ -524,6 +526,7 @@ void PassRunner::addDefaultGlobalOptimizationPrePasses() {
       options.optimizeLevel >= 2) {
     addIfNoDWARFIssues("cfp");
   }
+  // TODO: investigate enabling --gto
 }
 
 void PassRunner::addDefaultGlobalOptimizationPostPasses() {
diff --git a/src/passes/passes.h b/src/passes/passes.h
index 47b246bd4..a08e42f62 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -131,6 +131,7 @@ Pass* createTrapModeClamp();
 Pass* createTrapModeJS();
 Pass* createUnteePass();
 Pass* createVacuumPass();
+Pass* createGlobalTypeOptimizationPass();
 
 } // namespace wasm