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.

3 files changed, 37 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));
     }
   }
diff --git a/test/gtest/type-builder.cpp b/test/gtest/type-builder.cpp
index 78c8037a4..9cbd04298 100644
--- a/test/gtest/type-builder.cpp
+++ b/test/gtest/type-builder.cpp
@@ -242,3 +242,30 @@ TEST_F(IsorecursiveTest, ForwardReferencedChild) {
   EXPECT_EQ(error->reason, TypeBuilder::ErrorReason::ForwardChildReference);
   EXPECT_EQ(error->index, 1u);
 }
+
+TEST_F(IsorecursiveTest, RecGroupIndices) {
+  TypeBuilder builder(5);
+
+  builder.createRecGroup(0, 2);
+  builder[0] = Struct{};
+  builder[1] = Struct{};
+
+  builder.createRecGroup(2, 3);
+  builder[2] = Struct{};
+  builder[3] = Struct{};
+  builder[4] = Struct{};
+
+  auto result = builder.build();
+  ASSERT_TRUE(result);
+  auto built = *result;
+
+  EXPECT_EQ(built[0].getRecGroup(), built[1].getRecGroup());
+  EXPECT_EQ(built[0].getRecGroupIndex(), 0u);
+  EXPECT_EQ(built[1].getRecGroupIndex(), 1u);
+
+  EXPECT_EQ(built[2].getRecGroup(), built[3].getRecGroup());
+  EXPECT_EQ(built[3].getRecGroup(), built[4].getRecGroup());
+  EXPECT_EQ(built[2].getRecGroupIndex(), 0u);
+  EXPECT_EQ(built[3].getRecGroupIndex(), 1u);
+  EXPECT_EQ(built[4].getRecGroupIndex(), 2u);
+}