[llvm][CycleInfo] Quick look-up for block in cycle.

Use a SetVector to store blocks in a cycle to ensure a quick loop-up when
querying whether the cycle contains a given block. This is along the same lines
as the SmallPtrSet in LoopBase, introduced by commit
be640b28c0cb81b77015baaef20ca2941fc61dea.

To make this work, we also enhance SetVector to support vector operations with
pointers and set operations with const pointers in the same container.

Reviewed By: foad

Differential Revision: https://reviews.llvm.org/D146136

llvm/include/llvm/ADT/GenericCycleImpl.h

@@ -177,8 +177,7 @@ void GenericCycleInfo<ContextT>::moveTopLevelCycleToNewParent(CycleT *NewParent,
   CurrentContainer.pop_back();
   Child->ParentCycle = NewParent;
 
-  NewParent->Blocks.insert(NewParent->Blocks.end(), Child->block_begin(),
-                           Child->block_end());
+  NewParent->Blocks.insert(Child->block_begin(), Child->block_end());
 
   for (auto &It : BlockMapTopLevel)
     if (It.second == Child)
@@ -266,7 +265,7 @@ void GenericCycleInfoCompute<ContextT>::run(BlockT *EntryBlock) {
       } else {
         Info.BlockMap.try_emplace(Block, NewCycle.get());
         assert(!is_contained(NewCycle->Blocks, Block));
-        NewCycle->Blocks.push_back(Block);
+        NewCycle->Blocks.insert(Block);
         ProcessPredecessors(Block);
         Info.BlockMapTopLevel.try_emplace(Block, NewCycle.get());
       }

llvm/include/llvm/ADT/GenericCycleInfo.h

@@ -28,16 +28,12 @@
 #ifndef LLVM_ADT_GENERICCYCLEINFO_H
 #define LLVM_ADT_GENERICCYCLEINFO_H
 
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/GenericSSAContext.h"
 #include "llvm/ADT/GraphTraits.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/iterator.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
-#include "llvm/Support/Printable.h"
 #include "llvm/Support/raw_ostream.h"
-#include <vector>
 
 namespace llvm {
 
@@ -67,7 +63,9 @@ private:
 
   /// Basic blocks that are contained in the cycle, including entry blocks,
   /// and including blocks that are part of a child cycle.
-  std::vector<BlockT *> Blocks;
+  using BlockSetVectorT = SetVector<BlockT *, SmallVector<BlockT *, 8>,
+                                    SmallPtrSet<const BlockT *, 8>>;
+  BlockSetVectorT Blocks;
 
   /// Depth of the cycle in the tree. The root "cycle" is at depth 0.
   ///
@@ -85,7 +83,7 @@ private:
   }
 
   void appendEntry(BlockT *Block) { Entries.push_back(Block); }
-  void appendBlock(BlockT *Block) { Blocks.push_back(Block); }
+  void appendBlock(BlockT *Block) { Blocks.insert(Block); }
 
   GenericCycle(const GenericCycle &) = delete;
   GenericCycle &operator=(const GenericCycle &) = delete;
@@ -110,9 +108,7 @@ public:
   }
 
   /// \brief Return whether \p Block is contained in the cycle.
-  bool contains(const BlockT *Block) const {
-    return is_contained(Blocks, Block);
-  }
+  bool contains(const BlockT *Block) const { return Blocks.contains(Block); }
 
   /// \brief Returns true iff this cycle contains \p C.
   ///
@@ -171,7 +167,7 @@ public:
 
   /// Iteration over blocks in the cycle (including entry blocks).
   //@{
