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

[jcranmer-intel]

Summary

The method is specified to return the leader of the result of a join. If
the two members were already in an equivalence class, there is a chance
that it might fail to return the leader. For example, if the array looks
like this:

EC[0] = 0
EC[1] = 0
EC[2] = 1
EC[3] = 1

Then join(2, 3) would have returned 1, which is not the same as the
actual leader 0.

[MatzeB]

Nice catch!

llvm/lib/Support/IntEqClasses.cpp

36–53

Not sure how this would affect worst-case complexity given findLeader does not do any path compression. How about this instead?

[jcranmer-intel]

llvm/lib/Support/IntEqClasses.cpp

36–53

The changes I made wouldn't make the worst-case complexity any worse; it seems to me that it would work better to have findLeader do path compression instead (I further note that there appear to be no callers of IntEqClasses::findLeader in LLVM at present, so there's definitely no compile time regression here).

[jcranmer-intel]

ping?

[MatzeB]

llvm/lib/Support/IntEqClasses.cpp

36–53

The changes I made wouldn't make the worst-case complexity any worse

They obviously won't affect the case when we only call join on members of different sets. However without the changes I proposed (or adding path compression to findLeader) it's not obvious to me that we won't cannot hit bad complexity for the cases of calling join with on members of the same set (this is only a gut feeling without proof, but you're not presenting an argument against that either).

Anyway any reason not to go with my proposed changes? Anyway adding path-compression to findLeader should also work... Either way you should upload one variant :)

[jcranmer-intel]

Add support for path compression to IntEqClasses::findLeader.

llvm/include/llvm/ADT/IntEqClasses.h

 //===-- llvm/ADT/IntEqClasses.h - Equiv. Classes of Integers ----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 ///

 namespace llvm {
 
 class IntEqClasses {
   /// EC - When uncompressed, map each integer to a smaller member of its
   /// equivalence class. The class leader is the smallest member and maps to
   /// 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.
   unsigned NumClasses = 0;
 
 public:
   /// IntEqClasses - Create an equivalence class mapping for 0 .. N-1.
   IntEqClasses(unsigned N = 0) { grow(N); }

llvm/lib/Support/IntEqClasses.cpp

 //===-- llvm/ADT/IntEqClasses.cpp - Equivalence Classes of Integers -------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
void IntEqClasses::grow(unsigned N) {
   while (EC.size() < N)
     EC.push_back(EC.size());
 }
 
 unsigned IntEqClasses::join(unsigned a, unsigned b) {
   assert(NumClasses == 0 && "join() called after compress().");
   unsigned eca = EC[a];
   unsigned ecb = EC[b];
   // Update pointers while searching for the leaders, compressing the paths
   // incrementally. The larger leader will eventually be updated, joining the
   // classes.
   while (eca != ecb)
     if (eca < ecb) {
       EC[b] = eca;
       b = ecb;
       ecb = EC[b];
     } else {
       EC[a] = ecb;
       a = eca;
       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;
 }
 
 void IntEqClasses::compress() {
   if (NumClasses)
     return;
   for (unsigned i = 0, e = EC.size(); i != e; ++i)
     EC[i] = (EC[i] == i) ? NumClasses++ : EC[EC[i]];

llvm/unittests/ADT/IntEqClassesTest.cpp

 //===---- ADT/IntEqClassesTest.cpp - IntEqClasses unit tests ----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
TEST(IntEqClasses, Simple) {
   EXPECT_EQ(4u, ec.findLeader(4));
   EXPECT_EQ(4u, ec.findLeader(5));
   EXPECT_EQ(6u, ec.findLeader(6));
   EXPECT_EQ(6u, ec.findLeader(7));
   EXPECT_EQ(4u, ec.findLeader(8));
   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