Add SubTypes::getAllSubTypes variant which includes the type itself (#4649)

This also includes the type itself in the returned vector. This will be
useful in a future PR.

5 files changed, 46 insertions, 17 deletions

diff --git a/src/ir/struct-utils.h b/src/ir/struct-utils.h
index bc2eb5485..e7b587958 100644
--- a/src/ir/struct-utils.h
+++ b/src/ir/struct-utils.h
@@ -260,7 +260,7 @@ private:
       if (toSubTypes) {
         // Propagate shared fields to the subtypes.
         auto numFields = type.getStruct().fields.size();
-        for (auto subType : subTypes.getSubTypes(type)) {
+        for (auto subType : subTypes.getStrictSubTypes(type)) {
           auto& subInfos = combinedInfos[subType];
           for (Index i = 0; i < numFields; i++) {
             if (subInfos[i].combine(infos[i])) {
diff --git a/src/ir/subtypes.h b/src/ir/subtypes.h
index 2315934d7..9064d5ed0 100644
--- a/src/ir/subtypes.h
+++ b/src/ir/subtypes.h
@@ -26,24 +26,27 @@ namespace wasm {
 // them.
 struct SubTypes {
   SubTypes(Module& wasm) {
+    if (getTypeSystem() != TypeSystem::Nominal) {
+      Fatal() << "SubTypes requires nominal typing to find supers";
+    }
     types = ModuleUtils::collectHeapTypes(wasm);
     for (auto type : types) {
       note(type);
     }
   }
 
-  const std::vector<HeapType>& getSubTypes(HeapType type) {
+  const std::vector<HeapType>& getStrictSubTypes(HeapType type) {
     return typeSubTypes[type];
   }
 
   // Get all subtypes of a type, and their subtypes and so forth, recursively.
-  std::vector<HeapType> getAllSubTypes(HeapType type) {
+  std::vector<HeapType> getAllStrictSubTypes(HeapType type) {
     std::vector<HeapType> ret, work;
     work.push_back(type);
     while (!work.empty()) {
       auto curr = work.back();
       work.pop_back();
-      for (auto sub : getSubTypes(curr)) {
+      for (auto sub : getStrictSubTypes(curr)) {
         ret.push_back(sub);
         work.push_back(sub);
       }
@@ -51,20 +54,15 @@ struct SubTypes {
     return ret;
   }
 
-  // Get all supertypes of a type. The order in the output vector is with the
-  // immediate supertype first, then its supertype, and so forth.
-  std::vector<HeapType> getAllSuperTypes(HeapType type) {
-    std::vector<HeapType> ret;
-    while (1) {
-      auto super = type.getSuperType();
-      if (!super) {
-        return ret;
-      }
-      ret.push_back(*super);
-      type = *super;
-    }
+  // Like getAllStrictSubTypes, but also includes the type itself.
+  std::vector<HeapType> getAllSubTypes(HeapType type) {
+    auto ret = getAllStrictSubTypes(type);
+    ret.push_back(type);
+    return ret;
   }
 
+  // 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;
 
 private:
diff --git a/src/passes/TypeRefining.cpp b/src/passes/TypeRefining.cpp
index a2a37f750..0f5169f8f 100644
--- a/src/passes/TypeRefining.cpp
+++ b/src/passes/TypeRefining.cpp
@@ -191,7 +191,7 @@ struct TypeRefining : public Pass {
         }
       }
 
-      for (auto subType : subTypes.getSubTypes(type)) {
+      for (auto subType : subTypes.getStrictSubTypes(type)) {
         work.push(subType);
       }
     }
diff --git a/test/example/type-builder-nominal.cpp b/test/example/type-builder-nominal.cpp
index 34571c48e..8dbbfcf30 100644
--- a/test/example/type-builder-nominal.cpp
+++ b/test/example/type-builder-nominal.cpp
@@ -1,6 +1,7 @@
 #include <cassert>
 #include <iostream>
 
+#include "wasm-builder.h"
 #include "wasm-type-printing.h"
 #include "wasm-type.h"
 
diff --git a/test/gtest/type-builder.cpp b/test/gtest/type-builder.cpp
index 42b388532..5f3ed7384 100644
--- a/test/gtest/type-builder.cpp
+++ b/test/gtest/type-builder.cpp
@@ -1,3 +1,4 @@
+#include "ir/subtypes.h"
 #include "type-test.h"
 #include "wasm-type-printing.h"
 #include "wasm-type.h"
@@ -498,3 +499,32 @@ TEST_F(EquirecursiveTest, CanonicalizeBasicTypes) {
 TEST_F(IsorecursiveTest, CanonicalizeBasicTypes) {
   testCanonicalizeBasicTypes();
 }
+
+// Test SubTypes utility code.
+TEST_F(NominalTest, testSubTypes) {
+  // Build type types, the second of which is a subtype.
+  TypeBuilder builder(2);
+  builder[0] = Struct({Field(Type::anyref, Immutable)});
+  builder[1] = Struct({Field(Type::funcref, Immutable)});
+  builder[1].subTypeOf(builder[0]);
+  auto built = *builder.build();
+
+  // Build a tiny wasm module that uses the types, so that we can test the
+  // SubTypes utility code.
+  Module wasm;
+  Builder wasmBuilder(wasm);
+  wasm.addFunction(wasmBuilder.makeFunction(
+    "func",
+    Signature(Type::none, Type::none),
+    {Type(built[0], Nullable), Type(built[1], Nullable)},
+    wasmBuilder.makeNop()));
+  SubTypes subTypes(wasm);
+  auto subTypes0 = subTypes.getStrictSubTypes(built[0]);
+  EXPECT_TRUE(subTypes0.size() == 1 && subTypes0[0] == built[1]);
+  auto subTypes0Inclusive = subTypes.getAllSubTypes(built[0]);
+  EXPECT_TRUE(subTypes0Inclusive.size() == 2 &&
+              subTypes0Inclusive[0] == built[1] &&
+              subTypes0Inclusive[1] == built[0]);
+  auto subTypes1 = subTypes.getStrictSubTypes(built[1]);
+  EXPECT_EQ(subTypes1.size(), 0u);
+}