[Wasm GC] Update nulls to allow finding better LUBs (#4340)

It is common in GC code to have stuff like this:

x = null;
..
x = Data();

Nulls in wasm have a type, and if that initial null has say anyref then
before this PR we would keep the type of x as anyref. However,
while nulls have types, all null values are identical, and so we can
in fact change x's type to a nullable reference of Data, by also
changing the null's type to something more specific.

LUBFinder now has an API that can return the best possible LUB
so far, and that can be told to update nulls if we decide that the
new LUB is worth using. This updates the passes using LUBFinder
to use the new API. Note how TypeRefining becomes simpler
because the special logic it had in a subclass of LUBFinder is now
part of the main class (it used to remember if there was a null
default; LUBFinder now handles both a null default as well as
other nulls).

This requires some changes to existing tests to avoid them from
optimizing using nulls in ways that ends up not testing the
original intent. Specifically the dae-gc-refine-params.wast now
has calls to get a null of a type, instead of just having a ref.null
of that type (which could be optimized now). And
dae-gc-refine-return uses locals instead of ref.nulls.

4 files changed, 142 insertions, 75 deletions

diff --git a/src/ir/lubs.h b/src/ir/lubs.h
index a45b62c52..e03e693e0 100644
--- a/src/ir/lubs.h
+++ b/src/ir/lubs.h
@@ -22,31 +22,110 @@
 
 namespace wasm {
 
+//
 // Helper to find a LUB of a series of expressions. This works incrementally so
 // that if we see we are not improving on an existing type then we can stop
-// early.
+// early. It also notes null expressions that can be updated later, and if
+// updating them would allow a better LUB it can do so. That is, given this:
+//
+//   (ref.null any)  ;; an expression that we can update
+//   (.. something of type data ..)
+//
+// We can update that null to type (ref null data) which would allow setting
+// that as the LUB. This is important in cases where there is a null initial
+// value in a field, for example: we should not let the type there prevent us
+// from optimizing - after all, all nulls compare equally anyhow.
+//
 struct LUBFinder {
-  // The least upper bound we found so far.
-  Type lub = Type::unreachable;
+  LUBFinder() {}
+
+  LUBFinder(Type initialType) { note(initialType); }
 
-  // Note another type to take into account in the lub. Returns the new lub.
-  Type note(Type type) { return lub = Type::getLeastUpperBound(lub, type); }
+  // Note another type to take into account in the lub.
+  void note(Type type) { lub = Type::getLeastUpperBound(lub, type); }
 
-  Type note(Expression* curr) { return note(curr->type); }
+  // Note an expression that can be updated, that is, that we can modify in
+  // safe ways if doing so would allow us to get a better lub. The specific
+  // optimization possible here involves nulls, see the top comment.
+  void noteUpdatableExpression(Expression* curr) {
+    if (auto* null = curr->dynCast<RefNull>()) {
+      nulls.insert(null);
+    } else {
+      note(curr->type);
+    }
+  }
 
-  // Returns whether we noted any (reachable) value.
+  void noteNullDefault() {
+    // A default value is indicated by a null pointer.
+    nulls.insert(nullptr);
+  }
+
+  // Returns whether we noted any (reachable) value. This ignores nulls, as they
+  // do not contribute type information - we do not try to find a lub based on
+  // them (rather we update them to the LUB).
   bool noted() { return lub != Type::unreachable; }
 
-  // Returns the lub that we found.
-  Type get() { return lub; }
+  // Returns the best possible lub. This ignores updatable nulls for the reasons
+  // mentioned above, since they will not limit us, aside from making the type
+  // nullable if nulls exist. This does not update the nulls.
+  Type getBestPossible() {
+    if (lub == Type::unreachable) {
+      // Perhaps if we have seen nulls we could compute a lub on them, but it's
+      // not clear that is helpful.
+      return lub;
+    }
+
+    // We have a lub. Make it nullable if we need to.
+    if (!lub.isNullable() && !nulls.empty()) {
+      return Type(lub.getHeapType(), Nullable);
+    } else {
+      return lub;
+    }
+  }
+
+  // Update the nulls for the best possible LUB, if we found one.
+  void updateNulls() {
+    auto newType = getBestPossible();
+    if (newType != Type::unreachable) {
+      for (auto* null : nulls) {
+        // Default null values (represented as nullptr here) do not need to be
+        // updated. Aside from that, if this null is already of a more specific
+        // type, also do not update it - it's never worth making a type less
+        // specific. What we care about here is making sure the nulls are all
+        // specific enough given the LUB that is being applied.
+        if (null && !Type::isSubType(null->type, newType)) {
+          null->finalize(newType);
+        }
+      }
+    }
+  }
 
   // Combines the information in another LUBFinder into this one, and returns
   // whether we changed anything.
   bool combine(const LUBFinder& other) {
+    // Check if the lub was changed.
     auto old = lub;
     note(other.lub);
-    return lub != old;
+    bool changed = old != lub;
+
+    // Check if we added new updatable nulls.
+    for (auto* null : other.nulls) {
+      if (nulls.insert(null).second) {
+        changed = true;
+      }
+    }
+
+    return changed;
   }
+
+private:
+  // The least upper bound. As we go this always contains the latest value based
+  // on everything we've seen so far, except for nulls.
+  Type lub = Type::unreachable;
+
+  // Nulls that we can update. A nullptr here indicates an "implicit" null, that
+  // is, a null default value.
+  std::unordered_set<RefNull*> nulls;
 };
 
 } // namespace wasm
diff --git a/src/passes/DeadArgumentElimination.cpp b/src/passes/DeadArgumentElimination.cpp
index e997c591f..b7ce5ef5f 100644
--- a/src/passes/DeadArgumentElimination.cpp
+++ b/src/passes/DeadArgumentElimination.cpp
@@ -553,6 +553,7 @@ private:
     auto numParams = func->getNumParams();
     std::vector<Type> newParamTypes;
     newParamTypes.reserve(numParams);
+    std::vector<LUBFinder> lubs(numParams);
     for (Index i = 0; i < numParams; i++) {
       auto originalType = func->getLocalType(i);
       // If the parameter type is not a reference, there is nothing to refine.
@@ -565,10 +566,11 @@ private:
         newParamTypes.push_back(originalType);
         continue;
       }
-      LUBFinder lub;
+      auto& lub = lubs[i];
       for (auto* call : calls) {
         auto* operand = call->operands[i];
-        if (lub.note(operand) == originalType) {
+        lub.noteUpdatableExpression(operand);
+        if (lub.getBestPossible() == originalType) {
           // We failed to refine this parameter to anything more specific.
           break;
         }
@@ -578,7 +580,7 @@ private:
       if (!lub.noted()) {
         return;
       }
-      newParamTypes.push_back(lub.get());
+      newParamTypes.push_back(lub.getBestPossible());
     }
 
     // Check if we are able to optimize here before we do the work to scan the
@@ -591,6 +593,11 @@ private:
     // We can do this!
     TypeUpdating::updateParamTypes(func, newParamTypes, *module);
 
+    // Update anything the lubs need to update.
+    for (auto& lub : lubs) {
+      lub.updateNulls();
+    }
+
     // Also update the function's type.
     func->setParams(newParams);
   }
@@ -633,21 +640,28 @@ private:
     ReFinalize().walkFunctionInModule(func, module);
 
     LUBFinder lub;
-    if (lub.note(func->body) == originalType) {
+    lub.noteUpdatableExpression(func->body);
+    if (lub.getBestPossible() == originalType) {
       return false;
     }
 
-    // Scan the body and look at the returns.
-    auto processReturnType = [&](Type type) {
-      // Return whether we still look ok to do the optimization. If this is
-      // false then we can stop here.
-      return lub.note(type) != originalType;
-    };
+    // Scan the body and look at the returns. First, return expressions.
     for (auto* ret : FindAll<Return>(func->body).list) {
-      if (!processReturnType(ret->value->type)) {
+      lub.noteUpdatableExpression(ret->value);
+      if (lub.getBestPossible() == originalType) {
         return false;
       }
     }
+
+    // Process return_calls and call_refs. Unlike return expressions which we
+    // just handled, these only get a type to update, not a value.
+    auto processReturnType = [&](Type type) {
+      // Return whether we still look ok to do the optimization. If this is
+      // false then we can stop here.
+      lub.note(type);
+      return lub.getBestPossible() != originalType;
+    };
+
     for (auto* call : FindAll<Call>(func->body).list) {
       if (call->isReturn &&
           !processReturnType(module->getFunction(call->target)->getResults())) {
@@ -671,7 +685,6 @@ private:
         }
       }
     }
-    assert(lub.get() != originalType);
 
     // If the refined type is unreachable then nothing actually returns from
     // this function.
@@ -680,13 +693,19 @@ private:
       return false;
     }
 
+    auto newType = lub.getBestPossible();
+    if (newType == originalType) {
+      return false;
+    }
+
     // Success. Update the type, and the calls.
-    func->setResults(lub.get());
+    func->setResults(newType);
     for (auto* call : calls) {
       if (call->type != Type::unreachable) {
-        call->type = lub.get();
+        call->type = newType;
       }
     }
+    lub.updateNulls();
     return true;
   }
 };
