[ADT] Fix a bug in EquivalenceClasses::erase

[kazutakahirata]

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.

[llvmbot]

@llvm/pr-subscribers-llvm-adt

Author: Kazu Hirata (kazutakahirata)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/161121.diff

2 Files Affected:

• (modified) llvm/include/llvm/ADT/EquivalenceClasses.h (+2-1)
• (modified) llvm/unittests/ADT/EquivalenceClassesTest.cpp (+23)

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

[fhahn]

LGTM, thanks