-  using const_block_iterator = typename std::vector<BlockT *>::const_iterator;
+  using const_block_iterator = typename BlockSetVectorT::const_iterator;
 
   const_block_iterator block_begin() const {
     return const_block_iterator{Blocks.begin()};

llvm/include/llvm/ADT/SetVector.h

@@ -35,14 +35,25 @@ namespace llvm {
 /// This adapter class provides a way to keep a set of things that also has the
 /// property of a deterministic iteration order. The order of iteration is the
 /// order of insertion.
+///
+/// The key and value types are derived from the Set and Vector types
+/// respectively. This allows the vector-type operations and set-type operations
+/// to have different types. In particular, this is useful when storing pointers
+/// as "Foo *" values but looking them up as "const Foo *" keys.
+///
+/// No constraint is placed on the key and value types, although it is assumed
+/// that value_type can be converted into key_type for insertion. Users must be
+/// aware of any loss of information in this conversion. For example, setting
+/// value_type to float and key_type to int can produce very surprising results,
+/// but it is not explicitly disallowed.
 template <typename T, typename Vector = std::vector<T>,
           typename Set = DenseSet<T>>
 class SetVector {
 public:
-  using value_type = T;
-  using key_type = T;
-  using reference = T&;
-  using const_reference = const T&;
+  using value_type = typename Vector::value_type;
+  using key_type = typename Set::key_type;
+  using reference = value_type &;
+  using const_reference = const value_type &;
   using set_type = Set;
   using vector_type = Vector;
   using iterator = typename vector_type::const_iterator;
@@ -60,7 +71,7 @@ public:
     insert(Start, End);
   }
 
-  ArrayRef<T> getArrayRef() const { return vector_; }
+  ArrayRef<value_type> getArrayRef() const { return vector_; }
 
   /// Clear the SetVector and return the underlying vector.
   Vector takeVector() {
@@ -119,13 +130,13 @@ public:
   }
 
   /// Return the first element of the SetVector.
-  const T &front() const {
+  const value_type &front() const {
     assert(!empty() && "Cannot call front() on empty SetVector!");
     return vector_.front();
   }
 
   /// Return the last element of the SetVector.
-  const T &back() const {
+  const value_type &back() const {
     assert(!empty() && "Cannot call back() on empty SetVector!");
     return vector_.back();
   }
@@ -222,8 +233,8 @@ public:
     vector_.pop_back();
   }
 
-  [[nodiscard]] T pop_back_val() {
-    T Ret = back();
+  [[nodiscard]] value_type pop_back_val() {
+    value_type Ret = back();
     pop_back();
     return Ret;
   }

llvm/include/llvm/ADT/SmallSet.h

@@ -149,7 +149,9 @@ class SmallSet {
   static_assert(N <= 32, "N should be small");
 
 public:
+  using key_type = T;
   using size_type = size_t;
+  using value_type = T;
   using const_iterator = SmallSetIterator<T, N, C>;
 
   SmallSet() = default;

llvm/unittests/ADT/SetVectorTest.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
@@ -48,3 +49,40 @@ TEST(SetVector, ContainsTest) {
   S.remove(2);
   EXPECT_FALSE(S.contains(2));
 }
+
+TEST(SetVector, ConstPtrKeyTest) {
+  SetVector<int *, SmallVector<int *, 8>, SmallPtrSet<const int *, 8>> S, T;
+  int i, j, k, m, n;
+
+  S.insert(&i);
+  S.insert(&j);
+  S.insert(&k);
+
+  EXPECT_TRUE(S.contains(&i));
+  EXPECT_TRUE(S.contains(&j));
+  EXPECT_TRUE(S.contains(&k));
+
+  EXPECT_TRUE(S.contains((const int *)&i));
+  EXPECT_TRUE(S.contains((const int *)&j));
+  EXPECT_TRUE(S.contains((const int *)&k));
+
+  EXPECT_TRUE(S.contains(S[0]));
+  EXPECT_TRUE(S.contains(S[1]));
+  EXPECT_TRUE(S.contains(S[2]));
+
+  S.remove(&k);
+  EXPECT_FALSE(S.contains(&k));
+  EXPECT_FALSE(S.contains((const int *)&k));
+
+  T.insert(&j);
+  T.insert(&m);
+  T.insert(&n);
+
+  EXPECT_TRUE(S.set_union(T));
+  EXPECT_TRUE(S.contains(&m));
+  EXPECT_TRUE(S.contains((const int *)&m));
+
+  S.set_subtract(T);
+  EXPECT_FALSE(S.contains(&j));
+  EXPECT_FALSE(S.contains((const int *)&j));
+}