1 files changed, 100 insertions, 1 deletions

diff --git a/src/ir/module-utils.cpp b/src/ir/module-utils.cpp
index 22f07a2e8..f2c90dcfd 100644
--- a/src/ir/module-utils.cpp
+++ b/src/ir/module-utils.cpp
@@ -198,10 +198,86 @@ void setIndices(IndexedHeapTypes& indexedTypes) {
   }
 }
 
+InsertOrderedSet<HeapType> getPublicTypeSet(Module& wasm) {
+  InsertOrderedSet<HeapType> publicTypes;
+
+  auto notePublic = [&](HeapType type) {
+    if (type.isBasic()) {
+      return;
+    }
+    // All the rec group members are public as well.
+    for (auto member : type.getRecGroup()) {
+      if (!publicTypes.insert(member)) {
+        // We've already inserted this rec group.
+        break;
+      }
+    }
+  };
+
+  // TODO: Consider Tags as well, but they should store HeapTypes instead of
+  // Signatures first.
+  ModuleUtils::iterImportedTables(wasm, [&](Table* table) {
+    assert(table->type.isRef());
+    notePublic(table->type.getHeapType());
+  });
+  ModuleUtils::iterImportedGlobals(wasm, [&](Global* global) {
+    if (global->type.isRef()) {
+      notePublic(global->type.getHeapType());
+    }
+  });
+  ModuleUtils::iterImportedFunctions(
+    wasm, [&](Function* func) { notePublic(func->type); });
+  for (auto& ex : wasm.exports) {
+    switch (ex->kind) {
+      case ExternalKind::Function: {
+        auto* func = wasm.getFunction(ex->value);
+        notePublic(func->type);
+        continue;
+      }
+      case ExternalKind::Table: {
+        auto* table = wasm.getTable(ex->value);
+        assert(table->type.isRef());
+        notePublic(table->type.getHeapType());
+        continue;
+      }
+      case ExternalKind::Memory:
+        // Never a reference type.
+        continue;
+      case ExternalKind::Global: {
+        auto* global = wasm.getGlobal(ex->value);
+        if (global->type.isRef()) {
+          notePublic(global->type.getHeapType());
+        }
+        continue;
+      }
+      case ExternalKind::Tag:
+        // TODO
+        continue;
+      case ExternalKind::Invalid:
+        break;
+    }
+    WASM_UNREACHABLE("unexpected export kind");
+  }
+
+  // Find all the other public types reachable from directly publicized types.
+  std::vector<HeapType> workList(publicTypes.begin(), publicTypes.end());
+  while (workList.size()) {
+    auto curr = workList.back();
+    workList.pop_back();
+    for (auto t : curr.getReferencedHeapTypes()) {
+      if (!t.isBasic() && publicTypes.insert(t)) {
+        workList.push_back(t);
+      }
+    }
+  }
+
+  return publicTypes;
+}
+
 } // anonymous namespace
 
 std::vector<HeapType> collectHeapTypes(Module& wasm) {
-  Counts counts = getHeapTypeCounts(wasm);
+  auto counts = getHeapTypeCounts(wasm);
   std::vector<HeapType> types;
   types.reserve(counts.size());
   for (auto& [type, _] : counts) {
@@ -210,6 +286,29 @@ std::vector<HeapType> collectHeapTypes(Module& wasm) {
   return types;
 }
 
+std::vector<HeapType> getPublicHeapTypes(Module& wasm) {
+  auto publicTypes = getPublicTypeSet(wasm);
+  std::vector<HeapType> types;
+  types.reserve(publicTypes.size());
+  for (auto type : publicTypes) {
+    types.push_back(type);
+  }
+  return types;
+}
+
+std::vector<HeapType> getPrivateHeapTypes(Module& wasm) {
+  auto allTypes = getHeapTypeCounts(wasm);
+  auto publicTypes = getPublicTypeSet(wasm);
+  std::vector<HeapType> types;
+  types.reserve(allTypes.size() - publicTypes.size());
+  for (auto [type, _] : allTypes) {
+    if (!publicTypes.count(type)) {
+      types.push_back(type);
+    }
+  }
+  return types;
+}
+
 IndexedHeapTypes getOptimizedIndexedHeapTypes(Module& wasm) {
   TypeSystem system = getTypeSystem();
   Counts counts = getHeapTypeCounts(wasm);