[IR] Redesign the case iterator in SwitchInst to actually be an iterator

and to expose a handle to represent the actual case rather than having
the iterator return a reference to itself.

All of this allows the iterator to be used with common STL facilities,
standard algorithms, etc.

Doing this exposed some missing facilities in the iterator facade that
I've fixed and required some work to the actual iterator to fully
support the necessary API.

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

llvm-svn: 300032

3 files changed, 125 insertions, 75 deletions

diff --git a/llvm/include/llvm/ADT/iterator.h b/llvm/include/llvm/ADT/iterator.h
index 8f162da1b54..28dcdf9613e 100644
--- a/llvm/include/llvm/ADT/iterator.h
+++ b/llvm/include/llvm/ADT/iterator.h
@@ -165,9 +165,15 @@ public:
     return !static_cast<const DerivedT *>(this)->operator<(RHS);
   }
 
+  PointerT operator->() { return &static_cast<DerivedT *>(this)->operator*(); }
   PointerT operator->() const {
     return &static_cast<const DerivedT *>(this)->operator*();
   }
+  ReferenceProxy operator[](DifferenceTypeT n) {
+    static_assert(IsRandomAccess,
+                  "Subscripting is only defined for random access iterators.");
+    return ReferenceProxy(static_cast<DerivedT *>(this)->operator+(n));
+  }
   ReferenceProxy operator[](DifferenceTypeT n) const {
     static_assert(IsRandomAccess,
                   "Subscripting is only defined for random access iterators.");
diff --git a/llvm/include/llvm/Analysis/CFGPrinter.h b/llvm/include/llvm/Analysis/CFGPrinter.h
index efaa9d6df8e..5786769cc50 100644
--- a/llvm/include/llvm/Analysis/CFGPrinter.h
+++ b/llvm/include/llvm/Analysis/CFGPrinter.h
@@ -140,8 +140,7 @@ struct DOTGraphTraits<const Function*> : public DefaultDOTGraphTraits {
 
       std::string Str;
       raw_string_ostream OS(Str);
-      SwitchInst::ConstCaseIt Case =
-          SwitchInst::ConstCaseIt::fromSuccessorIndex(SI, SuccNo);
+      auto Case = *SwitchInst::ConstCaseIt::fromSuccessorIndex(SI, SuccNo);
       OS << Case.getCaseValue()->getValue();
       return OS.str();
     }
diff --git a/llvm/include/llvm/IR/Instructions.h b/llvm/include/llvm/IR/Instructions.h
index 2a55048c133..dbbe6eed114 100644
--- a/llvm/include/llvm/IR/Instructions.h
+++ b/llvm/include/llvm/IR/Instructions.h
@@ -3096,49 +3096,39 @@ public:
   // -2
   static const unsigned DefaultPseudoIndex = static_cast<unsigned>(~0L-1);
 
-  template <class SwitchInstT, class ConstantIntT, class BasicBlockT>
-  class CaseIteratorT
-      : public iterator_facade_base<
-            CaseIteratorT<SwitchInstT, ConstantIntT, BasicBlockT>,
-            std::random_access_iterator_tag,
-            CaseIteratorT<SwitchInstT, ConstantIntT, BasicBlockT>> {
+  template <typename CaseHandleT> class CaseIteratorT;
+
+  /// A handle to a particular switch case. It exposes a convenient interface
+  /// to both the case value and the successor block.
+  ///
+  /// We define this as a template and instantiate it to form both a const and
+  /// non-const handle.
+  template <typename SwitchInstT, typename ConstantIntT, typename BasicBlockT>
+  class CaseHandleT {
+    // Directly befriend both const and non-const iterators.
+    friend class SwitchInst::CaseIteratorT<
+        CaseHandleT<SwitchInstT, ConstantIntT, BasicBlockT>>;
+
   protected:
+    // Expose the switch type we're parameterized with to the iterator.
+    typedef SwitchInstT SwitchInstType;
+
     SwitchInstT *SI;
     ptrdiff_t Index;
 
-  public:
-    typedef CaseIteratorT<SwitchInstT, ConstantIntT, BasicBlockT> Self;
-
-    /// Default constructed iterator is in an invalid state until assigned to
-    /// a case for a particular switch.
-    CaseIteratorT() : SI(nullptr) {}
-
-    /// Initializes case iterator for given SwitchInst and for given
-    /// case number.
-    CaseIteratorT(SwitchInstT *SI, unsigned CaseNum) {
-      this->SI = SI;
-      Index = CaseNum;
-    }
-
-    /// Initializes case iterator for given SwitchInst and for given
-    /// TerminatorInst's successor index.
-    static Self fromSuccessorIndex(SwitchInstT *SI, unsigned SuccessorIndex) {
-      assert(SuccessorIndex < SI->getNumSuccessors() &&
-             "Successor index # out of range!");
-      return SuccessorIndex != 0 ?
-             Self(SI, SuccessorIndex - 1) :
-             Self(SI, DefaultPseudoIndex);
-    }
+    CaseHandleT() = default;
+    CaseHandleT(SwitchInstT *SI, ptrdiff_t Index) : SI(SI), Index(Index) {}
 
+  public:
     /// Resolves case value for current case.
-    ConstantIntT *getCaseValue() {
+    ConstantIntT *getCaseValue() const {
       assert((unsigned)Index < SI->getNumCases() &&
              "Index out the number of cases.");
       return reinterpret_cast<ConstantIntT *>(SI->getOperand(2 + Index * 2));
     }
 
     /// Resolves successor for current case.
-    BasicBlockT *getCaseSuccessor() {
+    BasicBlockT *getCaseSuccessor() const {
       assert(((unsigned)Index < SI->getNumCases() ||
               (unsigned)Index == DefaultPseudoIndex) &&
              "Index out the number of cases.");
@@ -3156,43 +3146,20 @@ public:
       return (unsigned)Index != DefaultPseudoIndex ? Index + 1 : 0;
     }
 
-    Self &operator+=(ptrdiff_t N) {
-      // Check index correctness after addition.
-      // Note: Index == getNumCases() means end().
-      assert(Index + N >= 0 && (unsigned)(Index + N) <= SI->getNumCases() &&
-             "Index out the number of cases.");
-      Index += N;
-      return *this;
-    }
-    Self &operator-=(ptrdiff_t N) {
-      // Check index correctness after subtraction.
-      // Note: Index == getNumCases() means end().
-      assert(Index - N >= 0 && (unsigned)(Index - N) <= SI->getNumCases() &&
-             "Index out the number of cases.");
-      Index -= N;
-      return *this;
-    }
-    bool operator==(const Self& RHS) const {
+    bool operator==(const CaseHandleT &RHS) const {
       assert(SI == RHS.SI && "Incompatible operators.");
       return Index == RHS.Index;
     }
-    bool operator<(const Self& RHS) const {
-      assert(SI == RHS.SI && "Incompatible operators.");
-      return Index < RHS.Index;
-    }
-    Self &operator*() { return *this; }
-    const Self &operator*() const { return *this; }
   };
 
-  typedef CaseIteratorT<const SwitchInst, const ConstantInt, const BasicBlock>
-    ConstCaseIt;
+  typedef CaseHandleT<const SwitchInst, const ConstantInt, const BasicBlock>
+      ConstCaseHandle;
 
-  class CaseIt : public CaseIteratorT<SwitchInst, ConstantInt, BasicBlock> {
-    typedef CaseIteratorT<SwitchInst, ConstantInt, BasicBlock> ParentTy;
+  class CaseHandle : public CaseHandleT<SwitchInst, ConstantInt, BasicBlock> {
+    friend class SwitchInst::CaseIteratorT<CaseHandle>;
 
   public:
-    CaseIt(const ParentTy &Src) : ParentTy(Src) {}
-    CaseIt(SwitchInst *SI, unsigned CaseNum) : ParentTy(SI, CaseNum) {}
+    CaseHandle(SwitchInst *SI, ptrdiff_t Index) : CaseHandleT(SI, Index) {}
 
     /// Sets the new value for current case.
     void setValue(ConstantInt *V) {
@@ -3207,6 +3174,74 @@ public:
     }
   };
 
+  template <typename CaseHandleT>
+  class CaseIteratorT
+      : public iterator_facade_base<CaseIteratorT<CaseHandleT>,
+                                    std::random_access_iterator_tag,
+                                    CaseHandleT> {
+    typedef typename CaseHandleT::SwitchInstType SwitchInstT;
+
+    CaseHandleT Case;
+
+  public:
+    /// Default constructed iterator is in an invalid state until assigned to
+    /// a case for a particular switch.
+    CaseIteratorT() = default;
+
+    /// Initializes case iterator for given SwitchInst and for given
+    /// case number.
+    CaseIteratorT(SwitchInstT *SI, unsigned CaseNum) : Case(SI, CaseNum) {}
+
+    /// Initializes case iterator for given SwitchInst and for given
+    /// TerminatorInst's successor index.
+    static CaseIteratorT fromSuccessorIndex(SwitchInstT *SI,
+                                            unsigned SuccessorIndex) {
+      assert(SuccessorIndex < SI->getNumSuccessors() &&
+             "Successor index # out of range!");
+      return SuccessorIndex != 0 ? CaseIteratorT(SI, SuccessorIndex - 1)
+                                 : CaseIteratorT(SI, DefaultPseudoIndex);
+    }
+
+    /// Support converting to the const variant. This will be a no-op for const
+    /// variant.
+    operator CaseIteratorT<ConstCaseHandle>() const {
+      return CaseIteratorT<ConstCaseHandle>(Case.SI, Case.Index);
+    }
+
+    CaseIteratorT &operator+=(ptrdiff_t N) {
+      // Check index correctness after addition.
+      // Note: Index == getNumCases() means end().
+      assert(Case.Index + N >= 0 &&
+             (unsigned)(Case.Index + N) <= Case.SI->getNumCases() &&
+             "Case.Index out the number of cases.");
+      Case.Index += N;
+      return *this;
+    }
+    CaseIteratorT &operator-=(ptrdiff_t N) {
+      // Check index correctness after subtraction.
+      // Note: Case.Index == getNumCases() means end().
+      assert(Case.Index - N >= 0 &&
+             (unsigned)(Case.Index - N) <= Case.SI->getNumCases() &&
+             "Case.Index out the number of cases.");
+      Case.Index -= N;
+      return *this;
+    }
+    ptrdiff_t operator-(const CaseIteratorT &RHS) const {
+      assert(Case.SI == RHS.Case.SI && "Incompatible operators.");
+      return Case.Index - RHS.Case.Index;
+    }
+    bool operator==(const CaseIteratorT &RHS) const { return Case == RHS.Case; }
+    bool operator<(const CaseIteratorT &RHS) const {
+      assert(Case.SI == RHS.Case.SI && "Incompatible operators.");
+      return Case.Index < RHS.Case.Index;
+    }
+    CaseHandleT &operator*() { return Case; }
+    const CaseHandleT &operator*() const { return Case; }
+  };
+
+  typedef CaseIteratorT<CaseHandle> CaseIt;
+  typedef CaseIteratorT<ConstCaseHandle> ConstCaseIt;
+
   static SwitchInst *Create(Value *Value, BasicBlock *Default,
                             unsigned NumCases,
                             Instruction *InsertBefore = nullptr) {
@@ -3290,30 +3325,40 @@ public:
   /// default case iterator to indicate that it is handled by the default
   /// handler.
   CaseIt findCaseValue(const ConstantInt *C) {
-    for (CaseIt i = case_begin(), e = case_end(); i != e; ++i)
-      if (i.getCaseValue() == C)
-        return i;
+    CaseIt I = llvm::find_if(
+        cases(), [C](CaseHandle &Case) { return Case.getCaseValue() == C; });
+    if (I != case_end())
+      return I;
+
     return case_default();
   }
   ConstCaseIt findCaseValue(const ConstantInt *C) const {
-    for (ConstCaseIt i = case_begin(), e = case_end(); i != e; ++i)
-      if (i.getCaseValue() == C)
-        return i;
+    ConstCaseIt I = llvm::find_if(cases(), [C](ConstCaseHandle &Case) {
+      return Case.getCaseValue() == C;
+    });
+    if (I != case_end())
+      return I;
+
     return case_default();
   }
 
   /// Finds the unique case value for a given successor. Returns null if the
   /// successor is not found, not unique, or is the default case.
   ConstantInt *findCaseDest(BasicBlock *BB) {
-    if (BB == getDefaultDest()) return nullptr;
+    if (BB == getDefaultDest())
+      return nullptr;
 
     ConstantInt *CI = nullptr;
-    for (CaseIt i = case_begin(), e = case_end(); i != e; ++i) {
-      if (i.getCaseSuccessor() == BB) {
-        if (CI) return nullptr;   // Multiple cases lead to BB.
-        else CI = i.getCaseValue();
-      }
+    for (auto Case : cases()) {
+      if (Case.getCaseSuccessor() != BB)
+        continue;
+
+      if (CI)
+        return nullptr; // Multiple cases lead to BB.
+
+      CI = Case.getCaseValue();
     }
+
     return CI;
   }
 
@@ -3330,7 +3375,7 @@ public:
   /// This action invalidates iterators for all cases following the one removed,
   /// including the case_end() iterator. It returns an iterator for the next
   /// case.
-  CaseIt removeCase(CaseIt i);
+  CaseIt removeCase(CaseIt I);
 
   unsigned getNumSuccessors() const { return getNumOperands()/2; }
   BasicBlock *getSuccessor(unsigned idx) const {