[EquivalenceClasses] Introduce erase member function (#134660)

Introduce 'erase(const ElemTy &V)' member function to allow the deletion
of a certain value from EquivClasses. This is essential for certain
scenarios that require modifying the contents of EquivClasses.

---------

Co-authored-by: Florian Hahn <[email protected]>

Author: cxy-1993

llvm/include/llvm/ADT/EquivalenceClasses.h

@@ -16,6 +16,7 @@
 #define LLVM_ADT_EQUIVALENCECLASSES_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/Allocator.h"
@@ -220,6 +221,57 @@ template <class ElemTy> class EquivalenceClasses {
     return *ECV;
   }
 
+  /// erase - Erase a value from the union/find set, return "true" if erase
+  /// succeeded, or "false" when the value was not found.
+  bool erase(const ElemTy &V) {
+    if (!TheMapping.contains(V))
+      return false;
+    const ECValue *Cur = TheMapping[V];
+    const ECValue *Next = Cur->getNext();
+    // If the current element is the leader and has a successor element,
+    // update the successor element's 'Leader' field to be the last element,
+    // set the successor element's stolen bit, and set the 'Leader' field of
+    // all other elements in same class to be the successor element.
+    if (Cur->isLeader() && Next) {
+      Next->Leader = Cur->Leader;
+      Next->Next = reinterpret_cast<const ECValue *>(
+          reinterpret_cast<intptr_t>(Next->Next) | static_cast<intptr_t>(1));
+
+      const ECValue *NewLeader = Next;
+      while ((Next = Next->getNext())) {
+        Next->Leader = NewLeader;
+      }
+    } else if (!Cur->isLeader()) {
+      const ECValue *Leader = findLeader(V).Node;
+      const ECValue *Pre = Leader;
+      while (Pre->getNext() != Cur) {
+        Pre = Pre->getNext();
+      }
+      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;
+        Leader->Leader = Pre;
+      } else {
+        // If the current element is in the middle of class, then simply
+        // connect the predecessor element and the successor element.
+        Pre->Next = reinterpret_cast<const ECValue *>(
+            reinterpret_cast<intptr_t>(Next) |
+            static_cast<intptr_t>(Pre->isLeader()));
+        Next->Leader = Pre;
+      }
+    }
+
+    // Update 'TheMapping' and 'Members'.
+    assert(TheMapping.contains(V) && "Can't find input in TheMapping!");
+    TheMapping.erase(V);
+    auto I = find(Members, Cur);
+    assert(I != Members.end() && "Can't find input in members!");
+    Members.erase(I);
+    return true;
+  }
+
   /// findLeader - Given a value in the set, return a member iterator for the
   /// equivalence class it is in.  This does the path-compression part that
   /// makes union-find "union findy".  This returns an end iterator if the value

llvm/unittests/ADT/EquivalenceClassesTest.cpp

@@ -59,6 +59,55 @@ TEST(EquivalenceClassesTest, SimpleMerge2) {
       EXPECT_TRUE(EqClasses.isEquivalent(i, j));
 }
 
+TEST(EquivalenceClassesTest, SimpleErase1) {
+  EquivalenceClasses<int> EqClasses;
+  // Check that erase head success.
+  // After erase A from (A, B ,C, D), <B, C, D> belong to one set.
+  EqClasses.unionSets(0, 1);
+  EqClasses.unionSets(2, 3);
+  EqClasses.unionSets(0, 2);
+  EXPECT_TRUE(EqClasses.erase(0));
+  for (int i = 1; i < 4; ++i)
+    for (int j = 1; j < 4; ++j)
+      EXPECT_TRUE(EqClasses.isEquivalent(i, j));
+}
+
+TEST(EquivalenceClassesTest, SimpleErase2) {
+  EquivalenceClasses<int> EqClasses;
+  // Check that erase tail success.
+  // After erase D from (A, B ,C, D), <A, B, C> belong to one set.
+  EqClasses.unionSets(0, 1);
+  EqClasses.unionSets(2, 3);
+  EqClasses.unionSets(0, 2);
+  EXPECT_TRUE(EqClasses.erase(3));
+  for (int i = 0; i < 3; ++i)
+    for (int j = 0; j < 3; ++j)
+      EXPECT_TRUE(EqClasses.isEquivalent(i, j));
+}
+
+TEST(EquivalenceClassesTest, SimpleErase3) {
+  EquivalenceClasses<int> EqClasses;
+  // Check that erase a value in the middle success.
+  // After erase B from (A, B ,C, D), <A, C, D> belong to one set.
+  EqClasses.unionSets(0, 1);
+  EqClasses.unionSets(2, 3);
+  EqClasses.unionSets(0, 2);
+  EXPECT_TRUE(EqClasses.erase(1));
+  for (int i = 0; i < 3; ++i)
+    for (int j = 0; j < 3; ++j)
+      EXPECT_TRUE(EqClasses.isEquivalent(i, j) ^ ((i == 1) ^ (j == 1)));
+}
+
+TEST(EquivalenceClassesTest, SimpleErase4) {
+  EquivalenceClasses<int> EqClasses;
+  // Check that erase a single class success.
+  EqClasses.insert(0);
+  EXPECT_TRUE(EqClasses.getNumClasses() == 1);
+  EXPECT_TRUE(EqClasses.erase(0));
+  EXPECT_TRUE(EqClasses.getNumClasses() == 0);
+  EXPECT_FALSE(EqClasses.erase(1));
+}
+
 TEST(EquivalenceClassesTest, TwoSets) {
   EquivalenceClasses<int> EqClasses;
   // Form sets of odd and even numbers, check that we split them into these