Use standard type traversal in module-utils.h (#3851)

Add new public `getHeapTypeChildren` methods to Type and HeapType, implemented
in using the standard machinery from #3844, and use them to simplify
`ModuleUtils::collectHeapTypes`.

Now that the type traversal code in wasm-type.cpp is not just used in
canonicalization, move it to a more appropriate place in the file. Also, since
the only users of `HeapTypePathWalker` were using it to visit top-level children
only, replace that with a more specialized `HeapTypeChildWalker` to reduce code
duplication.

3 files changed, 267 insertions, 272 deletions

diff --git a/src/ir/module-utils.h b/src/ir/module-utils.h
index 8d778bc3c..ef6bc2bf6 100644
--- a/src/ir/module-utils.h
+++ b/src/ir/module-utils.h
@@ -468,13 +468,14 @@ inline void collectHeapTypes(Module& wasm,
                              std::vector<HeapType>& types,
                              std::unordered_map<HeapType, Index>& typeIndices) {
   struct Counts : public std::unordered_map<HeapType, size_t> {
-    bool isRelevant(Type type) {
-      return (type.isRef() || type.isRtt()) && !type.getHeapType().isBasic();
+    void note(HeapType type) {
+      if (!type.isBasic()) {
+        (*this)[type]++;
+      }
     }
-    void note(HeapType type) { (*this)[type]++; }
-    void maybeNote(Type type) {
-      if (isRelevant(type)) {
-        note(type.getHeapType());
+    void note(Type type) {
+      for (HeapType ht : type.getHeapTypeChildren()) {
+        note(ht);
       }
     }
   };
@@ -489,19 +490,19 @@ inline void collectHeapTypes(Module& wasm,
       if (auto* call = curr->dynCast<CallIndirect>()) {
         counts.note(call->sig);
       } else if (curr->is<RefNull>()) {
-        counts.maybeNote(curr->type);
+        counts.note(curr->type);
       } else if (curr->is<RttCanon>() || curr->is<RttSub>()) {
         counts.note(curr->type.getRtt().heapType);
       } else if (auto* get = curr->dynCast<StructGet>()) {
-        counts.maybeNote(get->ref->type);
+        counts.note(get->ref->type);
       } else if (auto* set = curr->dynCast<StructSet>()) {
-        counts.maybeNote(set->ref->type);
+        counts.note(set->ref->type);
       } else if (Properties::isControlFlowStructure(curr)) {
         if (curr->type.isTuple()) {
           // TODO: Allow control flow to have input types as well
           counts.note(Signature(Type::none, curr->type));
         } else {
-          counts.maybeNote(curr->type);
+          counts.note(curr->type);
         }
       }
     }
@@ -514,10 +515,10 @@ inline void collectHeapTypes(Module& wasm,
     counts.note(curr->sig);
   }
   for (auto& curr : wasm.tables) {
-    counts.maybeNote(curr->type);
+    counts.note(curr->type);
   }
   for (auto& curr : wasm.elementSegments) {
-    counts.maybeNote(curr->type);
+    counts.note(curr->type);
   }
 
   // Collect info from functions in parallel.
@@ -525,9 +526,7 @@ inline void collectHeapTypes(Module& wasm,
     wasm, [&](Function* func, Counts& counts) {
       counts.note(func->sig);
       for (auto type : func->vars) {
-        for (auto t : type) {
-          counts.maybeNote(t);
-        }
+        counts.note(type);
       }
       if (!func->imported()) {
         CodeScanner(counts).walk(func->body);
@@ -542,30 +541,6 @@ inline void collectHeapTypes(Module& wasm,
     }
   }
 
-  // A generic utility to traverse the child types of a type.
-  // TODO: work with tlively to refactor this to a shared place
-  auto walkRelevantChildren = [&](HeapType type, auto callback) {
-    auto callIfRelevant = [&](Type type) {
-      if (counts.isRelevant(type)) {
-        callback(type.getHeapType());
-      }
-    };
-    if (type.isSignature()) {
-      auto sig = type.getSignature();
-      for (Type type : {sig.params, sig.results}) {
-        for (auto element : type) {
-          callIfRelevant(element);
-        }
-      }
-    } else if (type.isArray()) {
-      callIfRelevant(type.getArray().element.type);
-    } else if (type.isStruct()) {
-      auto fields = type.getStruct().fields;
-      for (auto field : fields) {
-        callIfRelevant(field.type);
-      }
-    }
-  };
   // Recursively traverse each reference type, which may have a child type that
   // is itself a reference type. This reflects an appearance in the binary
   // format that is in the type section itself.
@@ -578,14 +553,16 @@ inline void collectHeapTypes(Module& wasm,
   }
   while (!newTypes.empty()) {
     auto iter = newTypes.begin();
-    auto type = *iter;
+    auto ht = *iter;
     newTypes.erase(iter);
-    walkRelevantChildren(type, [&](HeapType type) {
-      if (!counts.count(type)) {
-        newTypes.insert(type);
+    for (HeapType child : ht.getHeapTypeChildren()) {
+      if (!child.isBasic()) {
+        if (!counts.count(child)) {
+          newTypes.insert(child);
+        }
+        counts.note(child);
       }
-      counts.note(type);
-    });
+    }
   }
 
   // Sort by frequency and then simplicity.
diff --git a/src/wasm-type.h b/src/wasm-type.h
index 1b7d389ca..41a89d4fe 100644
--- a/src/wasm-type.h
+++ b/src/wasm-type.h
@@ -204,6 +204,9 @@ public:
   // Returns true if left is a subtype of right. Subtype includes itself.
   static bool isSubType(Type left, Type right);
 
+  // Return the ordered HeapType children, looking through child Types.
+  std::vector<HeapType> getHeapTypeChildren();
+
   // Computes the least upper bound from the type lattice.
   // If one of the type is unreachable, the other type becomes the result. If
   // the common supertype does not exist, returns none, a poison value.
@@ -361,10 +364,14 @@ public:
 
   // Order heap types by some notion of simplicity.
   bool operator<(const HeapType& other) const;
-  std::string toString() const;
 
   // Returns true if left is a subtype of right. Subtype includes itself.
   static bool isSubType(HeapType left, HeapType right);
+
+  // Return the ordered HeapType children, looking through child Types.
+  std::vector<HeapType> getHeapTypeChildren();
+
+  std::string toString() const;
 };
 
 typedef std::vector<Type> TypeList;
diff --git a/src/wasm/wasm-type.cpp b/src/wasm/wasm-type.cpp
index 6f9671fcd..a00b5043b 100644
--- a/src/wasm/wasm-type.cpp
+++ b/src/wasm/wasm-type.cpp
@@ -276,6 +276,132 @@ struct FiniteShapeEquator {
   bool eq(const Rtt& a, const Rtt& b);
 };
 
+// Generic utility for traversing type graphs. The inserted roots must live as
+// long as the Walker because they are referenced by address. This base class
+// only has logic for traversing type graphs; figuring out when to stop
+// traversing the graph and doing useful work during the traversal is left to
+// subclasses.
+template<typename Self> struct TypeGraphWalkerBase {
+  void walkRoot(Type* type);
+  void walkRoot(HeapType* ht);
+
+  // Override these in subclasses to do useful work.
+  void preVisitType(Type* type) {}
+  void preVisitHeapType(HeapType* ht) {}
+  void postVisitType(Type* type) {}
+  void postVisitHeapType(HeapType* ht) {}
+
+  // This base walker does not know when to stop scanning, so at least one of
+  // these needs to be overridden with a method that calls the base scanning
+  // method only if some end condition isn't met.
+  void scanType(Type* type);
+  void scanHeapType(HeapType* ht);
+
+private:
+  struct Task {
+    enum Kind {
+      PreType,
+      PreHeapType,
+      ScanType,
+      ScanHeapType,
+      PostType,
+      PostHeapType,
+    } kind;
+    union {
+      Type* type;
+      HeapType* heapType;
+    };
+    static Task preVisit(Type* type) { return Task(type, PreType); }
+    static Task preVisit(HeapType* ht) { return Task(ht, PreHeapType); }
+    static Task scan(Type* type) { return Task(type, ScanType); }
+    static Task scan(HeapType* ht) { return Task(ht, ScanHeapType); }
+    static Task postVisit(Type* type) { return Task(type, PostType); }
+    static Task postVisit(HeapType* ht) { return Task(ht, PostHeapType); }
+
+  private:
+    Task(Type* type, Kind kind) : kind(kind), type(type) {}
+    Task(HeapType* ht, Kind kind) : kind(kind), heapType(ht) {}
+  };
+
+  void doWalk();
+
+  std::vector<Task> taskList;
+  void push(Type* type);
+  void push(HeapType* type);
+
+  Self& self() { return *static_cast<Self*>(this); }
+};
+
+// A type graph walker base class that still does no useful work, but at least
+// knows to scan each HeapType only once.
+template<typename Self> struct HeapTypeGraphWalker : TypeGraphWalkerBase<Self> {
+  // Override this.
+  void noteHeapType(HeapType ht) {}
+
+  void scanHeapType(HeapType* ht) {
+    if (scanned.insert(*ht).second) {
+      static_cast<Self*>(this)->noteHeapType(*ht);
+      TypeGraphWalkerBase<Self>::scanHeapType(ht);
+    }
+  }
+
+private:
+  std::unordered_set<HeapType> scanned;
+};
+
+// A type graph walker base class that still does no useful work, but at least
+// knows to scan each HeapType and Type only once.
+template<typename Self> struct TypeGraphWalker : TypeGraphWalkerBase<Self> {
+  // Override these.
+  void noteType(Type type) {}
+  void noteHeapType(HeapType ht) {}
+
+  void scanType(Type* type) {
+    if (scannedTypes.insert(*type).second) {
+      static_cast<Self*>(this)->noteType(*type);
+      TypeGraphWalkerBase<Self>::scanType(type);
+    }
+  }
+  void scanHeapType(HeapType* ht) {
+    if (scannedHeapTypes.insert(*ht).second) {
+      static_cast<Self*>(this)->noteHeapType(*ht);
+      TypeGraphWalkerBase<Self>::scanHeapType(ht);
+    }
+  }
+
+private:
+  std::unordered_set<HeapType> scannedHeapTypes;
+  std::unordered_set<Type> scannedTypes;
+};
+
+// A type graph walker that only traverses the direct HeapType children of the
+// root, looking through child Types. What to do with each child is left to
+// subclasses.
+template<typename Self> struct HeapTypeChildWalker : HeapTypeGraphWalker<Self> {
+  // Override this.
+  void noteChild(HeapType* child) {}
+
+  void scanType(Type* type) {
+    isTopLevel = false;
+    HeapTypeGraphWalker<Self>::scanType(type);
+  }
+  void scanHeapType(HeapType* ht) {
+    if (isTopLevel) {
+      HeapTypeGraphWalker<Self>::scanHeapType(ht);
+    } else {
+      static_cast<Self*>(this)->noteChild(ht);
+    }
+  }
+
+private:
+  bool isTopLevel = true;
+};
+
+struct HeapTypeChildCollector : HeapTypeChildWalker<HeapTypeChildCollector> {
+  std::vector<HeapType> children;
+  void noteChild(HeapType* child) { children.push_back(*child); }
+};
+
 } // anonymous namespace
 } // namespace wasm
 
@@ -878,6 +1004,12 @@ bool Type::isSubType(Type left, Type right) {
   return SubTyper().isSubType(left, right);
 }
 
+std::vector<HeapType> Type::getHeapTypeChildren() {
+  HeapTypeChildCollector collector;
+  collector.walkRoot(this);
+  return collector.children;
+}
+
 bool Type::hasLeastUpperBound(Type a, Type b) {
   return TypeBounder().hasLeastUpperBound(a, b);
 }
@@ -1007,6 +1139,12 @@ bool HeapType::isSubType(HeapType left, HeapType right) {
   return SubTyper().isSubType(left, right);
 }
 
+std::vector<HeapType> HeapType::getHeapTypeChildren() {
+  HeapTypeChildCollector collector;
+  collector.walkRoot(this);
+  return collector.children;
+}
+
 bool Signature::operator<(const Signature& other) const {
   return TypeComparator().lessThan(*this, other);
 }
@@ -1939,6 +2077,98 @@ bool FiniteShapeEquator::eq(const Rtt& a, const Rtt& b) {
   return a.depth == b.depth && eq(a.heapType, b.heapType);
 }
 
+template<typename Self> void TypeGraphWalkerBase<Self>::walkRoot(Type* type) {
+  assert(taskList.empty());
+  taskList.push_back(Task::scan(type));
+  doWalk();
+}
+
+template<typename Self> void TypeGraphWalkerBase<Self>::walkRoot(HeapType* ht) {
+  assert(taskList.empty());
+  taskList.push_back(Task::scan(ht));
+  doWalk();
+}
+
+template<typename Self> void TypeGraphWalkerBase<Self>::doWalk() {
+  while (!taskList.empty()) {
+    auto curr = taskList.back();
+    taskList.pop_back();
+    switch (curr.kind) {
+      case Task::PreType:
+        self().preVisitType(curr.type);
+        break;
+      case Task::PreHeapType:
+        self().preVisitHeapType(curr.heapType);
+        break;
+      case Task::ScanType:
+        taskList.push_back(Task::postVisit(curr.type));
+        self().scanType(curr.type);
+        taskList.push_back(Task::preVisit(curr.type));
+        break;
+      case Task::ScanHeapType:
+        taskList.push_back(Task::postVisit(curr.heapType));
+        self().scanHeapType(curr.heapType);
+        taskList.push_back(Task::preVisit(curr.heapType));
+        break;
+      case Task::PostType:
+        self().postVisitType(curr.type);
+        break;
+      case Task::PostHeapType:
+        self().postVisitHeapType(curr.heapType);
+        break;
+    }
+  }
+}
+
+template<typename Self> void TypeGraphWalkerBase<Self>::scanType(Type* type) {
+  if (type->isBasic()) {
+    return;
+  }
+  auto* info = getTypeInfo(*type);
+  switch (info->kind) {
+    case TypeInfo::TupleKind: {
+      auto& types = info->tuple.types;
+      for (auto it = types.rbegin(); it != types.rend(); ++it) {
+        taskList.push_back(Task::scan(&*it));
+      }
+      break;
+    }
+    case TypeInfo::RefKind: {
+      taskList.push_back(Task::scan(&info->ref.heapType));
+      break;
+    }
+    case TypeInfo::RttKind:
+      taskList.push_back(Task::scan(&info->rtt.heapType));
+      break;
+  }
+}
+
+template<typename Self>
+void TypeGraphWalkerBase<Self>::scanHeapType(HeapType* ht) {
+  if (ht->isBasic()) {
+    return;
+  }
+  auto* info = getHeapTypeInfo(*ht);
+  switch (info->kind) {
+    case HeapTypeInfo::BasicKind:
+      break;
+    case HeapTypeInfo::SignatureKind:
+      taskList.push_back(Task::scan(&info->signature.results));
+      taskList.push_back(Task::scan(&info->signature.params));
+      break;
+    case HeapTypeInfo::StructKind: {
+      auto& fields = info->struct_.fields;
+      for (auto field = fields.rbegin(); field != fields.rend(); ++field) {
+        taskList.push_back(Task::scan(&field->type));
+      }
+      break;
+    }
+    case HeapTypeInfo::ArrayKind:
+      taskList.push_back(Task::scan(&info->array.element.type));
+      break;
+  }
+}
+
 } // anonymous namespace
 
 struct TypeBuilder::Impl {
@@ -2031,217 +2261,6 @@ Type TypeBuilder::getTempRttType(Rtt rtt) {
 
 namespace {
 
-// Generic utility for traversing type graphs. The inserted roots must live as
-// long as the Walker because they are referenced by address. This base class
-// only has logic for traversing type graphs; figuring out when to stop
-// traversing the graph and doing useful work during the traversal is left to
-// subclasses.
-template<typename Self> struct TypeGraphWalkerBase {
-  void walkRoot(Type* type);
-  void walkRoot(HeapType* ht);
-
-  // Override these in subclasses to do useful work.
-  void preVisitType(Type* type) {}
-  void preVisitHeapType(HeapType* ht) {}
-  void postVisitType(Type* type) {}
-  void postVisitHeapType(HeapType* ht) {}
-
-  // This base walker does not know when to stop scanning, so at least one of
-  // these needs to be overridden with a method that calls the base scanning
-  // method only if some end condition isn't met.
-  void scanHeapType(HeapType* ht);
-  void scanType(Type* type);
-
-private:
-  struct Task {
-    enum Kind {
-      PreType,
-      PreHeapType,
-      ScanType,
-      ScanHeapType,
-      PostType,
-      PostHeapType,
-    } kind;
-    union {
-      Type* type;
-      HeapType* heapType;
-    };
-    static Task preVisit(Type* type) { return Task(type, PreType); }
-    static Task preVisit(HeapType* ht) { return Task(ht, PreHeapType); }
-    static Task scan(Type* type) { return Task(type, ScanType); }
-    static Task scan(HeapType* ht) { return Task(ht, ScanHeapType); }
-    static Task postVisit(Type* type) { return Task(type, PostType); }
-    static Task postVisit(HeapType* ht) { return Task(ht, PostHeapType); }
-
-  private:
-    Task(Type* type, Kind kind) : kind(kind), type(type) {}
-    Task(HeapType* ht, Kind kind) : kind(kind), heapType(ht) {}
-  };
-
-  void doWalk();
-
-  std::vector<Task> taskList;
-  void push(Type* type);
-  void push(HeapType* type);
-
-  Self& self() { return *static_cast<Self*>(this); }
-};
-
-template<typename Self> void TypeGraphWalkerBase<Self>::walkRoot(Type* type) {
-  assert(taskList.empty());
-  taskList.push_back(Task::scan(type));
-  doWalk();
-}
-
-template<typename Self> void TypeGraphWalkerBase<Self>::walkRoot(HeapType* ht) {
-  assert(taskList.empty());
-  taskList.push_back(Task::scan(ht));
-  doWalk();
-}
-
-template<typename Self> void TypeGraphWalkerBase<Self>::doWalk() {
-  while (!taskList.empty()) {
-    auto curr = taskList.back();
-    taskList.pop_back();
-    switch (curr.kind) {
-      case Task::PreType:
-        self().preVisitType(curr.type);
-        break;
-      case Task::PreHeapType:
-        self().preVisitHeapType(curr.heapType);
-        break;
-      case Task::ScanType:
-        taskList.push_back(Task::postVisit(curr.type));
-        self().scanType(curr.type);
-        taskList.push_back(Task::preVisit(curr.type));
-        break;
-      case Task::ScanHeapType:
-        taskList.push_back(Task::postVisit(curr.heapType));
-        self().scanHeapType(curr.heapType);
-        taskList.push_back(Task::preVisit(curr.heapType));
-        break;
-      case Task::PostType:
-        self().postVisitType(curr.type);
-        break;
-      case Task::PostHeapType:
-        self().postVisitHeapType(curr.heapType);
-        break;
-    }
-  }
-}
-
-template<typename Self>
-void TypeGraphWalkerBase<Self>::scanHeapType(HeapType* ht) {
-  if (ht->isBasic()) {
-    return;
-  }
-  auto* info = getHeapTypeInfo(*ht);
-  switch (info->kind) {
-    case HeapTypeInfo::BasicKind:
-      break;
-    case HeapTypeInfo::SignatureKind:
-      taskList.push_back(Task::scan(&info->signature.results));
-      taskList.push_back(Task::scan(&info->signature.params));
-      break;
-    case HeapTypeInfo::StructKind: {
-      auto& fields = info->struct_.fields;
-      for (auto field = fields.rbegin(); field != fields.rend(); ++field) {
-        taskList.push_back(Task::scan(&field->type));
-      }
-      break;
-    }
-    case HeapTypeInfo::ArrayKind:
-      taskList.push_back(Task::scan(&info->array.element.type));
-      break;
-  }
-}
-
-template<typename Self> void TypeGraphWalkerBase<Self>::scanType(Type* type) {
-  if (type->isBasic()) {
-    return;
-  }
-  auto* info = getTypeInfo(*type);
-  switch (info->kind) {
-    case TypeInfo::TupleKind: {
-      auto& types = info->tuple.types;
-      for (auto it = types.rbegin(); it != types.rend(); ++it) {
-        taskList.push_back(Task::scan(&*it));
-      }
-      break;
-    }
-    case TypeInfo::RefKind: {
-      taskList.push_back(Task::scan(&info->ref.heapType));
-      break;
-    }
-    case TypeInfo::RttKind:
-      taskList.push_back(Task::scan(&info->rtt.heapType));
-      break;
-  }
-}
-
-// A type graph walker base class that still does no useful work, but at least
-// knows to scan each HeapType only once.
-template<typename Self> struct HeapTypeGraphWalker : TypeGraphWalkerBase<Self> {
-  // Override this.
-  void noteHeapType(HeapType ht) {}
-
-  void scanHeapType(HeapType* ht) {
-    if (scanned.insert(*ht).second) {
-      static_cast<Self*>(this)->noteHeapType(*ht);
-      TypeGraphWalkerBase<Self>::scanHeapType(ht);
-    }
-  }
-
-private:
-  std::unordered_set<HeapType> scanned;
-};
-
-// A type graph walker base class that still does no useful work, but at least
-// knows to scan each HeapType and Type only once.
-template<typename Self> struct TypeGraphWalker : TypeGraphWalkerBase<Self> {
-  // Override these.
-  void noteType(Type type) {}
-  void noteHeapType(HeapType ht) {}
-
-  void scanType(Type* type) {
-    if (scannedTypes.insert(*type).second) {
-      static_cast<Self*>(this)->noteType(*type);
-      TypeGraphWalkerBase<Self>::scanType(type);
-    }
-  }
-  void scanHeapType(HeapType* ht) {
-    if (scannedHeapTypes.insert(*ht).second) {
-      static_cast<Self*>(this)->noteHeapType(*ht);
-      TypeGraphWalkerBase<Self>::scanHeapType(ht);
-    }
-  }
-
-private:
-  std::unordered_set<HeapType> scannedHeapTypes;
-  std::unordered_set<Type> scannedTypes;
-};
-
-// A type graph walker that notes parent-child relationships between HeapTypes.
-// What to do with each noted relationship is left to subclasses.
-template<typename Self> struct HeapTypePathWalker : HeapTypeGraphWalker<Self> {
-  // Override this.
-  void noteChild(HeapType parent, HeapType* child) {}
-
-  void preVisitHeapType(HeapType* ht) {
-    if (!path.empty()) {
-      static_cast<Self*>(this)->noteChild(path.back(), ht);
-    }
-    path.push_back(*ht);
-  }
-  void postVisitHeapType(HeapType* ht) {
-    assert(!path.empty() && path.back() == *ht);
-    path.pop_back();
-  }
-
-private:
-  std::vector<HeapType> path;
-};
-
 // A wrapper around a HeapType that provides equality and hashing based only on
 // its top-level shape, up to but not including its closest HeapType
 // descendants. This is the shape that determines the most fine-grained
@@ -2314,8 +2333,8 @@ private:
   void initializePartitions();
   void translatePartitionsToTypes();
 
-  // Return the non-basic HeapType children of `ht` (including BasicKind
-  // children).
+  // Return pointers to the non-basic HeapType children of `ht`, including
+  // BasicKind children.
   std::vector<HeapType*> getChildren(HeapType ht);
   const TypeSet& getPredsOf(HeapType type, size_t symbol);
   TypeSet getIntersection(const TypeSet& a, const TypeSet& b);
@@ -2542,24 +2561,16 @@ void ShapeCanonicalizer::translatePartitionsToTypes() {
 }
 
 std::vector<HeapType*> ShapeCanonicalizer::getChildren(HeapType ht) {
-  struct Walker : HeapTypePathWalker<Walker> {
+  struct Collector : HeapTypeChildWalker<Collector> {
     std::vector<HeapType*> children;
-    bool topLevel = true;
-    void noteChild(HeapType, HeapType* child) {
+    void noteChild(HeapType* child) {
       if (!child->isBasic()) {
         children.push_back(child);
       }
     }
-    void scanHeapType(HeapType* ht) {
-      if (topLevel) {
-        HeapTypePathWalker<Walker>::scanHeapType(ht);
-        topLevel = false;
-      }
-    }
-  };
-  Walker walker;
-  walker.walkRoot(&ht);
-  return walker.children;
+  } collector;
+  collector.walkRoot(&ht);
+  return collector.children;
 }
 
 const std::unordered_set<HeapType>&