SparseSet: Add support for key-derived indexes and arbitrary key types.

This nicely handles the most common case of virtual register sets, but
also handles anticipated cases where we will map pointers to IDs.

The goal is not to develop a completely generic SparseSet
template. Instead we want to handle the expected uses within llvm
without any template antics in the client code. I'm adding a bit of
template nastiness here, and some assumption about expected usage in
order to make the client code very clean.

The expected common uses cases I'm designing for:
- integer keys that need to be reindexed, and may map to additional
  data
- densely numbered objects where we want pointer keys because no
  number->object map exists.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@155227 91177308-0d34-0410-b5e6-96231b3b80d8

7 files changed, 95 insertions, 54 deletions

diff --git a/include/llvm/ADT/IndexedMap.h b/include/llvm/ADT/IndexedMap.h
index 87126ea491..2ffb5058e5 100644
--- a/include/llvm/ADT/IndexedMap.h
+++ b/include/llvm/ADT/IndexedMap.h
@@ -20,19 +20,14 @@
 #ifndef LLVM_ADT_INDEXEDMAP_H
 #define LLVM_ADT_INDEXEDMAP_H
 
+#include "llvm/ADT/STLExtras.h"
 #include <cassert>
 #include <functional>
 #include <vector>
 
 namespace llvm {
 
-  struct IdentityFunctor : public std::unary_function<unsigned, unsigned> {
-    unsigned operator()(unsigned Index) const {
-      return Index;
-    }
-  };
-
-  template <typename T, typename ToIndexT = IdentityFunctor>
+template <typename T, typename ToIndexT = llvm::identity<unsigned> >
   class IndexedMap {
     typedef typename ToIndexT::argument_type IndexT;
     typedef std::vector<T> StorageT;
diff --git a/include/llvm/ADT/STLExtras.h b/include/llvm/ADT/STLExtras.h
index 5da906dc8c..58a23d886f 100644
--- a/include/llvm/ADT/STLExtras.h
+++ b/include/llvm/ADT/STLExtras.h
@@ -30,6 +30,16 @@ namespace llvm {
 //===----------------------------------------------------------------------===//
 
 template<class Ty>
+struct identity : public std::unary_function<Ty, Ty> {
+  Ty &operator()(Ty &self) const {
+    return self;
+  }
+  const Ty &operator()(const Ty &self) const {
+    return self;
+  }
+};
+
+template<class Ty>
 struct less_ptr : public std::binary_function<Ty, Ty, bool> {
   bool operator()(const Ty* left, const Ty* right) const {
     return *left < *right;
@@ -274,7 +284,7 @@ static inline void array_pod_sort(IteratorTy Start, IteratorTy End,
   if (Start == End) return;
   qsort(&*Start, End-Start, sizeof(*Start), Compare);
 }
-  
+
 //===----------------------------------------------------------------------===//
 //     Extra additions to <algorithm>
 //===----------------------------------------------------------------------===//
diff --git a/include/llvm/ADT/SparseSet.h b/include/llvm/ADT/SparseSet.h
index 923c6a5954..5569633348 100644
--- a/include/llvm/ADT/SparseSet.h
+++ b/include/llvm/ADT/SparseSet.h
@@ -21,33 +21,62 @@
 #define LLVM_ADT_SPARSESET_H
 
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/DataTypes.h"
 #include <limits>
 
 namespace llvm {
 
-/// SparseSetFunctor - Objects in a SparseSet are identified by small integer
-/// keys.  A functor object is used to compute the key of an object.  The
-/// functor's operator() must return an unsigned smaller than the universe.
+/// SparseSetValTraits - Objects in a SparseSet are identified by keys that can
+/// be uniquely converted to a small integer less than the set's universe. This
+/// class allows the set to hold values that differ from the set's key type as
+/// long as an index can still be derived from the value. SparseSet never
+/// directly compares ValueT, only their indices, so it can map keys to
+/// arbitrary values. SparseSetValTraits computes the index from the value
+/// object. To compute the index from a key, SparseSet uses a separate
+/// KeyFunctorT template argument.
 ///
-/// The default functor implementation forwards to a getSparseSetKey() method
-/// on the object.  It is intended for sparse sets holding ad-hoc structs.
+/// A simple type declaration, SparseSet<Type>, handles these cases:
+/// - unsigned key, identity index, identity value
+/// - unsigned key, identity index, fat value providing getSparseSetIndex()
+///
+/// The type declaration SparseSet<Type, UnaryFunction> handles:
+/// - unsigned key, remapped index, identity value (virtual registers)
+/// - pointer key, pointer-derived index, identity value (node+ID)
+/// - pointer key, pointer-derived index, fat value with getSparseSetIndex()
+///
+/// Only other, unexpected cases require specializing SparseSetValTraits.
+///
+/// For best results, ValueT should not require a destructor.
 ///
 template<typename ValueT>
-struct SparseSetFunctor {
-  unsigned operator()(const ValueT &Val) {
-    return Val.getSparseSetKey();
+struct SparseSetValTraits {
+  static unsigned getValIndex(const ValueT &Val) {
+    return Val.getSparseSetIndex();
   }
 };
 
-/// SparseSetFunctor<unsigned> - Provide a trivial identity functor for
-/// SparseSet<unsigned>.
+/// SparseSetValFunctor - Helper class for selecting SparseSetValTraits. The
+/// generic implementation handles ValueT classes which either provide
+/// getSparseSetIndex() or specialize SparseSetValTraits<>.
 ///
-template<> struct SparseSetFunctor<unsigned> {
-  unsigned operator()(unsigned Val) { return Val; }
+template<typename KeyT, typename ValueT, typename KeyFunctorT>
+struct SparseSetValFunctor {
+  unsigned operator()(const ValueT &Val) const {
+    return SparseSetValTraits<ValueT>::getValIndex(Val);
+  }
 };
 
-/// SparseSet - Fast set implementation for objects that can be identified by
+/// SparseSetValFunctor<KeyT, KeyT> - Helper class for the common case of
+/// identity key/value sets.
+template<typename KeyT, typename KeyFunctorT>
+struct SparseSetValFunctor<KeyT, KeyT, KeyFunctorT> {
+  unsigned operator()(const KeyT &Key) const {
+    return KeyFunctorT()(Key);
+  }
+};
+
+/// SparseSet - Fast set implmentation for objects that can be identified by
 /// small unsigned keys.
 ///
 /// SparseSet allocates memory proportional to the size of the key universe, so
@@ -82,18 +111,20 @@ template<> struct SparseSetFunctor<unsigned> {
 /// uint16_t or uint32_t.
 ///
 /// @param ValueT      The type of objects in the set.
+/// @param KeyFunctorT A functor that computes an unsigned index from KeyT.
 /// @param SparseT     An unsigned integer type. See above.
-/// @param KeyFunctorT A functor that computes the unsigned key of a ValueT.
 ///
 template<typename ValueT,
-         typename SparseT = uint8_t,
-         typename KeyFunctorT = SparseSetFunctor<ValueT> >
+         typename KeyFunctorT = llvm::identity<unsigned>,
+         typename SparseT = uint8_t>
 class SparseSet {
+  typedef typename KeyFunctorT::argument_type KeyT;
   typedef SmallVector<ValueT, 8> DenseT;
   DenseT Dense;
   SparseT *Sparse;
   unsigned Universe;
-  KeyFunctorT KeyOf;
+  KeyFunctorT KeyIndexOf;
+  SparseSetValFunctor<KeyT, ValueT, KeyFunctorT> ValIndexOf;
 
   // Disable copy construction and assignment.
   // This data structure is not meant to be used that way.
@@ -160,21 +191,21 @@ public:
     Dense.clear();
   }
 
-  /// find - Find an element by its key.
+  /// findIndex - Find an element by its index.
   ///
-  /// @param   Key A valid key to find.
+  /// @param   Idx A valid index to find.
   /// @returns An iterator to the element identified by key, or end().
   ///
-  iterator find(unsigned Key) {
-    assert(Key < Universe && "Key out of range");
+  iterator findIndex(unsigned Idx) {
+    assert(Idx < Universe && "Key out of range");
     assert(std::numeric_limits<SparseT>::is_integer &&
            !std::numeric_limits<SparseT>::is_signed &&
            "SparseT must be an unsigned integer type");
     const unsigned Stride = std::numeric_limits<SparseT>::max() + 1u;
-    for (unsigned i = Sparse[Key], e = size(); i < e; i += Stride) {
-      const unsigned FoundKey = KeyOf(Dense[i]);
-      assert(FoundKey < Universe && "Invalid key in set. Did object mutate?");
-      if (Key == FoundKey)
+    for (unsigned i = Sparse[Idx], e = size(); i < e; i += Stride) {
+      const unsigned FoundIdx = ValIndexOf(Dense[i]);
+      assert(FoundIdx < Universe && "Invalid key in set. Did object mutate?");
+      if (Idx == FoundIdx)
         return begin() + i;
       // Stride is 0 when SparseT >= unsigned.  We don't need to loop.
       if (!Stride)
@@ -183,13 +214,22 @@ public:
     return end();
   }
 
-  const_iterator find(unsigned Key) const {
-    return const_cast<SparseSet*>(this)->find(Key);
+  /// find - Find an element by its key.
+  ///
+  /// @param   Key A valid key to find.
+  /// @returns An iterator to the element identified by key, or end().
+  ///
+  iterator find(const KeyT &Key) {
+    return findIndex(KeyIndexOf(Key));
+  }
+
+  const_iterator find(const KeyT &Key) const {
+    return const_cast<SparseSet*>(this)->findIndex(KeyIndexOf(Key));
   }
 
   /// count - Returns true if this set contains an element identified by Key.
   ///
-  bool count(unsigned Key) const {
+  bool count(const KeyT &Key) const {
     return find(Key) != end();
   }
 
@@ -204,11 +244,11 @@ public:
   /// Insertion invalidates all iterators.
   ///
   std::pair<iterator, bool> insert(const ValueT &Val) {
-    unsigned Key = KeyOf(Val);
-    iterator I = find(Key);
+    unsigned Idx = ValIndexOf(Val);
+    iterator I = findIndex(Idx);
     if (I != end())
       return std::make_pair(I, false);
-    Sparse[Key] = size();
+    Sparse[Idx] = size();
     Dense.push_back(Val);
     return std::make_pair(end() - 1, true);
   }
@@ -216,7 +256,7 @@ public:
   /// array subscript - If an element already exists with this key, return it.
   /// Otherwise, automatically construct a new value from Key, insert it,
   /// and return the newly inserted element.
-  ValueT &operator[](unsigned Key) {
+  ValueT &operator[](const KeyT &Key) {
     return *insert(ValueT(Key)).first;
   }
 
@@ -238,9 +278,9 @@ public:
     assert(unsigned(I - begin()) < size() && "Invalid iterator");
     if (I != end() - 1) {
       *I = Dense.back();
-      unsigned BackKey = KeyOf(Dense.back());
-      assert(BackKey < Universe && "Invalid key in set. Did object mutate?");
-      Sparse[BackKey] = I - begin();
+      unsigned BackIdx = ValIndexOf(Dense.back());
+      assert(BackIdx < Universe && "Invalid key in set. Did object mutate?");
+      Sparse[BackIdx] = I - begin();
     }
     // This depends on SmallVector::pop_back() not invalidating iterators.
     // std::vector::pop_back() doesn't give that guarantee.
@@ -253,7 +293,7 @@ public:
   /// @param   Key The key identifying the element to erase.
   /// @returns True when an element was erased, false if no element was found.
   ///
-  bool erase(unsigned Key) {
+  bool erase(const KeyT &Key) {
     iterator I = find(Key);
     if (I == end())
       return false;
diff --git a/include/llvm/CodeGen/ScheduleDAGInstrs.h b/include/llvm/CodeGen/ScheduleDAGInstrs.h
index 4fee108cd2..766c9b2f56 100644
--- a/include/llvm/CodeGen/ScheduleDAGInstrs.h
+++ b/include/llvm/CodeGen/ScheduleDAGInstrs.h
@@ -105,7 +105,7 @@ namespace llvm {
 
     VReg2SUnit(unsigned reg, SUnit *su): VirtReg(reg), SU(su) {}
 
-    unsigned getSparseSetKey() const {
+    unsigned getSparseSetIndex() const {
       return TargetRegisterInfo::virtReg2Index(VirtReg);
     }
   };
@@ -160,7 +160,7 @@ namespace llvm {
   /// compares ValueT's, only unsigned keys. This allows the set to be cleared
   /// between scheduling regions in constant time as long as ValueT does not
   /// require a destructor.
-  typedef SparseSet<VReg2SUnit> VReg2SUnitMap;
+  typedef SparseSet<VReg2SUnit, VirtReg2IndexFunctor> VReg2SUnitMap;
 
   /// ScheduleDAGInstrs - A ScheduleDAG subclass for scheduling lists of
   /// MachineInstrs.
@@ -321,10 +321,6 @@ namespace llvm {
     void addPhysRegDeps(SUnit *SU, unsigned OperIdx);
     void addVRegDefDeps(SUnit *SU, unsigned OperIdx);
     void addVRegUseDeps(SUnit *SU, unsigned OperIdx);
-
-    VReg2SUnitMap::iterator findVRegDef(unsigned VirtReg) {
-      return VRegDefs.find(TargetRegisterInfo::virtReg2Index(VirtReg));
-    }
   };
 
   /// newSUnit - Creates a new SUnit and return a ptr to it.
diff --git a/lib/CodeGen/RegAllocFast.cpp b/lib/CodeGen/RegAllocFast.cpp
index e09b7f8d26..70fbfd4c5c 100644
--- a/lib/CodeGen/RegAllocFast.cpp
+++ b/lib/CodeGen/RegAllocFast.cpp
@@ -77,7 +77,7 @@ namespace {
       explicit LiveReg(unsigned v)
         : LastUse(0), VirtReg(v), PhysReg(0), LastOpNum(0), Dirty(false) {}
 
-      unsigned getSparseSetKey() const {
+      unsigned getSparseSetIndex() const {
         return TargetRegisterInfo::virtReg2Index(VirtReg);
       }
     };
diff --git a/lib/CodeGen/ScheduleDAGInstrs.cpp b/lib/CodeGen/ScheduleDAGInstrs.cpp
index ae49b05858..20f8b78759 100644
--- a/lib/CodeGen/ScheduleDAGInstrs.cpp
+++ b/lib/CodeGen/ScheduleDAGInstrs.cpp
@@ -412,7 +412,7 @@ void ScheduleDAGInstrs::addVRegDefDeps(SUnit *SU, unsigned OperIdx) {
   // uses. We're conservative for now until we have a way to guarantee the uses
   // are not eliminated sometime during scheduling. The output dependence edge
   // is also useful if output latency exceeds def-use latency.
-  VReg2SUnitMap::iterator DefI = findVRegDef(Reg);
+  VReg2SUnitMap::iterator DefI = VRegDefs.find(Reg);
   if (DefI == VRegDefs.end())
     VRegDefs.insert(VReg2SUnit(Reg, SU));
   else {
@@ -464,7 +464,7 @@ void ScheduleDAGInstrs::addVRegUseDeps(SUnit *SU, unsigned OperIdx) {
   }
 
   // Add antidependence to the following def of the vreg it uses.
-  VReg2SUnitMap::iterator DefI = findVRegDef(Reg);
+  VReg2SUnitMap::iterator DefI = VRegDefs.find(Reg);
   if (DefI != VRegDefs.end() && DefI->SU != SU)
     DefI->SU->addPred(SDep(SU, SDep::Anti, 0, Reg));
 }
diff --git a/unittests/ADT/SparseSetTest.cpp b/unittests/ADT/SparseSetTest.cpp
index a6ea7572ce..eb0e0db283 100644
--- a/unittests/ADT/SparseSetTest.cpp
+++ b/unittests/ADT/SparseSetTest.cpp
@@ -161,7 +161,7 @@ TEST(SparseSetTest, MultipleEntrySet) {
 struct Alt {
   unsigned Value;
   explicit Alt(unsigned x) : Value(x) {}
-  unsigned getSparseSetKey() const { return Value - 1000; }
+  unsigned getSparseSetIndex() const { return Value - 1000; }
 };
 
 TEST(SparseSetTest, AltStructSet) {