[llvm] [ADT] Fix a bug in EquivalenceClasses::erase (PR #161121)

[Kazu Hirata via llvm-commits]

https://github.com/kazutakahirata updated https://github.com/llvm/llvm-project/pull/161121

>From eb806f1955f4dd35d663a8f31fbf19427193df5c Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Sun, 28 Sep 2025 11:35:15 -0700
Subject: [PATCH 1/2] [ADT] Fix a bug in EquivalenceClasses::erase

This patch fixes a bug in EquivalenceClasses::erase, where we lose a
leader bit in a certain scenario.

Here is some background.  In EquivalenceClasses, each equivalence
class is maintained as a singly linked list over its members.  When we
join two classes, we concatenate the two singly linked lists.  To
support path compression, each member points to the leader (through
lazy updates).  This is implemented with the two members:

  class ECValue {
    mutable const ECValue *Leader, *Next;
    :
  };

Each member stores its leader in Leader and its sibling in Next.  Now,
the leader uses the Leader field to to point the last element of the
singly linked list to accommodate the list concatenation.  We use the
LSB of the Next field to indicate whether a given member is a leader
or not.

Now, imagine we have an equivalence class:

  Elem 1 -> Elem 2 -> nullptr
  Leader

and wish to remove Elem 2.  We would like to end up with:

  Elem 1 -> nullptr
  Leader

but we mistakenly drop the leader bit when we update the Next field of
Elem 1 with:

  Pre->Next = nullptr;

This makes Elem 1 the end of the singly linked list, as intended, but
mistakenly clears its leader bit stored in the LSB of Next, so we end
up with an equivalence class with no leader.

This patch fixes the problem by preserving the leader bit:

  Pre->Next = reinterpret_cast<const ECValue *>(
      static_cast<intptr_t>(Pre->isLeader()));

The unit test closely follows the scenario above.
---
 llvm/include/llvm/ADT/EquivalenceClasses.h    |  3 ++-
 llvm/unittests/ADT/EquivalenceClassesTest.cpp | 23 +++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/ADT/EquivalenceClasses.h b/llvm/include/llvm/ADT/EquivalenceClasses.h
index 1a2331c1a0322..0010975a64106 100644
--- a/llvm/include/llvm/ADT/EquivalenceClasses.h
+++ b/llvm/include/llvm/ADT/EquivalenceClasses.h
@@ -258,7 +258,8 @@ template <class ElemTy> class EquivalenceClasses {
         // If the current element is the last element(not leader), set the
         // successor of the current element's predecessor to null, and set
         // the 'Leader' field of the class leader to the predecessor element.
-        Pre->Next = nullptr;
+        Pre->Next = reinterpret_cast<const ECValue *>(
+            static_cast<intptr_t>(Pre->isLeader()));
         Leader->Leader = Pre;
       } else {
         // If the current element is in the middle of class, then simply
diff --git a/llvm/unittests/ADT/EquivalenceClassesTest.cpp b/llvm/unittests/ADT/EquivalenceClassesTest.cpp
index 3d5c48eb8e1b6..725247e8384db 100644
--- a/llvm/unittests/ADT/EquivalenceClassesTest.cpp
+++ b/llvm/unittests/ADT/EquivalenceClassesTest.cpp
@@ -108,6 +108,29 @@ TEST(EquivalenceClassesTest, SimpleErase4) {
   EXPECT_FALSE(EqClasses.erase(1));
 }
 
+TEST(EquivalenceClassesTest, EraseKeepsLeaderBit) {
+  EquivalenceClasses<int> EC;
+
+  // Create a set {1, 2} where 1 is the leader.
+  EC.unionSets(1, 2);
+
+  // Verify initial state
+  EXPECT_EQ(EC.getLeaderValue(2), 1);
+
+  // Erase 2, the non-leader member.
+  EXPECT_TRUE(EC.erase(2));
+
+  // Verify that we have exactly one equivalence class.
+  ASSERT_NE(EC.begin(), EC.end());
+  ASSERT_EQ(std::next(EC.begin()), EC.end());
+
+  // Verify that 1 is still a leader after erasing 2.
+  const auto *Elem = *EC.begin();
+  ASSERT_NE(Elem, nullptr);
+  EXPECT_EQ(Elem->getData(), 1);
+  EXPECT_TRUE(Elem->isLeader()) << "The leader bit was lost!";
+}
+
 TEST(EquivalenceClassesTest, TwoSets) {
   EquivalenceClasses<int> EqClasses;
   // Form sets of odd and even numbers, check that we split them into these

>From 8becaf8e436ecf0ea59300c93e74db4c814c9b71 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Mon, 29 Sep 2025 07:50:29 -0700
Subject: [PATCH 2/2] Address a comment.

---
 llvm/include/llvm/ADT/EquivalenceClasses.h | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/ADT/EquivalenceClasses.h b/llvm/include/llvm/ADT/EquivalenceClasses.h
index 0010975a64106..0dd2cdd837d4a 100644
--- a/llvm/include/llvm/ADT/EquivalenceClasses.h
+++ b/llvm/include/llvm/ADT/EquivalenceClasses.h
@@ -256,8 +256,9 @@ template <class ElemTy> class EquivalenceClasses {
       }
       if (!Next) {
         // If the current element is the last element(not leader), set the
-        // successor of the current element's predecessor to null, and set
-        // the 'Leader' field of the class leader to the predecessor element.
+        // successor of the current element's predecessor to null while
+        // preserving the leader bit, and set the 'Leader' field of the class
+        // leader to the predecessor element.
         Pre->Next = reinterpret_cast<const ECValue *>(
             static_cast<intptr_t>(Pre->isLeader()));
         Leader->Leader = Pre;