Refactor ModuleUtils::collectHeapTypes (#4455)

Update the API to make both the type indices and optimized sorting optional.
It will become more important to avoid unnecessary sorting once isorecursive
types have been implemented because they will make the sorting more complicated.

1 files changed, 50 insertions, 15 deletions

diff --git a/src/ir/module-utils.cpp b/src/ir/module-utils.cpp
index c738c6143..9b2293984 100644
--- a/src/ir/module-utils.cpp
+++ b/src/ir/module-utils.cpp
@@ -19,22 +19,23 @@
 
 namespace wasm::ModuleUtils {
 
-void collectHeapTypes(Module& wasm,
-                      std::vector<HeapType>& types,
-                      std::unordered_map<HeapType, Index>& typeIndices) {
-  struct Counts : public InsertOrderedMap<HeapType, size_t> {
-    void note(HeapType type) {
-      if (!type.isBasic()) {
-        (*this)[type]++;
-      }
+namespace {
+
+// Helper for collecting HeapTypes and their frequencies.
+struct Counts : public InsertOrderedMap<HeapType, size_t> {
+  void note(HeapType type) {
+    if (!type.isBasic()) {
+      (*this)[type]++;
     }
-    void note(Type type) {
-      for (HeapType ht : type.getHeapTypeChildren()) {
-        note(ht);
-      }
+  }
+  void note(Type type) {
+    for (HeapType ht : type.getHeapTypeChildren()) {
+      note(ht);
     }
-  };
+  }
+};
 
+Counts getHeapTypeCounts(Module& wasm) {
   struct CodeScanner
     : PostWalker<CodeScanner, UnifiedExpressionVisitor<CodeScanner>> {
     Counts& counts;
@@ -160,15 +161,49 @@ void collectHeapTypes(Module& wasm,
     }
   }
 
+  return counts;
+}
+
+} // anonymous namespace
+
+std::vector<HeapType> collectHeapTypes(Module& wasm) {
+  Counts counts = getHeapTypeCounts(wasm);
+  std::vector<HeapType> types;
+  types.reserve(counts.size());
+  for (auto& [type, _] : counts) {
+    types.push_back(type);
+  }
+  return types;
+}
+
+IndexedHeapTypes getIndexedHeapTypes(Module& wasm) {
+  Counts counts = getHeapTypeCounts(wasm);
+  IndexedHeapTypes indexedTypes;
+  Index i = 0;
+  for (auto& [type, _] : counts) {
+    indexedTypes.types.push_back(type);
+    indexedTypes.indices[type] = i++;
+  }
+  return indexedTypes;
+}
+
+IndexedHeapTypes getOptimizedIndexedHeapTypes(Module& wasm) {
+  Counts counts = getHeapTypeCounts(wasm);
+
   // Sort by frequency and then original insertion order.
   std::vector<std::pair<HeapType, size_t>> sorted(counts.begin(), counts.end());
   std::stable_sort(sorted.begin(), sorted.end(), [&](auto a, auto b) {
     return a.second > b.second;
   });
+
+  // Collect the results.
+  IndexedHeapTypes indexedTypes;
   for (Index i = 0; i < sorted.size(); ++i) {
-    typeIndices[sorted[i].first] = i;
-    types.push_back(sorted[i].first);
+    HeapType type = sorted[i].first;
+    indexedTypes.types.push_back(type);
+    indexedTypes.indices[type] = i;
   }
+  return indexedTypes;
 }
 
 } // namespace wasm::ModuleUtils