From 27ca4837eec1734d23cc628fcaddf03b94afc845 Mon Sep 17 00:00:00 2001 From: donald chen Date: Thu, 10 Apr 2025 12:42:10 +0800 Subject: [PATCH] [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 --- llvm/include/llvm/ADT/EquivalenceClasses.h | 52 +++++++++++++++++++ llvm/unittests/ADT/EquivalenceClassesTest.cpp | 49 +++++++++++++++++ 2 files changed, 101 insertions(+) diff --git a/llvm/include/llvm/ADT/EquivalenceClasses.h b/llvm/include/llvm/ADT/EquivalenceClasses.h index e0a7af9421c3..0d44d2c79dd2 100644 --- a/llvm/include/llvm/ADT/EquivalenceClasses.h +++ b/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 @@ public: 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( + reinterpret_cast(Next->Next) | static_cast(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( + reinterpret_cast(Next) | + static_cast(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 diff --git a/llvm/unittests/ADT/EquivalenceClassesTest.cpp b/llvm/unittests/ADT/EquivalenceClassesTest.cpp index ff243f51102f..3d5c48eb8e1b 100644 --- a/llvm/unittests/ADT/EquivalenceClassesTest.cpp +++ b/llvm/unittests/ADT/EquivalenceClassesTest.cpp @@ -59,6 +59,55 @@ TEST(EquivalenceClassesTest, SimpleMerge2) { EXPECT_TRUE(EqClasses.isEquivalent(i, j)); } +TEST(EquivalenceClassesTest, SimpleErase1) { + EquivalenceClasses EqClasses; + // Check that erase head success. + // After erase A from (A, 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 EqClasses; + // Check that erase tail success. + // After erase D from (A, B ,C, D), 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 EqClasses; + // Check that erase a value in the middle success. + // After erase B from (A, B ,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 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 EqClasses; // Form sets of odd and even numbers, check that we split them into these