[Wasm GC] Global Type Optimization: Remove unread fields (#4255)

Add struct.get tracking, and if a field is never read from, simply remove
it.

This will error if a field is written using struct.new with a value with side
effects. It is not clear we can handle that, as if the struct.new is in a
global then we can't save the other values to locals etc. to reorder
things. We could perhaps use other globals for it (ugh) but at least for
now, that corner case does not happen on any code I can see.

This allows a quite large code size reduction on j2wasm output (20%). The
reason is that many vtable fields are not actually read, and so removing
them and the ref.func they hold allows us to get rid of those functions,
and code that they reach.

5 files changed, 273 insertions, 56 deletions

diff --git a/src/ir/struct-utils.h b/src/ir/struct-utils.h
index 442d990fb..5d0f5a9b8 100644
--- a/src/ir/struct-utils.h
+++ b/src/ir/struct-utils.h
@@ -119,17 +119,22 @@ struct FunctionStructValuesMap
 //
 //   void noteCopy(HeapType type, Index index, T& info);
 //
-// We track information from struct.new and struct.set separately, because in
-// struct.new we know more about the type - we know the actual exact type being
-// written to, and not just that it is of a subtype of the instruction's type,
-// which helps later.
+// * Note a read
+//
+//   void noteRead(HeapType type, Index index, T& info);
+//
+// We track information from struct.new and struct.set/struct.get separately,
+// because in struct.new we know more about the type - we know the actual exact
+// type being written to, and not just that it is of a subtype of the
+// instruction's type, which helps later.
 template<typename T, typename SubType>
 struct Scanner : public WalkerPass<PostWalker<Scanner<T, SubType>>> {
   bool isFunctionParallel() override { return true; }
 
   Scanner(FunctionStructValuesMap<T>& functionNewInfos,
-          FunctionStructValuesMap<T>& functionSetInfos)
-    : functionNewInfos(functionNewInfos), functionSetInfos(functionSetInfos) {}
+          FunctionStructValuesMap<T>& functionSetGetInfos)
+    : functionNewInfos(functionNewInfos),
+      functionSetGetInfos(functionSetGetInfos) {}
 
   void visitStructNew(StructNew* curr) {
     auto type = curr->type;
@@ -158,11 +163,25 @@ struct Scanner : public WalkerPass<PostWalker<Scanner<T, SubType>>> {
     }
 
     // Note a write to this field of the struct.
-    noteExpressionOrCopy(
-      curr->value,
-      type.getHeapType(),
-      curr->index,
-      functionSetInfos[this->getFunction()][type.getHeapType()][curr->index]);
+    noteExpressionOrCopy(curr->value,
+                         type.getHeapType(),
+                         curr->index,
+                         functionSetGetInfos[this->getFunction()]
+                                            [type.getHeapType()][curr->index]);
+  }
+
+  void visitStructGet(StructGet* curr) {
+    auto type = curr->ref->type;
+    if (type == Type::unreachable) {
+      return;
+    }
+
+    auto heapType = type.getHeapType();
+    auto index = curr->index;
+    static_cast<SubType*>(this)->noteRead(
+      heapType,
+      index,
+      functionSetGetInfos[this->getFunction()][heapType][index]);
   }
 
   void
@@ -186,7 +205,7 @@ struct Scanner : public WalkerPass<PostWalker<Scanner<T, SubType>>> {
   }
 
   FunctionStructValuesMap<T>& functionNewInfos;
-  FunctionStructValuesMap<T>& functionSetInfos;
+  FunctionStructValuesMap<T>& functionSetGetInfos;
 };
 
 // Helper class to propagate information about fields to sub- and/or super-
