2 files changed, 82 insertions, 1 deletions

diff --git a/src/ir/subtypes.h b/src/ir/subtypes.h
index 969685292..0d8fccdbc 100644
--- a/src/ir/subtypes.h
+++ b/src/ir/subtypes.h
@@ -18,6 +18,7 @@
 #define wasm_ir_subtypes_h
 
 #include "ir/module-utils.h"
+#include "support/topological_sort.h"
 #include "wasm.h"
 
 namespace wasm {
@@ -39,7 +40,7 @@ struct SubTypes {
 
   SubTypes(Module& wasm) : SubTypes(ModuleUtils::collectHeapTypes(wasm)) {}
 
-  const std::vector<HeapType>& getStrictSubTypes(HeapType type) {
+  const std::vector<HeapType>& getStrictSubTypes(HeapType type) const {
     assert(!type.isBasic());
     if (auto iter = typeSubTypes.find(type); iter != typeSubTypes.end()) {
       return iter->second;
@@ -73,6 +74,57 @@ struct SubTypes {
     return ret;
   }
 
+  // Computes the depth of children for each type. This is 0 if the type has no
+  // subtypes, 1 if it has subtypes but none of those have subtypes themselves,
+  // and so forth.
+  //
+  // This depth ignores bottom types.
+  std::unordered_map<HeapType, Index> getMaxDepths() {
+    struct DepthSort : TopologicalSort<HeapType, DepthSort> {
+      const SubTypes& parent;
+
+      DepthSort(const SubTypes& parent) : parent(parent) {
+        for (auto type : parent.types) {
+          // The roots are types with no supertype.
+          if (!type.getSuperType()) {
+            push(type);
+          }
+        }
+      }
+
+      void pushPredecessors(HeapType type) {
+        // Things we need to process before each type are its subtypes. Once we
+        // know their depth, we can easily compute our own.
+        for (auto pred : parent.getStrictSubTypes(type)) {
+          push(pred);
+        }
+      }
+    };
+
+    std::unordered_map<HeapType, Index> depths;
+
+    for (auto type : DepthSort(*this)) {
+      // Begin with depth 0, then take into account the subtype depths.
+      Index depth = 0;
+      for (auto subType : getStrictSubTypes(type)) {
+        depth = std::max(depth, depths[subType] + 1);
+      }
+      depths[type] = depth;
+    }
+
+    // Add the max depths of basic types.
+    // TODO: update when we get structtype and arraytype
+    for (auto type : types) {
+      HeapType basic = type.isData() ? HeapType::data : HeapType::func;
+      depths[basic] = std::max(depths[basic], depths[type] + 1);
+    }
+
+    depths[HeapType::eq] = std::max(Index(1), depths[HeapType::data] + 1);
+    depths[HeapType::any] = depths[HeapType::eq] + 1;
+
+    return depths;
+  }
+
   // 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).
diff --git a/test/gtest/type-builder.cpp b/test/gtest/type-builder.cpp
index cb2150974..d92057502 100644
--- a/test/gtest/type-builder.cpp
+++ b/test/gtest/type-builder.cpp
@@ -620,6 +620,35 @@ TEST_F(IsorecursiveTest, TestExistingSuperType) {
   EXPECT_EQ(B1.getHeapType(), B2.getHeapType());
 }
 
+// Test .getMaxDepths() helper.
+TEST_F(NominalTest, TestMaxDepths) {
+  /*
+      A
+      |
+      B
+  */
+  HeapType A, B;
+  {
+    TypeBuilder builder(2);
+    builder[0] = Struct();
+    builder[1] = Struct();
+    builder.setSubType(1, builder.getTempHeapType(0));
+    auto result = builder.build();
+    ASSERT_TRUE(result);
+    auto built = *result;
+    A = built[0];
+    B = built[1];
+  }
+
+  SubTypes subTypes({A, B});
+  auto maxDepths = subTypes.getMaxDepths();
+
+  EXPECT_EQ(maxDepths[B], Index(0));
+  EXPECT_EQ(maxDepths[A], Index(1));
+  EXPECT_EQ(maxDepths[HeapType::data], Index(2));
+  EXPECT_EQ(maxDepths[HeapType::eq], Index(3));
+}
+
 // Test .depth() helper.
 TEST_F(NominalTest, TestDepth) {
   HeapType A, B, C;