GlobalTypeOptimization: Reorder fields in order to remove them (#6820)

Before, we only removed fields from the end of a struct. If we had, say

struct Foo {
  int x;
  int y;
  int z;
};

// Add no fields but inherit the parent's.
struct Bar : Foo {};

If y is only used in Bar, but never Foo, then we still kept it around, because
if we removed it from Foo we'd end up with Foo = {x, z}, Bar = {x, y, z} which
is invalid - Bar no longer extends Foo. But we can do this if we first reorder
the two:

struct Foo {
  int x;
  int z;
  int y; // now y is at the end
};

struct Bar : Foo {};

And the optimized form is

struct Foo {
  int x;
  int z;
};

struct Bar : Foo {
  int y; // now y is added in Bar
};

This lets us remove all fields possible in all cases AFAIK.

This situation is not super-common, as most fields are actually used both
up and down the hierarchy (if they are used at all), but testing on some
large real-world codebases, I see 10 fields removed in Java, 45 in Kotlin,
and 31 in Dart testcases.

The NFC change to src/wasm-type-ordering.h was needed for this to
compile.

3 files changed, 151 insertions, 55 deletions

diff --git a/src/ir/localize.h b/src/ir/localize.h
index 85e4415f5..b36fe6396 100644
--- a/src/ir/localize.h
+++ b/src/ir/localize.h
@@ -130,6 +130,8 @@ struct ChildLocalizer {
       // effects we can't remove, or if it interacts with other children.
       bool needLocal = effects[i].hasUnremovableSideEffects();
       if (!needLocal) {
+        // TODO: Avoid quadratic time here by accumulating effects and checking
+        //       vs the accumulation.
         for (Index j = 0; j < num; j++) {
           if (j != i && effects[i].invalidates(effects[j])) {
             needLocal = true;
diff --git a/src/passes/GlobalTypeOptimization.cpp b/src/passes/GlobalTypeOptimization.cpp
index 536213561..eec8e206d 100644
--- a/src/passes/GlobalTypeOptimization.cpp
+++ b/src/passes/GlobalTypeOptimization.cpp
@@ -29,7 +29,9 @@
 #include "ir/type-updating.h"
 #include "ir/utils.h"
 #include "pass.h"
+#include "support/permutations.h"
 #include "wasm-builder.h"
+#include "wasm-type-ordering.h"
 #include "wasm-type.h"
 #include "wasm.h"
 
@@ -160,19 +162,23 @@ struct GlobalTypeOptimization : public Pass {
     //    immutable). Note that by making more things immutable we therefore
     //    make it possible to apply more specific subtypes in subtype fields.
     StructUtils::TypeHierarchyPropagator<FieldInfo> propagator(*module);
-    auto subSupers = combinedSetGetInfos;
-    propagator.propagateToSuperAndSubTypes(subSupers);
-    auto subs = std::move(combinedSetGetInfos);
-    propagator.propagateToSubTypes(subs);
-
-    // Process the propagated info.
-    for (auto type : propagator.subTypes.types) {
+    auto dataFromSubsAndSupersMap = combinedSetGetInfos;
+    propagator.propagateToSuperAndSubTypes(dataFromSubsAndSupersMap);
+    auto dataFromSupersMap = std::move(combinedSetGetInfos);
+    propagator.propagateToSubTypes(dataFromSupersMap);
+
+    // Process the propagated info. We look at supertypes first, as the order of
+    // fields in a supertype is a constraint on what subtypes can do. That is,
+    // we decide for each supertype what the optimal order is, and consider that
+    // fixed, and then subtypes can decide how to sort fields that they append.
+    HeapTypeOrdering::SupertypesFirst sorted;
+    for (auto type : sorted.sort(propagator.subTypes.types)) {
       if (!type.isStruct()) {
         continue;
       }
       auto& fields = type.getStruct().fields;
-      auto& subSuper = subSupers[type];
-      auto& sub = subs[type];
+      auto& dataFromSubsAndSupers = dataFromSubsAndSupersMap[type];
+      auto& dataFromSupers = dataFromSupersMap[type];
 
       // Process immutability.
       for (Index i = 0; i < fields.size(); i++) {
@@ -181,7 +187,7 @@ struct GlobalTypeOptimization : public Pass {
           continue;
         }
 
-        if (subSuper[i].hasWrite) {
+        if (dataFromSubsAndSupers[i].hasWrite) {
           // A set exists.
           continue;
         }
@@ -192,48 +198,132 @@ struct GlobalTypeOptimization : public Pass {
         vec[i] = true;
       }
 
-      // Process removability. We check separately for the ability to
-      // remove in a general way based on sub+super-propagated info (that is,
-      // fields that are not used in sub- or super-types, and so we can
-      // definitely remove them from all the relevant types) and also in the
-      // specific way that only works for removing at the end, which as
-      // mentioned above only looks at super-types.
+      // Process removability.
       std::set<Index> removableIndexes;
       for (Index i = 0; i < fields.size(); i++) {
-        if (!subSuper[i].hasRead) {
-          removableIndexes.insert(i);
-        }
-      }
-      for (int i = int(fields.size()) - 1; i >= 0; i--) {
-        // Unlike above, a write would stop us here: above we propagated to both
-        // sub- and super-types, which means if we see no reads then there is no
-        // possible read of the data at all. But here we just propagated to
-        // subtypes, and so we need to care about the case where the parent
-        // writes to a field but does not read from it - we still need those
-        // writes to happen as children may read them. (Note that if no child
-        // reads this field, and since we check for reads in parents here, that
-        // means the field is not read anywhere at all, and we would have
-        // handled that case in the previous loop anyhow.)
-        if (!sub[i].hasRead && !sub[i].hasWrite) {
+        // If there is no read whatsoever, in either subs or supers, then we can
+        // remove the field. That is so even if there are writes (it would be a
+        // pointless "write-only field").
+        auto hasNoReadsAnywhere = !dataFromSubsAndSupers[i].hasRead;
+
+        // Check for reads or writes in ourselves and our supers. If there are
+        // none, then operations only happen in our strict subtypes, and those
+        // subtypes can define the field there, and we don't need it here.
+        auto hasNoReadsOrWritesInSupers =
+          !dataFromSupers[i].hasRead && !dataFromSupers[i].hasWrite;
+
+        if (hasNoReadsAnywhere || hasNoReadsOrWritesInSupers) {
           removableIndexes.insert(i);
-        } else {
-          // Once we see something we can't remove, we must stop, as we can only
-          // remove from the end in this case.
-          break;
         }
       }
-      if (!removableIndexes.empty()) {
-        auto& indexesAfterRemoval = indexesAfterRemovals[type];
-        indexesAfterRemoval.resize(fields.size());
-        Index skip = 0;
-        for (Index i = 0; i < fields.size(); i++) {
-          if (!removableIndexes.count(i)) {
-            indexesAfterRemoval[i] = i - skip;
+
+      // We need to compute the new set of indexes if we are removing fields, or
+      // if our parent removed fields. In the latter case, our parent may have
+      // reordered fields even if we ourselves are not removing anything, and we
+      // must update to match the parent's order.
+      auto super = type.getDeclaredSuperType();
+      auto superHasUpdates = super && indexesAfterRemovals.count(*super);
+      if (!removableIndexes.empty() || superHasUpdates) {
+        // We are removing fields. Reorder them to allow that, as in the general
+        // case we can only remove fields from the end, so that if our subtypes
+        // still need the fields they can append them. For example:
+        //
+        //  type A     = { x: i32, y: f64 };
+        //  type B : A = { x: 132, y: f64, z: v128 };
+        //
+        // If field x is used in B but never in A then we want to remove it, but
+        // we cannot end up with this:
+        //
+        //  type A     = { y: f64 };
+        //  type B : A = { x: 132, y: f64, z: v128 };
+        //
+        // Here B no longer extends A's fields. Instead, we reorder A, which
+        // then imposes the same order on B's fields:
+        //
+        //  type A     = { y: f64, x: i32 };
+        //  type B : A = { y: f64, x: i32, z: v128 };
+        //
+        // And after that, it is safe to remove x in A: B will then append it,
+        // just like it appends z, leading to this:
+        //
+        //  type A     = { y: f64 };
+        //  type B : A = { y: f64, x: i32, z: v128 };
+        //
+        std::vector<Index> indexesAfterRemoval(fields.size());
+
+        // The next new index to use.
+        Index next = 0;
+
+        // If we have a super, then we extend it, and must match its fields.
+        // That is, we can only append fields: we cannot reorder or remove any
+        // field that is in the super.
+        Index numSuperFields = 0;
+        if (super) {
+          // We have visited the super before. Get the information about its
+          // fields.
+          std::vector<Index> superIndexes;
+          auto iter = indexesAfterRemovals.find(*super);
+          if (iter != indexesAfterRemovals.end()) {
+            superIndexes = iter->second;
           } else {
+            // We did not store any information about the parent, because we
+            // found nothing to optimize there. That means it is not removing or
+            // reordering anything, so its new indexes are trivial.
+            superIndexes = makeIdentity(super->getStruct().fields.size());
+          }
+
+          numSuperFields = superIndexes.size();
+
+          // Fields we keep but the super removed will be handled at the end.
+          std::vector<Index> keptFieldsNotInSuper;
+
+          // Go over the super fields and handle them.
+          for (Index i = 0; i < superIndexes.size(); ++i) {
+            auto superIndex = superIndexes[i];
+            if (superIndex == RemovedField) {
+              if (removableIndexes.count(i)) {
+                // This was removed in the super, and in us as well.
+                indexesAfterRemoval[i] = RemovedField;
+              } else {
+                // This was removed in the super, but we actually need it. It
+                // must appear after all other super fields, when we get to the
+                // proper index for that, later. That is, we are reordering.
+                keptFieldsNotInSuper.push_back(i);
+              }
+            } else {
+              // The super kept this field, so we must keep it as well.
+              assert(!removableIndexes.count(i));
+              // We need to keep it at the same index so we remain compatible.
+              indexesAfterRemoval[i] = superIndex;
+              // Update |next| to refer to the next available index. Due to
+              // possible reordering in the parent, we may not see indexes in
+              // order here, so just take the max at each point in time.
+              next = std::max(next, superIndex + 1);
+            }
+          }
+
+          // Handle fields we keep but the super removed.
+          for (auto i : keptFieldsNotInSuper) {
+            indexesAfterRemoval[i] = next++;
+          }
+        }
+
+        // Go over the fields only defined in us, and not in any super.
+        for (Index i = numSuperFields; i < fields.size(); ++i) {
+          if (removableIndexes.count(i)) {
             indexesAfterRemoval[i] = RemovedField;
-            skip++;
+          } else {
+            indexesAfterRemoval[i] = next++;
           }
         }
+
+        // Only store the new indexes we computed if we found something
+        // interesting. We might not, if e.g. our parent removes fields and we
+        // add them back in the exact order we started with. In such cases,
+        // avoid wasting memory and also time later.
+        if (indexesAfterRemoval != makeIdentity(indexesAfterRemoval.size())) {
+          indexesAfterRemovals[type] = indexesAfterRemoval;
+        }
       }
     }
 
@@ -273,15 +363,16 @@ struct GlobalTypeOptimization : public Pass {
           }
         }
 
-        // Remove fields where we can.
+        // Remove/reorder fields where we can.
         auto remIter = parent.indexesAfterRemovals.find(oldStructType);
         if (remIter != parent.indexesAfterRemovals.end()) {
           auto& indexesAfterRemoval = remIter->second;
           Index removed = 0;
+          auto copy = newFields;
           for (Index i = 0; i < newFields.size(); i++) {
             auto newIndex = indexesAfterRemoval[i];
             if (newIndex != RemovedField) {
-              newFields[newIndex] = newFields[i];
+              newFields[newIndex] = copy[i];
             } else {
               removed++;
             }
@@ -347,26 +438,32 @@ struct GlobalTypeOptimization : public Pass {
         auto& operands = curr->operands;
         assert(indexesAfterRemoval.size() == operands.size());
 
-        // Localize things so that we can simply remove the operands we no
-        // longer need.
+        // Ensure any children with non-trivial effects are replaced with
+        // local.gets, so that we can remove/reorder to our hearts' content.
         ChildLocalizer localizer(
           curr, getFunction(), *getModule(), getPassOptions());
         replaceCurrent(localizer.getReplacement());
 
-        // Remove the unneeded operands.
+        // Remove and reorder operands.
         Index removed = 0;
+        std::vector<Expression*> old(operands.begin(), operands.end());
         for (Index i = 0; i < operands.size(); i++) {
           auto newIndex = indexesAfterRemoval[i];
           if (newIndex != RemovedField) {
             assert(newIndex < operands.size());
-            operands[newIndex] = operands[i];
+            operands[newIndex] = old[i];
           } else {
             removed++;
           }
         }
-        operands.resize(operands.size() - removed);
-        // We should only get here if we did actual work.
-        assert(removed > 0);
+        if (removed) {
+          operands.resize(operands.size() - removed);
+        } else {
+          // If we didn't remove anything then we must have reordered (or else
+          // we have done pointless work).
+          assert(indexesAfterRemoval !=
+                 makeIdentity(indexesAfterRemoval.size()));
+        }
       }
 
       void visitStructSet(StructSet* curr) {
diff --git a/src/wasm-type-ordering.h b/src/wasm-type-ordering.h
index 72b1818f7..981ac004d 100644
--- a/src/wasm-type-ordering.h
+++ b/src/wasm-type-ordering.h
@@ -68,9 +68,6 @@ struct SupertypesFirstBase
 };
 
 struct SupertypesFirst : SupertypesFirstBase<SupertypesFirst> {
-  template<typename T>
-  SupertypesFirst(const T& types) : SupertypesFirstBase(types) {}
-
   std::optional<HeapType> getDeclaredSuperType(HeapType type) {
     return type.getDeclaredSuperType();
   }