[immutable-ptr-set] Change data structure to use T instead of T *, but require T to conform to PointerLikeTypeTraits.

This allows one to place types like PointerIntPair and value types that wrap a
pointer into the pointer set.

I am making this change before I use this data structure in TransferNonSendable
and using ARC to validate that I haven't broken anything.

Author: gottesmm

include/swift/Basic/ImmutablePointerSet.h

@@ -58,19 +58,22 @@
 
 namespace swift {
 
-template <typename PtrTy> class ImmutablePointerSetFactory;
+template <typename T>
+class ImmutablePointerSetFactory;
 
 /// An immutable set of pointers. It is backed by a tail allocated sorted array
 /// ref.
-template <typename T> class ImmutablePointerSet : public llvm::FoldingSetNode {
-  using PtrTy = typename std::add_pointer<T>::type;
+template <typename T>
+class ImmutablePointerSet : public llvm::FoldingSetNode {
   friend class ImmutablePointerSetFactory<T>;
 
+  using PtrTraits = llvm::PointerLikeTypeTraits<T>;
+
   NullablePtr<ImmutablePointerSetFactory<T>> ParentFactory;
-  llvm::ArrayRef<PtrTy> Data;
+  llvm::ArrayRef<T> Data;
 
   ImmutablePointerSet(ImmutablePointerSetFactory<T> *ParentFactory,
-                      llvm::ArrayRef<PtrTy> NewData)
+                      llvm::ArrayRef<T> NewData)
       : ParentFactory(ParentFactory), Data(NewData) {}
 
 public:
@@ -100,7 +103,7 @@ template <typename T> class ImmutablePointerSet : public llvm::FoldingSetNode {
     return !(*this == P);
   }
 
-  unsigned count(PtrTy Ptr) const {
+  unsigned count(T Ptr) const {
     // This returns the first element >= Ptr. Since we know that our array is
     // sorted and uniqued, Ptr must be that element.
     auto LowerBound = std::lower_bound(begin(), end(), Ptr);
@@ -113,7 +116,7 @@ template <typename T> class ImmutablePointerSet : public llvm::FoldingSetNode {
     return *LowerBound == Ptr;
   }
 
-  using iterator = typename llvm::ArrayRef<PtrTy>::iterator;
+  using iterator = typename llvm::ArrayRef<T>::iterator;
   iterator begin() const { return Data.begin(); }
   iterator end() const { return Data.end(); }
 
@@ -122,8 +125,8 @@ template <typename T> class ImmutablePointerSet : public llvm::FoldingSetNode {
 
   void Profile(llvm::FoldingSetNodeID &ID) const {
     assert(!Data.empty() && "Should not profile empty ImmutablePointerSet");
-    for (PtrTy P : Data) {
-      ID.AddPointer(P);
+    for (T P : Data) {
+      ID.AddPointer(PtrTraits::getAsVoidPointer(P));
     }
   }
 
@@ -177,10 +180,11 @@ template <typename T> class ImmutablePointerSet : public llvm::FoldingSetNode {
   }
 };
 
-template <typename T> class ImmutablePointerSetFactory {
-  using PtrTy = typename std::add_pointer<T>::type;
-
+template <typename T>
+class ImmutablePointerSetFactory {
   using PtrSet = ImmutablePointerSet<T>;
+  using PtrTraits = typename PtrSet::PtrTraits;
+
   // This is computed out-of-line so that ImmutablePointerSetFactory is
   // treated as a complete type.
   static const unsigned AllocAlignment;
@@ -205,7 +209,7 @@ template <typename T> class ImmutablePointerSetFactory {
 
   /// Given a sorted and uniqued list \p Array, return the ImmutablePointerSet
   /// containing Array. Asserts if \p Array is not sorted and uniqued.
-  PtrSet *get(llvm::ArrayRef<PtrTy> Array) {
+  PtrSet *get(llvm::ArrayRef<T> Array) {
     if (Array.empty())
       return ImmutablePointerSetFactory::getEmptySet();
 
@@ -215,8 +219,8 @@ template <typename T> class ImmutablePointerSetFactory {
     assert(is_sorted_and_uniqued(Array));
 
     llvm::FoldingSetNodeID ID;
-    for (PtrTy Ptr : Array) {
-      ID.AddPointer(Ptr);
+    for (T Ptr : Array) {
+      ID.AddPointer(PtrTraits::getAsVoidPointer(Ptr));
     }
 
     void *InsertPt;
@@ -225,7 +229,7 @@ template <typename T> class ImmutablePointerSetFactory {
     }
 
     size_t NumElts = Array.size();
-    size_t MemSize = sizeof(PtrSet) + sizeof(PtrTy) * NumElts;
+    size_t MemSize = sizeof(PtrSet) + sizeof(T) * NumElts;
 
     // Allocate the memory.
     auto *Mem =
@@ -234,8 +238,8 @@ template <typename T> class ImmutablePointerSetFactory {
     // Copy in the pointers into the tail allocated memory. We do not need to do
     // any sorting/uniquing ourselves since we assume that our users perform
     // this task for us.
-    llvm::MutableArrayRef<PtrTy> DataMem(reinterpret_cast<PtrTy *>(&Mem[1]),
-                                         NumElts);
+    llvm::MutableArrayRef<void *> DataMem(reinterpret_cast<void **>(&Mem[1]),
+                                          NumElts);
     std::copy(Array.begin(), Array.end(), DataMem.begin());
 
     // Allocate the new node and insert it into the Set.
@@ -244,7 +248,35 @@ template <typename T> class ImmutablePointerSetFactory {
     return NewNode;
   }
 
-  PtrSet *merge(PtrSet *S1, llvm::ArrayRef<PtrTy> S2) {
+  PtrSet *get(T value) {
+    llvm::FoldingSetNodeID ID;
+    ID.AddPointer(PtrTraits::getAsVoidPointer(value));
+
+    void *InsertPt;
+    if (auto *PSet = Set.FindNodeOrInsertPos(ID, InsertPt)) {
+      return PSet;
+    }
+
+    size_t NumElts = 1;
+    size_t MemSize = sizeof(PtrSet) + sizeof(T) * NumElts;
+
+    // Allocate the memory.
+    auto *Mem =
+        reinterpret_cast<PtrSet *>(Allocator.Allocate(MemSize, AllocAlignment));
+
+    // Copy in the pointers into the tail allocated memory. We do not need to do
+    // any sorting/uniquing ourselves since we assume that our users perform
+    // this task for us.
+    llvm::MutableArrayRef<T> DataMem(reinterpret_cast<T *>(&Mem[1]), NumElts);
+    DataMem[0] = value;
+
+    // Allocate the new node and insert it into the Set.
+    auto *NewNode = new (Mem) PtrSet(this, DataMem);
+    Set.InsertNode(NewNode, InsertPt);
+    return NewNode;
+  }
+
+  PtrSet *merge(PtrSet *S1, llvm::ArrayRef<T> S2) {
     if (S1->empty())
       return get(S2);
 
@@ -266,7 +298,7 @@ template <typename T> class ImmutablePointerSetFactory {
     // perform a sorted set merging algorithm to create the ID. We also count
     // the number of unique elements for allocation purposes.
     unsigned NumElts = 0;
-    set_union_for_each(*S1, S2, [&ID, &NumElts](const PtrTy Ptr) -> void {
+    set_union_for_each(*S1, S2, [&ID, &NumElts](const T Ptr) -> void {
       ID.AddPointer(Ptr);
       NumElts++;
     });
@@ -277,7 +309,7 @@ template <typename T> class ImmutablePointerSetFactory {
       return PSet;
     }
 
-    unsigned MemSize = sizeof(PtrSet) + sizeof(PtrTy) * NumElts;
+    unsigned MemSize = sizeof(PtrSet) + sizeof(T) * NumElts;
 
     // Allocate the memory.
     auto *Mem =
@@ -286,8 +318,7 @@ template <typename T> class ImmutablePointerSetFactory {
     // Copy in the union of the two pointer sets into the tail allocated
     // memory. Since we know that our sorted arrays are uniqued, we can use
     // set_union to get the uniqued sorted array that we want.
-    llvm::MutableArrayRef<PtrTy> DataMem(reinterpret_cast<PtrTy *>(&Mem[1]),
-                                         NumElts);
+    llvm::MutableArrayRef<T> DataMem(reinterpret_cast<T *>(&Mem[1]), NumElts);
     std::set_union(S1->begin(), S1->end(), S2.begin(), S2.end(),
                    DataMem.begin());
 
@@ -316,8 +347,8 @@ template <typename T> class ImmutablePointerSetFactory {
     // perform a sorted set merging algorithm to create the ID. We also count
     // the number of unique elements for allocation purposes.
     unsigned NumElts = 0;
-    set_union_for_each(*S1, *S2, [&ID, &NumElts](const PtrTy Ptr) -> void {
-      ID.AddPointer(Ptr);
+    set_union_for_each(*S1, *S2, [&ID, &NumElts](const T Ptr) -> void {
+      ID.AddPointer(PtrTraits::getAsVoidPointer(Ptr));
       NumElts++;
     });
 
@@ -327,7 +358,7 @@ template <typename T> class ImmutablePointerSetFactory {
       return PSet;
     }
 
-    unsigned MemSize = sizeof(PtrSet) + sizeof(PtrTy) * NumElts;
+    unsigned MemSize = sizeof(PtrSet) + sizeof(T) * NumElts;
 
     // Allocate the memory.
     auto *Mem =
@@ -336,8 +367,7 @@ template <typename T> class ImmutablePointerSetFactory {
     // Copy in the union of the two pointer sets into the tail allocated
     // memory. Since we know that our sorted arrays are uniqued, we can use
     // set_union to get the uniqued sorted array that we want.
-    llvm::MutableArrayRef<PtrTy> DataMem(reinterpret_cast<PtrTy *>(&Mem[1]),
-                                         NumElts);
+    llvm::MutableArrayRef<T> DataMem(reinterpret_cast<T *>(&Mem[1]), NumElts);
     std::set_union(S1->begin(), S1->end(), S2->begin(), S2->end(),
                    DataMem.begin());
 
@@ -356,8 +386,8 @@ template <typename T>
 #if !defined(_MSC_VER) || defined(__clang__)
 constexpr
 #endif
-const unsigned ImmutablePointerSetFactory<T>::AllocAlignment =
-    (alignof(PtrSet) > alignof(PtrTy)) ? alignof(PtrSet) : alignof(PtrTy);
+    const unsigned ImmutablePointerSetFactory<T>::AllocAlignment =
+        (alignof(PtrSet) > alignof(T)) ? alignof(PtrSet) : alignof(T);
 
 } // end swift namespace
 

lib/SILOptimizer/ARC/ARCRegionState.cpp

@@ -166,7 +166,7 @@ bool ARCRegionState::processBlockBottomUp(
     EpilogueARCFunctionInfo *EAFI, LoopRegionFunctionInfo *LRFI,
     bool FreezeOwnedArgEpilogueReleases,
     BlotMapVector<SILInstruction *, BottomUpRefCountState> &IncToDecStateMap,
-    ImmutablePointerSetFactory<SILInstruction> &SetFactory) {
+    ImmutablePointerSetFactory<SILInstruction *> &SetFactory) {
   LLVM_DEBUG(llvm::dbgs() << ">>>> Bottom Up!\n");
 
   SILBasicBlock &BB = *R->getBlock();
@@ -311,7 +311,7 @@ bool ARCRegionState::processBottomUp(
     llvm::DenseSet<SILInstruction *> &UnmatchedRefCountInsts,
     BlotMapVector<SILInstruction *, BottomUpRefCountState> &IncToDecStateMap,
     llvm::DenseMap<const LoopRegion *, ARCRegionState *> &RegionStateInfo,
-    ImmutablePointerSetFactory<SILInstruction> &SetFactory) {
+    ImmutablePointerSetFactory<SILInstruction *> &SetFactory) {
   const LoopRegion *R = getRegion();
 
   // We only process basic blocks for now. This ensures that we always propagate
@@ -331,7 +331,7 @@ bool ARCRegionState::processBottomUp(
 bool ARCRegionState::processBlockTopDown(
     SILBasicBlock &BB, AliasAnalysis *AA, RCIdentityFunctionInfo *RCIA,
     BlotMapVector<SILInstruction *, TopDownRefCountState> &DecToIncStateMap,
-    ImmutablePointerSetFactory<SILInstruction> &SetFactory) {
+    ImmutablePointerSetFactory<SILInstruction *> &SetFactory) {
   LLVM_DEBUG(llvm::dbgs() << ">>>> Top Down!\n");
 
   bool NestingDetected = false;
@@ -460,7 +460,7 @@ bool ARCRegionState::processTopDown(
     llvm::DenseSet<SILInstruction *> &UnmatchedRefCountInsts,
     BlotMapVector<SILInstruction *, TopDownRefCountState> &DecToIncStateMap,
     llvm::DenseMap<const LoopRegion *, ARCRegionState *> &RegionStateInfo,
-    ImmutablePointerSetFactory<SILInstruction> &SetFactory) {
+    ImmutablePointerSetFactory<SILInstruction *> &SetFactory) {
   const LoopRegion *R = getRegion();
 
   // We only process basic blocks for now. This ensures that we always propagate

lib/SILOptimizer/ARC/ARCRegionState.h

@@ -191,7 +191,7 @@ class ARCRegionState {
       llvm::DenseSet<SILInstruction *> &UnmatchedRefCountInsts,
       BlotMapVector<SILInstruction *, TopDownRefCountState> &DecToIncStateMap,
       llvm::DenseMap<const LoopRegion *, ARCRegionState *> &LoopRegionState,
-      ImmutablePointerSetFactory<SILInstruction> &SetFactory);
+      ImmutablePointerSetFactory<SILInstruction *> &SetFactory);
 
   /// If this region is a block, process all instructions bottom up. Otherwise,
   /// apply the summarized bottom up information to the merged bottom up
@@ -204,7 +204,7 @@ class ARCRegionState {
       llvm::DenseSet<SILInstruction *> &UnmatchedRefCountInsts,
       BlotMapVector<SILInstruction *, BottomUpRefCountState> &IncToDecStateMap,
       llvm::DenseMap<const LoopRegion *, ARCRegionState *> &RegionStateInfo,
-      ImmutablePointerSetFactory<SILInstruction> &SetFactory);
+      ImmutablePointerSetFactory<SILInstruction *> &SetFactory);
 
   void summarizeBlock(SILBasicBlock *BB);
 
@@ -223,10 +223,10 @@ class ARCRegionState {
 private:
   bool processBlockBottomUp(
       const LoopRegion *R, AliasAnalysis *AA, RCIdentityFunctionInfo *RCIA,
-      EpilogueARCFunctionInfo *EAFI,
-      LoopRegionFunctionInfo *LRFI, bool FreezeOwnedArgEpilogueReleases,
+      EpilogueARCFunctionInfo *EAFI, LoopRegionFunctionInfo *LRFI,
+      bool FreezeOwnedArgEpilogueReleases,
       BlotMapVector<SILInstruction *, BottomUpRefCountState> &IncToDecStateMap,
-      ImmutablePointerSetFactory<SILInstruction> &SetFactory);
+      ImmutablePointerSetFactory<SILInstruction *> &SetFactory);
   bool processLoopBottomUp(
       const LoopRegion *R, AliasAnalysis *AA, LoopRegionFunctionInfo *LRFI,
       RCIdentityFunctionInfo *RCIA,
@@ -236,7 +236,7 @@ class ARCRegionState {
   bool processBlockTopDown(
       SILBasicBlock &BB, AliasAnalysis *AA, RCIdentityFunctionInfo *RCIA,
       BlotMapVector<SILInstruction *, TopDownRefCountState> &DecToIncStateMap,
-      ImmutablePointerSetFactory<SILInstruction> &SetFactory);
+      ImmutablePointerSetFactory<SILInstruction *> &SetFactory);
   bool processLoopTopDown(
       const LoopRegion *R, ARCRegionState *State, AliasAnalysis *AA,
       LoopRegionFunctionInfo *LRFI, RCIdentityFunctionInfo *RCIA,

lib/SILOptimizer/ARC/GlobalARCSequenceDataflow.h

@@ -64,7 +64,7 @@ class ARCSequenceDataflowEvaluator {
   BlotMapVector<SILInstruction *, BottomUpRefCountState> &IncToDecStateMap;
 
   llvm::BumpPtrAllocator Allocator;
-  ImmutablePointerSetFactory<SILInstruction> SetFactory;
+  ImmutablePointerSetFactory<SILInstruction *> SetFactory;
 
   /// Stashed BB information.
   std::unique_ptr<ARCBBStateInfo> BBStateInfo;

lib/SILOptimizer/ARC/GlobalLoopARCSequenceDataflow.h

@@ -40,7 +40,7 @@ class LoopARCSequenceDataflowEvaluator {
   llvm::BumpPtrAllocator Allocator;
 
   /// The factory that we use to generate immutable pointer sets.
-  ImmutablePointerSetFactory<SILInstruction> SetFactory;
+  ImmutablePointerSetFactory<SILInstruction *> SetFactory;
 
   /// The SILFunction that we are applying the dataflow to.
   SILFunction &F;

lib/SILOptimizer/ARC/RCStateTransition.h

@@ -73,8 +73,8 @@ class RCStateTransition {
   /// An RCStateTransition can represent either an RC end point (i.e. an initial
   /// or terminal RC transition) or a ptr set of Mutators.
   SILNode *EndPoint;
-  ImmutablePointerSet<SILInstruction> *Mutators =
-      ImmutablePointerSetFactory<SILInstruction>::getEmptySet();
+  ImmutablePointerSet<SILInstruction *> *Mutators =
+      ImmutablePointerSetFactory<SILInstruction *>::getEmptySet();
   RCStateTransitionKind Kind;
 
   // Should only be constructed be default RefCountState.
@@ -84,9 +84,9 @@ class RCStateTransition {
   ~RCStateTransition() = default;
   RCStateTransition(const RCStateTransition &R) = default;
 
-  RCStateTransition(ImmutablePointerSet<SILInstruction> *I) {
+  RCStateTransition(ImmutablePointerSet<SILInstruction *> *I) {
     assert(I->size() == 1);
-    SILInstruction *Inst = *I->begin();
+    auto *Inst = *I->begin();
     Kind = getRCStateTransitionKind(Inst->asSILNode());
     if (isRCStateTransitionEndPoint(Kind)) {
       EndPoint = Inst->asSILNode();

lib/SILOptimizer/ARC/RCStateTransitionVisitors.cpp

@@ -57,7 +57,7 @@ BottomUpDataflowRCStateVisitor<ARCState>::BottomUpDataflowRCStateVisitor(
     RCIdentityFunctionInfo *RCFI, EpilogueARCFunctionInfo *EAFI,
     ARCState &State, bool FreezeOwnedArgEpilogueReleases,
     IncToDecStateMapTy &IncToDecStateMap,
-    ImmutablePointerSetFactory<SILInstruction> &SetFactory)
+    ImmutablePointerSetFactory<SILInstruction *> &SetFactory)
     : RCFI(RCFI), EAFI(EAFI), DataflowState(State),
       FreezeOwnedArgEpilogueReleases(FreezeOwnedArgEpilogueReleases),
       IncToDecStateMap(IncToDecStateMap), SetFactory(SetFactory) {}
@@ -184,7 +184,7 @@ template <class ARCState>
 TopDownDataflowRCStateVisitor<ARCState>::TopDownDataflowRCStateVisitor(
     RCIdentityFunctionInfo *RCFI, ARCState &DataflowState,
     DecToIncStateMapTy &DecToIncStateMap,
-    ImmutablePointerSetFactory<SILInstruction> &SetFactory)
+    ImmutablePointerSetFactory<SILInstruction *> &SetFactory)
     : RCFI(RCFI), DataflowState(DataflowState),
       DecToIncStateMap(DecToIncStateMap), SetFactory(SetFactory) {}
 

lib/SILOptimizer/ARC/RCStateTransitionVisitors.h

@@ -120,14 +120,14 @@ class BottomUpDataflowRCStateVisitor
   ARCState &DataflowState;
   bool FreezeOwnedArgEpilogueReleases;
   BlotMapVector<SILInstruction *, BottomUpRefCountState> &IncToDecStateMap;
-  ImmutablePointerSetFactory<SILInstruction> &SetFactory;
+  ImmutablePointerSetFactory<SILInstruction *> &SetFactory;
 
 public:
   BottomUpDataflowRCStateVisitor(
       RCIdentityFunctionInfo *RCFI, EpilogueARCFunctionInfo *EAFI,
       ARCState &DataflowState, bool FreezeOwnedArgEpilogueReleases,
       IncToDecStateMapTy &IncToDecStateMap,
-      ImmutablePointerSetFactory<SILInstruction> &SetFactory);
+      ImmutablePointerSetFactory<SILInstruction *> &SetFactory);
   DataflowResult visitAutoreleasePoolCall(SILNode *N);
   DataflowResult visitStrongDecrement(SILNode *N);
   DataflowResult visitStrongIncrement(SILNode *N);
@@ -156,13 +156,13 @@ class TopDownDataflowRCStateVisitor
   RCIdentityFunctionInfo *RCFI;
   ARCState &DataflowState;
   DecToIncStateMapTy &DecToIncStateMap;
-  ImmutablePointerSetFactory<SILInstruction> &SetFactory;
+  ImmutablePointerSetFactory<SILInstruction *> &SetFactory;
 
 public:
   TopDownDataflowRCStateVisitor(
       RCIdentityFunctionInfo *RCFI, ARCState &State,
       DecToIncStateMapTy &DecToIncStateMap,
-      ImmutablePointerSetFactory<SILInstruction> &SetFactory);
+      ImmutablePointerSetFactory<SILInstruction *> &SetFactory);
   DataflowResult visitAutoreleasePoolCall(SILNode *N);
   DataflowResult visitStrongDecrement(SILNode *N);
   DataflowResult visitStrongIncrement(SILNode *N);