2 files changed, 100 insertions, 2 deletions

diff --git a/src/ir/subtypes.h b/src/ir/subtypes.h
index 5a4ff87fc..969685292 100644
--- a/src/ir/subtypes.h
+++ b/src/ir/subtypes.h
@@ -24,19 +24,23 @@ namespace wasm {
 
 // Analyze subtyping relationships and provide useful interfaces to discover
 // them.
+//
+// This only scans user types, and not basic types like HeapType::eq.
 struct SubTypes {
-  SubTypes(Module& wasm) {
+  SubTypes(const std::vector<HeapType>& types) : types(types) {
     if (getTypeSystem() != TypeSystem::Nominal &&
         getTypeSystem() != TypeSystem::Isorecursive) {
       Fatal() << "SubTypes requires explicit supers";
     }
-    types = ModuleUtils::collectHeapTypes(wasm);
     for (auto type : types) {
       note(type);
     }
   }
 
+  SubTypes(Module& wasm) : SubTypes(ModuleUtils::collectHeapTypes(wasm)) {}
+
   const std::vector<HeapType>& getStrictSubTypes(HeapType type) {
+    assert(!type.isBasic());
     if (auto iter = typeSubTypes.find(type); iter != typeSubTypes.end()) {
       return iter->second;
     }
@@ -69,6 +73,49 @@ struct SubTypes {
     return ret;
   }
 
+  // Efficiently iterate on subtypes of a type, up to a particular depth (depth
+  // 0 means not to traverse subtypes, etc.). The callback function receives
+  // (type, depth).
+  template<typename F> void iterSubTypes(HeapType type, Index depth, F func) {
+    // Start by traversing the type itself.
+    func(type, 0);
+
+    if (depth == 0) {
+      // Nothing else to scan.
+      return;
+    }
+
+    // getStrictSubTypes() returns vectors of subtypes, so for efficiency store
+    // pointers to those in our work queue to avoid allocations. See the note
+    // below on typeSubTypes for why this is safe.
+    struct Item {
+      const std::vector<HeapType>* vec;
+      Index depth;
+    };
+
+    // Real-world type hierarchies tend to have a limited depth, so try to avoid
+    // allocations in our work queue with a SmallVector.
+    SmallVector<Item, 10> work;
+
+    // Start with the subtypes of the base type. Those have depth 1.
+    work.push_back({&getStrictSubTypes(type), 1});
+
+    while (!work.empty()) {
+      auto& item = work.back();
+      work.pop_back();
+      auto currDepth = item.depth;
+      auto& currVec = *item.vec;
+      assert(currDepth <= depth);
+      for (auto type : currVec) {
+        func(type, currDepth);
+        auto* subVec = &getStrictSubTypes(type);
+        if (currDepth + 1 <= depth && !subVec->empty()) {
+          work.push_back({subVec, currDepth + 1});
+        }
+      }
+    }
+  }
+
   // All the types in the program. This is computed here anyhow, and can be
   // useful for callers to iterate on, so it is public.
   std::vector<HeapType> types;
diff --git a/test/gtest/type-builder.cpp b/test/gtest/type-builder.cpp
index d642893c3..cb2150974 100644
--- a/test/gtest/type-builder.cpp
+++ b/test/gtest/type-builder.cpp
@@ -661,3 +661,54 @@ TEST_F(NominalTest, TestDepth) {
   EXPECT_EQ(HeapType(HeapType::nofunc).getDepth(), size_t(-1));
   EXPECT_EQ(HeapType(HeapType::noext).getDepth(), size_t(-1));
 }
+
+// Test .iterSubTypes() helper.
+TEST_F(NominalTest, TestIterSubTypes) {
+  /*
+        A
+       / \
+      B   C
+           \
+            D
+  */
+  HeapType A, B, C, D;
+  {
+    TypeBuilder builder(4);
+    builder[0] = Struct();
+    builder[1] = Struct();
+    builder[2] = Struct();
+    builder[3] = Struct();
+    builder.setSubType(1, builder.getTempHeapType(0));
+    builder.setSubType(2, builder.getTempHeapType(0));
+    builder.setSubType(3, builder.getTempHeapType(2));
+    auto result = builder.build();
+    ASSERT_TRUE(result);
+    auto built = *result;
+    A = built[0];
+    B = built[1];
+    C = built[2];
+    D = built[3];
+  }
+
+  SubTypes subTypes({A, B, C, D});
+
+  // Helper for comparing sets of types + their corresponding depth.
+  using TypeDepths = std::unordered_set<std::pair<HeapType, Index>>;
+
+  auto getSubTypes = [&](HeapType type, Index depth) {
+    TypeDepths ret;
+    subTypes.iterSubTypes(type, depth, [&](HeapType subType, Index depth) {
+      ret.insert({subType, depth});
+    });
+    return ret;
+  };
+
+  EXPECT_EQ(getSubTypes(A, 0), TypeDepths({{A, 0}}));
+  EXPECT_EQ(getSubTypes(A, 1), TypeDepths({{A, 0}, {B, 1}, {C, 1}}));
+  EXPECT_EQ(getSubTypes(A, 2), TypeDepths({{A, 0}, {B, 1}, {C, 1}, {D, 2}}));
+  EXPECT_EQ(getSubTypes(A, 3), TypeDepths({{A, 0}, {B, 1}, {C, 1}, {D, 2}}));
+
+  EXPECT_EQ(getSubTypes(C, 0), TypeDepths({{C, 0}}));
+  EXPECT_EQ(getSubTypes(C, 1), TypeDepths({{C, 0}, {D, 1}}));
+  EXPECT_EQ(getSubTypes(C, 2), TypeDepths({{C, 0}, {D, 1}}));
+}