Add a HeapType method for getting the rec group index (#4480)

Storing the rec group index on the HeapTypeInfo avoids having to do a linear
scan through the rec group to find the index for a particular type. This will
be important for isorecursive canonicalization, which uses rec group indices.

2 files changed, 10 insertions, 2 deletions

diff --git a/src/wasm-type.h b/src/wasm-type.h
index 7b63adf9d..99194b81d 100644
--- a/src/wasm-type.h
+++ b/src/wasm-type.h
@@ -356,6 +356,7 @@ public:
 
   // Get the recursion group for this non-basic type.
   RecGroup getRecGroup() const;
+  size_t getRecGroupIndex() const;
 
   constexpr TypeID getID() const { return id; }
   constexpr BasicHeapType getBasic() const {
@@ -392,8 +393,8 @@ class RecGroup {
 
 public:
   explicit RecGroup(uintptr_t id) : id(id) {}
-  bool operator==(const RecGroup& other) { return id == other.id; }
-  bool operator!=(const RecGroup& other) { return id != other.id; }
+  bool operator==(const RecGroup& other) const { return id == other.id; }
+  bool operator!=(const RecGroup& other) const { return id != other.id; }
   size_t size() const;
 
   struct Iterator : ParentIndexIterator<const RecGroup*, Iterator> {
diff --git a/src/wasm/wasm-type.cpp b/src/wasm/wasm-type.cpp
index 53de89aa5..f63505b20 100644
--- a/src/wasm/wasm-type.cpp
+++ b/src/wasm/wasm-type.cpp
@@ -113,6 +113,7 @@ struct HeapTypeInfo {
   // In isorecursive mode, the recursion group of this type or null if the
   // recursion group is trivial (i.e. contains only this type).
   RecGroupInfo* recGroup = nullptr;
+  size_t recGroupIndex = 0;
   enum Kind {
     BasicKind,
     SignatureKind,
@@ -1263,6 +1264,11 @@ RecGroup HeapType::getRecGroup() const {
   }
 }
 
+size_t HeapType::getRecGroupIndex() const {
+  assert(!isBasic());
+  return getHeapTypeInfo(*this)->recGroupIndex;
+}
+
 HeapType RecGroup::Iterator::operator*() const {
   if (parent->id & 1) {
     // This is a trivial recursion group. Mask off the low bit to recover the
@@ -3242,6 +3248,7 @@ std::optional<TypeBuilder::Error> canonicalizeIsorecursive(
   // Fill out the recursion groups.
   for (auto& info : state.newInfos) {
     if (info->recGroup != nullptr) {
+      info->recGroupIndex = info->recGroup->size();
       info->recGroup->push_back(asHeapType(info));
     }
   }