diff --git a/src/ir/type-updating.h b/src/ir/type-updating.h
index 83c1e1aa1..d3be7c3ca 100644
--- a/src/ir/type-updating.h
+++ b/src/ir/type-updating.h
@@ -309,6 +309,8 @@ struct TypeUpdater
 // made while doing so.
 class GlobalTypeRewriter {
 public:
+  Module& wasm;
+
   GlobalTypeRewriter(Module& wasm);
   virtual ~GlobalTypeRewriter() {}
 
@@ -332,7 +334,6 @@ public:
   Type getTempType(Type type);
 
 private:
-  Module& wasm;
   TypeBuilder typeBuilder;
 
   // The list of old types.
diff --git a/src/passes/ConstantFieldPropagation.cpp b/src/passes/ConstantFieldPropagation.cpp
index 2fdd8333c..f6bad04b6 100644
--- a/src/passes/ConstantFieldPropagation.cpp
+++ b/src/passes/ConstantFieldPropagation.cpp
@@ -214,7 +214,7 @@ private:
 
 struct PCVScanner : public Scanner<PossibleConstantValues, PCVScanner> {
   Pass* create() override {
-    return new PCVScanner(functionNewInfos, functionSetInfos);
+    return new PCVScanner(functionNewInfos, functionSetGetInfos);
   }
 
   PCVScanner(FunctionStructValuesMap<PossibleConstantValues>& functionNewInfos,
@@ -261,6 +261,10 @@ struct PCVScanner : public Scanner<PossibleConstantValues, PCVScanner> {
     //       analysis (but this is already entering the realm of diminishing
     //       returns).
   }
+
+  void noteRead(HeapType type, Index index, PossibleConstantValues& info) {
+    // Reads do not interest us.
+  }
 };
 
 struct ConstantFieldPropagation : public Pass {
diff --git a/src/passes/GlobalTypeOptimization.cpp b/src/passes/GlobalTypeOptimization.cpp
index 35b10bfff..442b7570a 100644
--- a/src/passes/GlobalTypeOptimization.cpp
+++ b/src/passes/GlobalTypeOptimization.cpp
@@ -19,11 +19,13 @@
 // types defined in the module.
 //
 //  * Immutability: If a field has no struct.set, it can become immutable.
+//  * Fields that are never read from can be removed entirely.
 //
 // TODO: Specialize field types.
 // TODO: Remove unused fields.
 //
 
+#include "ir/effects.h"
 #include "ir/struct-utils.h"
 #include "ir/subtypes.h"
 #include "ir/type-updating.h"
@@ -43,27 +45,34 @@ namespace {
 // Information about usage of a field.
 struct FieldInfo {
   bool hasWrite = false;
+  bool hasRead = false;
 
   void noteWrite() { hasWrite = true; }
+  void noteRead() { hasRead = true; }
 
   bool combine(const FieldInfo& other) {
+    bool changed = false;
     if (!hasWrite && other.hasWrite) {
       hasWrite = true;
-      return true;
+      changed = true;
     }
-    return false;
+    if (!hasRead && other.hasRead) {
+      hasRead = true;
+      changed = true;
+    }
+    return changed;
   }
 };
 
 struct FieldInfoScanner : public Scanner<FieldInfo, FieldInfoScanner> {
   Pass* create() override {
-    return new FieldInfoScanner(functionNewInfos, functionSetInfos);
+    return new FieldInfoScanner(functionNewInfos, functionSetGetInfos);
   }
 
   FieldInfoScanner(FunctionStructValuesMap<FieldInfo>& functionNewInfos,
-                   FunctionStructValuesMap<FieldInfo>& functionSetInfos)
-    : Scanner<FieldInfo, FieldInfoScanner>(functionNewInfos, functionSetInfos) {
-  }
+                   FunctionStructValuesMap<FieldInfo>& functionSetGetInfos)
+    : Scanner<FieldInfo, FieldInfoScanner>(functionNewInfos,
+                                           functionSetGetInfos) {}
 
   void noteExpression(Expression* expr,
                       HeapType type,
@@ -80,9 +89,31 @@ struct FieldInfoScanner : public Scanner<FieldInfo, FieldInfoScanner> {
   void noteCopy(HeapType type, Index index, FieldInfo& info) {
     info.noteWrite();
   }
+
+  void noteRead(HeapType type, Index index, FieldInfo& info) {
+    info.noteRead();
+  }
 };
 
 struct GlobalTypeOptimization : public Pass {
+  StructValuesMap<FieldInfo> combinedSetGetInfos;
+
+  // Maps types to a vector of booleans that indicate whether a field can
+  // become 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).
+  std::unordered_map<HeapType, std::vector<bool>> canBecomeImmutable;
+
+  // Maps each field to its new index after field removals. That is, this
+  // takes into account that fields before this one may have been removed,
+  // which would then reduce this field's index. If a field itself is removed,
+  // it has the special value |RemovedField|. This is always of the full size
+  // of the number of fields, unlike canBecomeImmutable which is lazily
+  // allocated, as if we remove one field that affects the indexes of all the
+  // others anyhow.
+  static const Index RemovedField = Index(-1);
+  std::unordered_map<HeapType, std::vector<Index>> indexesAfterRemovals;
+
   void run(PassRunner* runner, Module* module) override {
     if (getTypeSystem() != TypeSystem::Nominal) {
       Fatal() << "GlobalTypeOptimization requires nominal typing";
@@ -90,50 +121,50 @@ struct GlobalTypeOptimization : public Pass {
 
     // Find and analyze struct operations inside each function.
     FunctionStructValuesMap<FieldInfo> functionNewInfos(*module),
-      functionSetInfos(*module);
-    FieldInfoScanner scanner(functionNewInfos, functionSetInfos);
+      functionSetGetInfos(*module);
+    FieldInfoScanner scanner(functionNewInfos, functionSetGetInfos);
     scanner.run(runner, module);
     scanner.runOnModuleCode(runner, module);
 
     // Combine the data from the functions.
-    StructValuesMap<FieldInfo> combinedNewInfos, combinedSetInfos;
-    functionSetInfos.combineInto(combinedSetInfos);
+    functionSetGetInfos.combineInto(combinedSetGetInfos);
     // 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.
+    // Propagate information to super and subtypes on set/get infos:
+    //
+    //  * For removing unread fields, we can only remove a field if it is never
+    //    read in any sub or supertype, as such a read may alias any of those
+    //    types (where the field is present).
     //
-    // 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.
+    //  * For immutability, 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;
+    propagator.propagateToSuperAndSubTypes(combinedSetGetInfos);
 
+    // Process the propagated info.
     for (auto type : propagator.subTypes.types) {
       if (!type.isStruct()) {
         continue;
       }
-
       auto& fields = type.getStruct().fields;
+      auto& infos = combinedSetGetInfos[type];
+
+      // Process immutability.
       for (Index i = 0; i < fields.size(); i++) {
         if (fields[i].mutable_ == Immutable) {
           // Already immutable; nothing to do.
           continue;
         }
 
-        if (combinedSetInfos[type][i].hasWrite) {
+        if (infos[i].hasWrite) {
           // A set exists.
           continue;
         }
@@ -143,33 +174,195 @@ struct GlobalTypeOptimization : public Pass {
         vec.resize(i + 1);
         vec[i] = true;
       }
+
+      // Process removability. First, see if we can remove anything before we
+      // start to allocate info for that.
+      if (std::any_of(infos.begin(), infos.end(), [&](const FieldInfo& info) {
+            return !info.hasRead;
+          })) {
+        auto& indexesAfterRemoval = indexesAfterRemovals[type];
+        indexesAfterRemoval.resize(fields.size());
+        Index skip = 0;
+        for (Index i = 0; i < fields.size(); i++) {
+          if (infos[i].hasRead) {
+            indexesAfterRemoval[i] = i - skip;
+          } else {
+            indexesAfterRemoval[i] = RemovedField;
+            skip++;
+          }
+        }
+      }
     }
 
-    // The types are now generally correct, except for their internals, which we
-    // rewrite now.
+    // If we found fields that can be removed, remove them from instructions.
+    // (Note that we must do this first, while we still have the old heap types
+    // that we can identify, and only after this should we update all the types
+    // throughout the module.)
+    if (!indexesAfterRemovals.empty()) {
+      removeFieldsInInstructions(runner, *module);
+    }
+
+    // Update the types in the entire module.
+    updateTypes(*module);
+  }
+
+  void updateTypes(Module& wasm) {
     class TypeRewriter : public GlobalTypeRewriter {
-      CanBecomeImmutable& canBecomeImmutable;
+      GlobalTypeOptimization& parent;
 
     public:
-      TypeRewriter(Module& wasm, CanBecomeImmutable& canBecomeImmutable)
-        : GlobalTypeRewriter(wasm), canBecomeImmutable(canBecomeImmutable) {}
+      TypeRewriter(Module& wasm, GlobalTypeOptimization& parent)
+        : GlobalTypeRewriter(wasm), parent(parent) {}
 
       virtual void modifyStruct(HeapType oldStructType, Struct& struct_) {
-        if (!canBecomeImmutable.count(oldStructType)) {
+        auto& newFields = struct_.fields;
+
+        // Adjust immutability.
+        auto immIter = parent.canBecomeImmutable.find(oldStructType);
+        if (immIter != parent.canBecomeImmutable.end()) {
+          auto& immutableVec = immIter->second;
+          for (Index i = 0; i < immutableVec.size(); i++) {
+            if (immutableVec[i]) {
+              newFields[i].mutable_ = Immutable;
+            }
+          }
+        }
+
+        // Remove fields where we can.
+        auto remIter = parent.indexesAfterRemovals.find(oldStructType);
+        if (remIter != parent.indexesAfterRemovals.end()) {
+          auto& indexesAfterRemoval = remIter->second;
+          Index removed = 0;
+          for (Index i = 0; i < newFields.size(); i++) {
+            auto newIndex = indexesAfterRemoval[i];
+            if (newIndex != RemovedField) {
+              newFields[newIndex] = newFields[i];
+            } else {
+              removed++;
+            }
+          }
+          newFields.resize(newFields.size() - removed);
+
+          // Update field names as well. The Type Rewriter cannot do this for
+          // us, as it does not know which old fields map to which new ones (it
+          // just keeps the names in sequence).
+          auto iter = wasm.typeNames.find(oldStructType);
+          if (iter != wasm.typeNames.end()) {
+            auto& nameInfo = iter->second;
+
+            // Make a copy of the old ones to base ourselves off of as we do so.
+            auto oldFieldNames = nameInfo.fieldNames;
+
+            // Clear the old names and write the new ones.
+            nameInfo.fieldNames.clear();
+            for (Index i = 0; i < oldFieldNames.size(); i++) {
+              auto newIndex = indexesAfterRemoval[i];
+              if (newIndex != RemovedField && oldFieldNames.count(i)) {
+                assert(oldFieldNames[i].is());
+                nameInfo.fieldNames[newIndex] = oldFieldNames[i];
+              }
+            }
+          }
+        }
+      }
+    };
+
+    TypeRewriter(wasm, *this).update();
+  }
+
+  // After updating the types to remove certain fields, we must also remove
+  // them from struct instructions.
+  void removeFieldsInInstructions(PassRunner* runner, Module& wasm) {
+    struct FieldRemover : public WalkerPass<PostWalker<FieldRemover>> {
+      bool isFunctionParallel() override { return true; }
+
+      GlobalTypeOptimization& parent;
+
+      FieldRemover(GlobalTypeOptimization& parent) : parent(parent) {}
+
+      FieldRemover* create() override { return new FieldRemover(parent); }
+
+      void visitStructNew(StructNew* curr) {
+        if (curr->type == Type::unreachable) {
+          return;
+        }
+        if (curr->isWithDefault()) {
+          // Nothing to do, a default was written and will no longer be.
           return;
         }
 
-        auto& newFields = struct_.fields;
-        auto& immutableVec = canBecomeImmutable[oldStructType];
-        for (Index i = 0; i < immutableVec.size(); i++) {
-          if (immutableVec[i]) {
-            newFields[i].mutable_ = Immutable;
+        auto iter = parent.indexesAfterRemovals.find(curr->type.getHeapType());
+        if (iter == parent.indexesAfterRemovals.end()) {
+          return;
+        }
+        auto& indexesAfterRemoval = iter->second;
+
+        auto& operands = curr->operands;
+        assert(indexesAfterRemoval.size() == operands.size());
+
+        // Remove the unneeded operands.
+        Index removed = 0;
+        for (Index i = 0; i < operands.size(); i++) {
+          auto newIndex = indexesAfterRemoval[i];
+          if (newIndex != RemovedField) {
+            assert(newIndex < operands.size());
+            operands[newIndex] = operands[i];
+          } else {
+            if (EffectAnalyzer(getPassOptions(), *getModule(), operands[i])
+                  .hasUnremovableSideEffects()) {
+              Fatal() << "TODO: handle side effects in field removal "
+                         "(impossible in global locations?)";
+            }
+            removed++;
           }
         }
+        operands.resize(operands.size() - removed);
+      }
+
+      void visitStructSet(StructSet* curr) {
+        if (curr->ref->type == Type::unreachable) {
+          return;
+        }
+
+        auto newIndex = getNewIndex(curr->ref->type.getHeapType(), curr->index);
+        if (newIndex != RemovedField) {
+          // Map to the new index.
+          curr->index = newIndex;
+        } else {
+          // This field was removed, so just emit drops of our children.
+          Builder builder(*getModule());
+          replaceCurrent(builder.makeSequence(builder.makeDrop(curr->ref),
+                                              builder.makeDrop(curr->value)));
+        }
+      }
+
+      void visitStructGet(StructGet* curr) {
+        if (curr->ref->type == Type::unreachable) {
+          return;
+        }
+
+        auto newIndex = getNewIndex(curr->ref->type.getHeapType(), curr->index);
+        // We must not remove a field that is read from.
+        assert(newIndex != RemovedField);
+        curr->index = newIndex;
+      }
+
+      Index getNewIndex(HeapType type, Index index) {
+        auto iter = parent.indexesAfterRemovals.find(type);
+        if (iter == parent.indexesAfterRemovals.end()) {
+          return index;
+        }
+        auto& indexesAfterRemoval = iter->second;
+        auto newIndex = indexesAfterRemoval[index];
+        assert(newIndex < indexesAfterRemoval.size() ||
+               newIndex == RemovedField);
+        return newIndex;
       }
     };
 
-    TypeRewriter(*module, canBecomeImmutable).update();
+    FieldRemover remover(*this);
+    remover.run(runner, &wasm);
+    remover.runOnModuleCode(runner, &wasm);
   }
 };
 
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index 67087c328..fd31dada8 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -524,9 +524,9 @@ void PassRunner::addDefaultGlobalOptimizationPrePasses() {
   }
   if (wasm->features.hasGC() && getTypeSystem() == TypeSystem::Nominal &&
       options.optimizeLevel >= 2) {
+    // TODO: investigate enabling --gto and --remove-module-elements before cfp
     addIfNoDWARFIssues("cfp");
   }
-  // TODO: investigate enabling --gto
 }
 
 void PassRunner::addDefaultGlobalOptimizationPostPasses() {