[NFC] Use principal types in ChildTyper (#7816)

Refactor ChildTyper to report type constraints using a single callback
that takes a sequence of VarTypes from principal-type.h rather than
using several callbacks for the various kinds of constraints. The
previous approach using several callbacks did not scale well to the
various kinds of polymorphism introduced in new Wasm proposals, such
as polymorphism over sharedness in `ref.eq` or polymorphism over
exactness in `ref.get_desc`.

NFC despite better modeling these contraints because ChildTyper is so
far only used from TypeSSA, which only uses it to look for exact types,
and from IRBuilder, which only uses it to control the construction of
unreachable IR. Follow-on changes will use ChildTyper in new ways that
require this improved precision.

Author: tlively

src/ir/child-typer.h

@@ -16,39 +16,28 @@
 #ifndef wasm_ir_child_typer_h
 #define wasm_ir_child_typer_h
 
+#include "ir/principal-type.h"
 #include "wasm-traversal.h"
 #include "wasm.h"
 
 namespace wasm {
 
-// CRTP visitor for determining constraints on the types of expression children.
-// For each child of the visited expression, calls a callback with a pointer to
-// the child and information on how the child is constrained. The possible
-// callbacks are:
+// CRTP visitor for determining constaints on the types of expression children.
+// For each child of the visited expression, calls a callback with the VarTypes
+// giving the constraint on the child:
 //
-//   noteSubtype(Expression** childp, Type type) - The child must be a subtype
-//   of `type`, which may be a tuple type. For children that must not produce
-//   values, this may be `Type::none`. This accounts for most type constraints.
+//   note(Expression** childp, SmallVector<VarType, 1> type)
 //
-//   noteAnyType(Expression** childp) - The child may have any non-tuple type.
-//   Used for the children of polymorphic instructions like `drop` and `select`.
-//
-//   noteAnyReferenceType(Expression** childp) - The child may have any
-//   reference type. Used for the children of polymorphic reference instructions
-//   like `ref.is_null`.
-//
-//   noteAnyTupleType(Expression** childp, size_t arity) - The child may have
-//   any tuple type with the given arity. Used for the children of polymorphic
-//   tuple instructions like `tuple.drop` and `tuple.extract`.
+// Multiple VarTypes are provided if and only if the child must be a tuple. In
+// that case, each vartype gives the constraint on a single element of the
+// tuple.
 //
 // Subclasses must additionally implement a callback for getting the type of a
-// branch target. This callback will only be used when the label type is not
+// branch target. This callback will only be used when a the label type is not
 // passed directly as an argument to the branch visitor method (see below).
 //
 //   Type getLabelType(Name label)
 //
-// Children with type `unreachable` satisfy all constraints.
-//
 // Constraints are determined using information that would be present in the
 // binary, e.g. type annotation immediates. Many of the visitor methods take
 // optional additional parameter for passing this information directly, and if
@@ -70,33 +59,49 @@ namespace wasm {
 // information to generate constraints for all the children. Some users may wish
 // to ignore this situation and others may want to assert that it never occurs.
 template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
+  using Constraints = SmallVector<VarType, 1>;
   Module& wasm;
   Function* func;
 
   ChildTyper(Module& wasm, Function* func) : wasm(wasm), func(func) {}
 
   Subtype& self() { return *static_cast<Subtype*>(this); }
 
-  void note(Expression** childp, Type type) {
-    self().noteSubtype(childp, type);
+  void note(Expression** childp, VarType type) {
+    self().note(childp, Constraints{type});
   }
 
-  void notePointer(Expression** ptrp, Name mem) {
-    note(ptrp, wasm.getMemory(mem)->addressType);
+  void note(Expression** childp, Type type) {
+    if (type.isTuple()) {
+      Constraints tuple;
+      for (auto t : type.getTuple()) {
+        tuple.push_back(t);
+      }
+      self().note(childp, tuple);
+    } else {
+      note(childp, VarType{type});
+    }
   }
 
-  void noteTableIndex(Expression** indexp, Name table) {
-    note(indexp, wasm.getTable(table)->addressType);
+  // Disambiguate betwween Type and VarType.
+  void note(Expression** childp, Type::BasicType type) {
+    note(childp, VarType{Type(type)});
   }
 
-  void noteAny(Expression** childp) { self().noteAnyType(childp); }
+  void noteAnyTuple(Expression** childp, size_t n) {
+    Constraints tuple;
+    for (; n != 0; --n) {
+      tuple.push_back(VarType{Index(n - 1)});
+    }
+    self().note(childp, tuple);
+  }
 
-  void noteAnyReference(Expression** childp) {
-    self().noteAnyReferenceType(childp);
+  void notePointer(Expression** ptrp, Name mem) {
+    note(ptrp, wasm.getMemory(mem)->addressType);
   }
 
-  void noteAnyTuple(Expression** childp, size_t arity) {
-    self().noteAnyTupleType(childp, arity);
+  void noteTableIndex(Expression** indexp, Name table) {
+    note(indexp, wasm.getTable(table)->addressType);
   }
 
   Type getLabelType(Name label) { return self().getLabelType(label); }
@@ -115,6 +120,7 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
   }
 
   void visitIf(If* curr) {
+    // TODO: if the condition is unreachable, use a type variable for the arms?
     note(&curr->condition, Type::i32);
     note(&curr->ifTrue, curr->type);
     if (curr->ifFalse) {
@@ -200,8 +206,6 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
 
   void visitAtomicCmpxchg(AtomicCmpxchg* curr,
                           std::optional<Type> type = std::nullopt) {
-    assert(!type || *type == Type::i32 || *type == Type::i64);
-    notePointer(&curr->ptr, curr->memory);
     if (!type) {
       if (curr->expected->type == Type::i64 ||
           curr->replacement->type == Type::i64) {
@@ -210,6 +214,8 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
         type = Type::i32;
       }
     }
+    assert(*type == Type::i32 || *type == Type::i64);
+    notePointer(&curr->ptr, curr->memory);
     note(&curr->expected, *type);
     note(&curr->replacement, *type);
   }
@@ -701,8 +707,10 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
       note(&curr->ifTrue, *type);
       note(&curr->ifFalse, *type);
     } else {
-      noteAny(&curr->ifTrue);
-      noteAny(&curr->ifFalse);
+      // Polymorphic over types.
+      // TODO: Model the constraint that this must be a numeric or vector type.
+      note(&curr->ifTrue, VarType{0u});
+      note(&curr->ifFalse, VarType{0u});
     }
     note(&curr->condition, Type::i32);
   }
@@ -711,11 +719,9 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
     if (!arity) {
       arity = curr->value->type.size();
     }
-    if (*arity > 1) {
-      noteAnyTuple(&curr->value, *arity);
-    } else {
-      noteAny(&curr->value);
-    }
+    assert(*arity > 0);
+    // Might not actually be a tuple, but works either way.
+    noteAnyTuple(&curr->value, *arity);
   }
 
   void visitReturn(Return* curr) {
@@ -738,14 +744,19 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
 
   void visitRefNull(RefNull* curr) {}
 
-  void visitRefIsNull(RefIsNull* curr) { noteAnyReference(&curr->value); }
+  void visitRefIsNull(RefIsNull* curr) {
+    // Polymorphic over heap types.
+    note(&curr->value, VarRef{Nullable, VarHeapType{0u}});
+  }
 
   void visitRefFunc(RefFunc* curr) {}
 
   void visitRefEq(RefEq* curr) {
-    Type eqref(HeapType::eq, Nullable);
-    note(&curr->left, eqref);
-    note(&curr->right, eqref);
+    // Polymorphic over sharedness.
+    VarRef maybeShareEqref{Nullable,
+                           VarAbsHeapType{VarSharedness{0u}, HeapType::eq}};
+    note(&curr->left, maybeShareEqref);
+    note(&curr->right, maybeShareEqref);
   }
 
   void visitTableGet(TableGet* curr) {
@@ -815,8 +826,8 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
   }
 
   void visitTupleMake(TupleMake* curr) {
-    for (auto& expr : curr->operands) {
-      noteAny(&expr);
+    for (Index i = 0; i < curr->operands.size(); ++i) {
+      note(&curr->operands[i], VarType{Index(curr->operands.size() - i - 1)});
     }
   }
 
@@ -848,10 +859,7 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
       ht = curr->target->type.getHeapType().getSignature();
     }
     auto params = ht->getSignature().params;
-    assert(curr->operands.size() == params.size());
-    for (size_t i = 0; i < params.size(); ++i) {
-      note(&curr->operands[i], params[i]);
-    }
+    handleCall(curr, params);
     note(&curr->target, Type(*ht, Nullable));
   }
 
@@ -886,7 +894,9 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
       }
       ht = curr->ref->type.getHeapType();
     }
-    note(&curr->ref, Type(*ht, Nullable));
+    // Polymorphic over exactness.
+    assert(!ht->isBasic());
+    note(&curr->ref, VarRef{Nullable, VarDefHeapType{VarExactness{0u}, *ht}});
   }
 
   void visitBrOn(BrOn* curr, std::optional<Type> target = std::nullopt) {
@@ -896,7 +906,8 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
         // br_on(_non)_null is polymorphic over reference types and does not
         // take a type immediate.
         assert(!target);
-        noteAnyReference(&curr->ref);
+        // Polymorphic over heap types.
+        note(&curr->ref, VarRef{Nullable, VarHeapType{0u}});
         return;
       case BrOnCast:
       case BrOnCastFail:
@@ -931,8 +942,7 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
         note(&curr->operands[i], fields[i].type);
       }
     }
-    auto desc = curr->type.getHeapType().getDescriptorType();
-    if (desc) {
+    if (auto desc = curr->type.getHeapType().getDescriptorType()) {
       note(&curr->desc, Type(*desc, NonNullable, Exact));
     }
   }
@@ -992,6 +1002,7 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
     assert(curr->index < fields.size());
     note(&curr->ref, Type(*ht, Nullable));
     auto type = fields[curr->index].type;
+    // TODO: (shared eq) as appropriate.
     note(&curr->expected, type.isRef() ? Type(HeapType::eq, Nullable) : type);
     note(&curr->replacement, type);
   }
@@ -1154,17 +1165,19 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
       }
       ht = curr->ref->type.getHeapType();
     }
