[PATCH] D158875: [ADT] Fix IntEqClasses::join to return the leader in all cases.

[Joshua Cranmer via Phabricator via llvm-commits]

jcranmer-intel updated this revision to Diff 556481.
jcranmer-intel added a comment.

Add support for path compression to IntEqClasses::findLeader.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D158875/new/

https://reviews.llvm.org/D158875

Files:
  llvm/include/llvm/ADT/IntEqClasses.h
  llvm/lib/Support/IntEqClasses.cpp
  llvm/unittests/ADT/IntEqClassesTest.cpp


Index: llvm/unittests/ADT/IntEqClassesTest.cpp
===================================================================
--- llvm/unittests/ADT/IntEqClassesTest.cpp
+++ llvm/unittests/ADT/IntEqClassesTest.cpp
@@ -103,4 +103,46 @@
   EXPECT_EQ(0u, ec.findLeader(9));
 }
 
+TEST(IntEqClasses, JoinReturnsLeader) {
+  IntEqClasses ec(10);
+  EXPECT_EQ(1u, ec.join(1, 2));
+  EXPECT_EQ(1u, ec.join(1, 3));
+  EXPECT_EQ(0u, ec.join(0, 1));
+  EXPECT_EQ(0u, ec.findLeader(1));
+  // At this point, EC[2] = EC[3] = 1, but EC[1] = 0. Be sure that trying to
+  // join EC[2] and EC[3] returns the correct leader, 0, and not their common
+  // node 1.
+  EXPECT_EQ(0u, ec.join(2, 3));
+  EXPECT_EQ(0u, ec.findLeader(2));
+}
+
+TEST(IntEqClasses, FindLeader) {
+  IntEqClasses ec(10);
+
+  // Join all the elements in such a way that most nodes will not point to the
+  // root immediately.
+  ec.join(8, 9);
+  ec.join(7, 8);
+  ec.join(6, 7);
+  ec.join(5, 6);
+  ec.join(4, 5);
+  ec.join(3, 4);
+  ec.join(2, 3);
+  ec.join(1, 2);
+  ec.join(0, 1);
+
+  // Ensure that any path compression that happens doesn't mess up the results
+  // of finding the leader.
+  EXPECT_EQ(0u, ec.findLeader(9));
+  EXPECT_EQ(0u, ec.findLeader(8));
+  EXPECT_EQ(0u, ec.findLeader(7));
+  EXPECT_EQ(0u, ec.findLeader(6));
+  EXPECT_EQ(0u, ec.findLeader(5));
+  EXPECT_EQ(0u, ec.findLeader(4));
+  EXPECT_EQ(0u, ec.findLeader(3));
+  EXPECT_EQ(0u, ec.findLeader(2));
+  EXPECT_EQ(0u, ec.findLeader(1));
+  EXPECT_EQ(0u, ec.findLeader(0));
+}
+
 } // end anonymous namespace
Index: llvm/lib/Support/IntEqClasses.cpp
===================================================================
--- llvm/lib/Support/IntEqClasses.cpp
+++ llvm/lib/Support/IntEqClasses.cpp
@@ -47,13 +47,19 @@
       eca = EC[a];
     }
 
+  // Make sure that eca is actually the leader of the equivalence class (this
+  // can happen if a and b were already in the same equivalence class).
+  eca = findLeader(eca);
   return eca;
 }
 
 unsigned IntEqClasses::findLeader(unsigned a) const {
   assert(NumClasses == 0 && "findLeader() called after compress().");
-  while (a != EC[a])
-    a = EC[a];
+  while (a != EC[a]) {
+    unsigned eca = EC[a];
+    EC[a] = EC[eca];
+    a = eca;
+  }
   return a;
 }
 
Index: llvm/include/llvm/ADT/IntEqClasses.h
===================================================================
--- llvm/include/llvm/ADT/IntEqClasses.h
+++ llvm/include/llvm/ADT/IntEqClasses.h
@@ -31,7 +31,7 @@
   /// itself.
   ///
   /// When compressed, EC[i] is the equivalence class of i.
-  SmallVector<unsigned, 8> EC;
+  mutable SmallVector<unsigned, 8> EC;
 
   /// NumClasses - The number of equivalence classes when compressed, or 0 when
   /// uncompressed.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D158875.556481.patch
Type: text/x-patch
Size: 2719 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20230911/81267b5e/attachment.bin>