diff --git a/src/passes/LocalSubtyping.cpp b/src/passes/LocalSubtyping.cpp
index 57119a253..a8174aa2d 100644
--- a/src/passes/LocalSubtyping.cpp
+++ b/src/passes/LocalSubtyping.cpp
@@ -124,7 +124,8 @@ struct LocalSubtyping : public WalkerPass<PostWalker<LocalSubtyping>> {
         // Find all the types assigned to the var, and compute the optimal LUB.
         LUBFinder lub;
         for (auto* set : setsForLocal[i]) {
-          if (lub.note(set->value) == oldType) {
+          lub.noteUpdatableExpression(set->value);
+          if (lub.getBestPossible() == oldType) {
             break;
           }
         }
@@ -133,7 +134,7 @@ struct LocalSubtyping : public WalkerPass<PostWalker<LocalSubtyping>> {
           continue;
         }
 
-        auto newType = lub.get();
+        auto newType = lub.getBestPossible();
         assert(newType != Type::none); // in valid wasm there must be a LUB
 
         // Remove non-nullability if we disallow that in locals.
@@ -157,6 +158,7 @@ struct LocalSubtyping : public WalkerPass<PostWalker<LocalSubtyping>> {
           func->vars[i - varBase] = newType;
           more = true;
           optimized = true;
+          lub.updateNulls();
 
           // Update gets and tees.
           for (auto* get : getsForLocal[i]) {
diff --git a/src/passes/TypeRefining.cpp b/src/passes/TypeRefining.cpp
index 33e382bc5..9480322bb 100644
--- a/src/passes/TypeRefining.cpp
+++ b/src/passes/TypeRefining.cpp
@@ -36,38 +36,7 @@ namespace {
 // possible LUB so far. The only extra functionality we need here is whether
 // there is a default null value (which would force us to keep the type
 // nullable).
-struct FieldInfo : public LUBFinder {
-  bool nullDefault = false;
-
-  void noteNullDefault() { nullDefault = true; }
-
-  bool hasNullDefault() { return nullDefault; }
-
-  bool noted() { return LUBFinder::noted() || nullDefault; }
-
-  Type get() {
-    auto ret = LUBFinder::get();
-    if (nullDefault && !ret.isNullable()) {
-      ret = Type(ret.getHeapType(), Nullable);
-    }
-    return ret;
-  }
-
-  bool combine(const FieldInfo& other) {
-    auto old = nullDefault;
-    if (other.nullDefault) {
-      noteNullDefault();
-    }
-    return LUBFinder::combine(other) || nullDefault != old;
-  }
-
-  // Force the lub to a particular type.
-  void set(Type type) { lub = type; }
-
-  void dump(std::ostream& o) {
-    std::cout << "FieldInfo(" << lub << ", " << nullDefault << ")";
-  }
-};
+using FieldInfo = LUBFinder;
 
 struct FieldInfoScanner
   : public StructUtils::StructScanner<FieldInfo, FieldInfoScanner> {
@@ -85,7 +54,7 @@ struct FieldInfoScanner
                       HeapType type,
                       Index index,
                       FieldInfo& info) {
-    info.note(expr);
+    info.noteUpdatableExpression(expr);
   }
 
   void
@@ -171,9 +140,8 @@ struct TypeRefining : public Pass {
       for (Index i = 0; i < fields.size(); i++) {
         auto oldType = fields[i].type;
         auto& info = finalInfos[type][i];
-        auto newType = info.get();
-        if (newType == Type::unreachable) {
-          info.set(oldType);
+        if (!info.noted()) {
+          info = LUBFinder(oldType);
         }
       }
 
@@ -181,9 +149,9 @@ struct TypeRefining : public Pass {
       if (auto super = type.getSuperType()) {
         auto& superFields = super->getStruct().fields;
         for (Index i = 0; i < superFields.size(); i++) {
-          auto newSuperType = finalInfos[*super][i].get();
+          auto newSuperType = finalInfos[*super][i].getBestPossible();
           auto& info = finalInfos[type][i];
-          auto newType = info.get();
+          auto newType = info.getBestPossible();
           if (!Type::isSubType(newType, newSuperType)) {
             // To ensure we are a subtype of the super's field, simply copy that
             // value, which is more specific than us.
@@ -201,7 +169,7 @@ struct TypeRefining : public Pass {
             // to something more specific than $C's old type, we end up with the
             // problem that this code path fixes: we just need to get $C's type
             // to be identical to its super so that validation works.
-            info.set(newSuperType);
+            info = LUBFinder(newSuperType);
           } else if (fields[i].mutable_ == Mutable) {
             // Mutable fields must have identical types, so we cannot
             // specialize.
@@ -209,20 +177,19 @@ struct TypeRefining : public Pass {
             //       here? This entire analysis might be done on fields, and not
             //       types, which would also handle more things added to fields
             //       in the future.
-            info.set(newSuperType);
+            info = LUBFinder(newSuperType);
           }
         }
       }
 
       // After all those decisions, see if we found anything to optimize.
-      if (!canOptimize) {
-        for (Index i = 0; i < fields.size(); i++) {
-          auto oldType = fields[i].type;
-          auto newType = finalInfos[type][i].get();
-          if (newType != oldType) {
-            canOptimize = true;
-            break;
-          }
+      for (Index i = 0; i < fields.size(); i++) {
+        auto oldType = fields[i].type;
+        auto& lub = finalInfos[type][i];
+        auto newType = lub.getBestPossible();
+        if (newType != oldType) {
+          canOptimize = true;
+          lub.updateNulls();
         }
       }
 
@@ -253,7 +220,7 @@ struct TypeRefining : public Pass {
           if (!oldType.isRef()) {
             continue;
           }
-          auto newType = parent.finalInfos[oldStructType][i].get();
+          auto newType = parent.finalInfos[oldStructType][i].getBestPossible();
           newFields[i].type = getTempType(newType);
         }
       }