-    auto type = ht->getArray().element.type;
+    Type type = ht->getArray().element.type;
     note(&curr->ref, Type(*ht, Nullable));
     note(&curr->index, Type::i32);
+    // TODO: (shared eq) as appropriate.
     note(&curr->expected, type.isRef() ? Type(HeapType::eq, Nullable) : type);
     note(&curr->replacement, type);
   }
 
   void visitRefAs(RefAs* curr) {
     switch (curr->op) {
       case RefAsNonNull:
-        noteAnyReference(&curr->value);
+        // Polymorphic over heap types.
+        note(&curr->value, VarRef{Nullable, VarHeapType{0u}});
         return;
       case AnyConvertExtern:
         note(&curr->value, Type(HeapType::ext, Nullable));
@@ -1200,6 +1213,7 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
   void visitStringConst(StringConst* curr) {}
 
   void visitStringMeasure(StringMeasure* curr) {
+    // TODO: extern instead of string? For other ops as well.
     note(&curr->ref, Type(HeapType::string, Nullable));
   }
 

src/ir/principal-type.cpp

@@ -1069,6 +1069,11 @@ std::optional<Signature> PrincipalType::getSignature() const {
   return Signature(sigParams, sigResults);
 }
 
+bool PrincipalType::matches(Type type, VarType constraint) {
+  VarAssignments assignments;
+  return match(assignments, VarType{type}, constraint);
+}
+
 std::ostream& operator<<(std::ostream& o, const PrincipalType& type) {
   print(o, type);
   return o;

src/ir/principal-type.h

@@ -130,6 +130,8 @@ struct PrincipalType {
   // variables.
   std::optional<Signature> getSignature() const;
 
+  static bool matches(Type type, VarType constraint);
+
   bool operator==(const PrincipalType& other) const {
     return rparams == other.rparams && results == other.results &&
            unreachable == other.unreachable;

src/passes/TypeSSA.cpp

@@ -48,14 +48,12 @@
 //
 
 #include "ir/child-typer.h"
-#include "ir/find_all.h"
 #include "ir/module-utils.h"
 #include "ir/names.h"
 #include "ir/possible-constant.h"
 #include "ir/utils.h"
 #include "pass.h"
 #include "support/hash.h"
-#include "wasm-builder.h"
 #include "wasm.h"
 
 namespace wasm {
@@ -221,21 +219,19 @@ struct Analyzer
         : ChildTyper(*parent.getModule(), parent.getFunction()),
           parent(parent) {}
 
-      void noteSubtype(Expression**, Type type) {
-        for (Type t : type) {
-          if (t.isExact()) {
-            parent.disallowedTypes.insert(t.getHeapType());
+      void note(Expression**, Constraints type) {
+        // Check closed type constraints for exactness. Other kinds of type
+        // constaints do not concern us.
+        // TODO: Handle tuples?
+        for (auto varType : type) {
+          if (auto* t = std::get_if<Type>(&varType)) {
+            if (t->isExact()) {
+              parent.disallowedTypes.insert(t->getHeapType());
+            }
           }
         }
       }
 
-      // Other constraints do not matter to us.
-      void noteAnyType(Expression**) {}
-      void noteAnyReferenceType(Expression**) {}
-      void noteAnyTupleType(Expression**, size_t) {}
-      void noteAnyI8ArrayReferenceType(Expression**) {}
-      void noteAnyI16ArrayReferenceType(Expression**) {}
-
       Type getLabelType(Name label) { WASM_UNREACHABLE("unexpected branch"); }
 
       // We don't mind if we cannot compute a constraint due to unreachability.