3 files changed, 96 insertions, 0 deletions
diff --git a/llvm/include/llvm/ADT/EquivalenceClasses.h b/llvm/include/llvm/ADT/EquivalenceClasses.h
index af293d4c142..e3f48433c69 100644
--- a/llvm/include/llvm/ADT/EquivalenceClasses.h
+++ b/llvm/include/llvm/ADT/EquivalenceClasses.h
@@ -239,6 +239,16 @@ public:
     return L1;
   }
 
+  // isEquivalent - Return true if V1 is equivalent to V2. This can happen if
+  // V1 is equal to V2 or if they belong to one equivalence class.
+  bool isEquivalent(const ElemTy &V1, const ElemTy &V2) const {
+    // Fast path: any element is equivalent to itself.
+    if (V1 == V2)
+      return true;
+    auto It = findLeader(V1);
+    return It != member_end() && It == findLeader(V2);
+  }
+
   class member_iterator : public std::iterator<std::forward_iterator_tag,
                                                const ElemTy, ptrdiff_t> {
     friend class EquivalenceClasses;
diff --git a/llvm/unittests/ADT/CMakeLists.txt b/llvm/unittests/ADT/CMakeLists.txt
index e5b6bbf2a23..c0d511000f6 100644
--- a/llvm/unittests/ADT/CMakeLists.txt
+++ b/llvm/unittests/ADT/CMakeLists.txt
@@ -16,6 +16,7 @@ set(ADTSources
   DenseMapTest.cpp
   DenseSetTest.cpp
   DepthFirstIteratorTest.cpp
+  EquivalenceClassesTest.cpp
   FoldingSet.cpp
   FunctionRefTest.cpp
   HashingTest.cpp
diff --git a/llvm/unittests/ADT/EquivalenceClassesTest.cpp b/llvm/unittests/ADT/EquivalenceClassesTest.cpp
new file mode 100644
index 00000000000..57d588a4c51
--- /dev/null
+++ b/llvm/unittests/ADT/EquivalenceClassesTest.cpp
@@ -0,0 +1,85 @@
+//=== llvm/unittest/ADT/EquivalenceClassesTest.cpp - the structure tests --===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/EquivalenceClasses.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace llvm {
+
+TEST(EquivalenceClassesTest, NoMerges) {
+  EquivalenceClasses<int> EqClasses;
+  // Until we merged any sets, check that every element is only equivalent to
+  // itself.
+  for (int i = 0; i < 3; i++)
+    for (int j = 0; j < 3; j++)
+      if (i == j)
+        EXPECT_TRUE(EqClasses.isEquivalent(i, j));
+      else
+        EXPECT_FALSE(EqClasses.isEquivalent(i, j));
+}
+
+TEST(EquivalenceClassesTest, SimpleMerge1) {
+  EquivalenceClasses<int> EqClasses;
+  // Check that once we merge (A, B), (B, C), (C, D), then all elements belong
+  // to one set.
+  EqClasses.unionSets(0, 1);
+  EqClasses.unionSets(1, 2);
+  EqClasses.unionSets(2, 3);
+  for (int i = 0; i < 4; ++i)
+    for (int j = 0; j < 4; ++j)
+      EXPECT_TRUE(EqClasses.isEquivalent(i, j));
+}
+
+TEST(EquivalenceClassesTest, SimpleMerge2) {
+  EquivalenceClasses<int> EqClasses;
+  // Check that once we merge (A, B), (C, D), (A, C), then all elements belong
+  // to one set.
+  EqClasses.unionSets(0, 1);
+  EqClasses.unionSets(2, 3);
+  EqClasses.unionSets(0, 2);
+  for (int i = 0; i < 4; ++i)
+    for (int j = 0; j < 4; ++j)
+      EXPECT_TRUE(EqClasses.isEquivalent(i, j));
+}
+
+TEST(EquivalenceClassesTest, TwoSets) {
+  EquivalenceClasses<int> EqClasses;
+  // Form sets of odd and even numbers, check that we split them into these
+  // two sets correcrly.
+  for (int i = 0; i < 30; i += 2)
+    EqClasses.unionSets(0, i);
+  for (int i = 1; i < 30; i += 2)
+    EqClasses.unionSets(1, i);
+
+  for (int i = 0; i < 30; i++)
+    for (int j = 0; j < 30; j++)
+      if (i % 2 == j % 2)
+        EXPECT_TRUE(EqClasses.isEquivalent(i, j));
+      else
+        EXPECT_FALSE(EqClasses.isEquivalent(i, j));
+}
+
+TEST(EquivalenceClassesTest, MultipleSets) {
+  EquivalenceClasses<int> EqClasses;
+  // Split numbers from [0, 100) into sets so that values in the same set have
+  // equal remainders (mod 17).
+  for (int i = 0; i < 100; i++)
+    EqClasses.unionSets(i % 17, i);
+
+  for (int i = 0; i < 100; i++)
+    for (int j = 0; j < 100; j++)
+      if (i % 17 == j % 17)
+        EXPECT_TRUE(EqClasses.isEquivalent(i, j));
+      else
+        EXPECT_FALSE(EqClasses.isEquivalent(i, j));
+}
+
+} // llvm