[llvm] 3e54505 - [ADT] Fix a bug in EquivalenceClasses::erase (#161121)

[via llvm-commits]

Author: Kazu Hirata
Date: 2025-09-29T09:43:52-07:00
New Revision: 3e54505b439923a34fe5cbf27743d883fb62ad9f

URL: https://github.com/llvm/llvm-project/commit/3e54505b439923a34fe5cbf27743d883fb62ad9f
DIFF: https://github.com/llvm/llvm-project/commit/3e54505b439923a34fe5cbf27743d883fb62ad9f.diff

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

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.

Added: 
    

Modified: 
    llvm/include/llvm/ADT/EquivalenceClasses.h
    llvm/unittests/ADT/EquivalenceClassesTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ADT/EquivalenceClasses.h b/llvm/include/llvm/ADT/EquivalenceClasses.h
index 7df0d1557af87..90d8948734729 100644
--- a/llvm/include/llvm/ADT/EquivalenceClasses.h
+++ b/llvm/include/llvm/ADT/EquivalenceClasses.h
@@ -256,9 +256,11 @@ 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.
-        Pre->Next = nullptr;
+        // 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;
       } 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..8172ff97e5169 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