2 files changed, 70 insertions, 25 deletions

diff --git a/src/ir/type-updating.cpp b/src/ir/type-updating.cpp
index 9f9692a7b..00f313a2a 100644
--- a/src/ir/type-updating.cpp
+++ b/src/ir/type-updating.cpp
@@ -20,6 +20,7 @@
 #include "ir/module-utils.h"
 #include "ir/utils.h"
 #include "support/topological_sort.h"
+#include "wasm-type-ordering.h"
 #include "wasm-type.h"
 #include "wasm.h"
 
@@ -33,32 +34,9 @@ void GlobalTypeRewriter::update() {
   // be reflected on or used for linking. Figure out where each private type
   // will be located in the builder. Sort the private types so that supertypes
   // come before their subtypes.
-  struct SortedPrivateTypes : TopologicalSort<HeapType, SortedPrivateTypes> {
-    SortedPrivateTypes(Module& wasm) {
-      auto privateTypes = ModuleUtils::getPrivateHeapTypes(wasm);
-      std::unordered_set<HeapType> supertypes;
-      for (auto type : privateTypes) {
-        if (auto super = type.getSuperType()) {
-          supertypes.insert(*super);
-        }
-      }
-      // Types that are not supertypes of others are the roots.
-      for (auto type : privateTypes) {
-        if (!supertypes.count(type)) {
-          push(type);
-        }
-      }
-    }
-
-    void pushPredecessors(HeapType type) {
-      if (auto super = type.getSuperType()) {
-        push(*super);
-      }
-    }
-  };
-
   Index i = 0;
-  for (auto type : SortedPrivateTypes(wasm)) {
+  auto privateTypes = ModuleUtils::getPrivateHeapTypes(wasm);
+  for (auto type : HeapTypeOrdering::SupertypesFirst(privateTypes)) {
     typeIndices[type] = i++;
   }
 
diff --git a/src/wasm-type-ordering.h b/src/wasm-type-ordering.h
new file mode 100644
index 000000000..0b0042ca3
--- /dev/null
+++ b/src/wasm-type-ordering.h
@@ -0,0 +1,67 @@
+/*
+ * Copyright 2023 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef wasm_wasm_type_ordering_h
+#define wasm_wasm_type_ordering_h
+
+#include <unordered_set>
+
+#include "support/insert_ordered.h"
+#include "support/topological_sort.h"
+#include "wasm-type.h"
+
+namespace wasm::HeapTypeOrdering {
+
+// Given a collection of types, iterate through it such that each type in the
+// collection is visited only after its immediate supertype in the collection is
+// visited.
+template<typename T>
+struct SupertypesFirst : TopologicalSort<HeapType, SupertypesFirst<T>> {
+  // For each type in the input collection, whether it is a supertype. Used to
+  // track membership in the input collection.
+  InsertOrderedMap<HeapType, bool> typeSet;
+
+  SupertypesFirst(const T& types) {
+    for (auto type : types) {
+      typeSet[type] = false;
+    }
+    // Find the supertypes that are in the collection.
+    for (auto [type, _] : typeSet) {
+      if (auto super = type.getSuperType()) {
+        if (auto it = typeSet.find(*super); it != typeSet.end()) {
+          it->second = true;
+        }
+      }
+    }
+    // Types that are not supertypes of others are the roots.
+    for (auto [type, isSuper] : typeSet) {
+      if (!isSuper) {
+        this->push(type);
+      }
+    }
+  }
+
+  void pushPredecessors(HeapType type) {
+    // Do not visit types that weren't in the input collection.
+    if (auto super = type.getSuperType(); super && typeSet.count(*super)) {
+      this->push(*super);
+    }
+  }
+};
+
+} // namespace wasm::HeapTypeOrdering
+
+#endif // wasm_wasm_type_ordering_h