[ADT] Add a generic concatenating iterator and range (take 2).

This recommits r290512 that was reverted when MSVC failed to compile it. Since
then I've played with various approaches using rextester.com (where I was able
to reproduce the failure) and think that I have a solution thanks in part to
the help of Dave Blaikie! It seems MSVC just has a defective `decltype` in this
version. Manually writing out the type seems to do the trick, even though it is
.... quite complicated.

Original commit message:
This allows both defining convenience iterator/range accessors on types
which walk across N different independent ranges within the object, and
more direct and simple usages with range based for loops such as shown
in the unittest. The same facilities are used for both. They end up
quite small and simple as it happens.

I've also switched an iterator on `Module` to use this. I would like to
add another convenience iterator that includes even more sequences as
part of it and seeing this one already present motivated me to actually
abstract it away and introduce a general utility.

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

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

include/llvm/ADT/STLExtras.h

@@ -31,6 +31,7 @@
 #include "llvm/ADT/iterator.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/Compiler.h"
+#include "llvm/Support/ErrorHandling.h"
 
 namespace llvm {
 
@@ -435,6 +436,151 @@ detail::zippy<detail::zip_first, T, U, Args...> zip_first(T &&t, U &&u,
       std::forward<T>(t), std::forward<U>(u), std::forward<Args>(args)...);
 }
 
+/// Iterator wrapper that concatenates sequences together.
+///
+/// This can concatenate different iterators, even with different types, into
+/// a single iterator provided the value types of all the concatenated
+/// iterators expose `reference` and `pointer` types that can be converted to
+/// `ValueT &` and `ValueT *` respectively. It doesn't support more
+/// interesting/customized pointer or reference types.
+///
+/// Currently this only supports forward or higher iterator categories as
+/// inputs and always exposes a forward iterator interface.
+template <typename ValueT, typename... IterTs>
+class concat_iterator
+    : public iterator_facade_base<concat_iterator<ValueT, IterTs...>,
+                                  std::forward_iterator_tag, ValueT> {
+  typedef typename concat_iterator::iterator_facade_base BaseT;
+
+  /// We store both the current and end iterators for each concatenated
+  /// sequence in a tuple of pairs.
+  ///
+  /// Note that something like iterator_range seems nice at first here, but the
+  /// range properties are of little benefit and end up getting in the way
+  /// because we need to do mutation on the current iterators.
+  std::tuple<std::pair<IterTs, IterTs>...> IterPairs;
+
+  /// Attempts to increment a specific iterator.
+  ///
+  /// Returns true if it was able to increment the iterator. Returns false if
+  /// the iterator is already at the end iterator.
+  template <size_t Index> bool incrementHelper() {
+    auto &IterPair = std::get<Index>(IterPairs);
+    if (IterPair.first == IterPair.second)
+      return false;
+
+    ++IterPair.first;
+    return true;
+  }
+
+  /// Increments the first non-end iterator.
+  ///
+  /// It is an error to call this with all iterators at the end.
+  template <size_t... Ns> void increment(index_sequence<Ns...>) {
+    // Build a sequence of functions to increment each iterator if possible.
+    bool (concat_iterator::*IncrementHelperFns[])() = {
+        &concat_iterator::incrementHelper<Ns>...};
+
+    // Loop over them, and stop as soon as we succeed at incrementing one.
+    for (auto &IncrementHelperFn : IncrementHelperFns)
+      if ((this->*IncrementHelperFn)())
+        return;
+
+    llvm_unreachable("Attempted to increment an end concat iterator!");
+  }
+
+  /// Returns null if the specified iterator is at the end. Otherwise,
+  /// dereferences the iterator and returns the address of the resulting
+  /// reference.
+  template <size_t Index> ValueT *getHelper() const {
+    auto &IterPair = std::get<Index>(IterPairs);
+    if (IterPair.first == IterPair.second)
+      return nullptr;
+
+    return &*IterPair.first;
+  }
+
+  /// Finds the first non-end iterator, dereferences, and returns the resulting
+  /// reference.
+  ///
+  /// It is an error to call this with all iterators at the end.
+  template <size_t... Ns> ValueT &get(index_sequence<Ns...>) const {
+    // Build a sequence of functions to get from iterator if possible.
+    ValueT *(concat_iterator::*GetHelperFns[])() const = {
+        &concat_iterator::getHelper<Ns>...};
+
+    // Loop over them, and return the first result we find.
+    for (auto &GetHelperFn : GetHelperFns)
+      if (ValueT *P = (this->*GetHelperFn)())
+        return *P;
+
+    llvm_unreachable("Attempted to get a pointer from an end concat iterator!");
+  }
+
+public:
+  /// Constructs an iterator from a squence of ranges.
+  ///
+  /// We need the full range to know how to switch between each of the
+  /// iterators.
+  template <typename... RangeTs>
+  explicit concat_iterator(RangeTs &&... Ranges)
+      : IterPairs({std::begin(Ranges), std::end(Ranges)}...) {}
+
+  using BaseT::operator++;
+  concat_iterator &operator++() {
+    increment(index_sequence_for<IterTs...>());
+    return *this;
+  }
+
+  ValueT &operator*() const { return get(index_sequence_for<IterTs...>()); }
+
+  bool operator==(const concat_iterator &RHS) const {
+    return IterPairs == RHS.IterPairs;
+  }
+};
+
+namespace detail {
+/// Helper to store a sequence of ranges being concatenated and access them.
+///
+/// This is designed to facilitate providing actual storage when temporaries
+/// are passed into the constructor such that we can use it as part of range
+/// based for loops.
+template <typename ValueT, typename... RangeTs> class concat_range {
+public:
+  typedef concat_iterator<ValueT,
+                          decltype(std::begin(std::declval<RangeTs &>()))...>
+      iterator;
+
+private:
+  std::tuple<RangeTs...> Ranges;
+
+  template <size_t... Ns> iterator begin_impl(index_sequence<Ns...>) {
+    return iterator(std::get<Ns>(Ranges)...);
+  }
+  template <size_t... Ns> iterator end_impl(index_sequence<Ns...>) {
+    return iterator(make_range(std::end(std::get<Ns>(Ranges)),
+                               std::end(std::get<Ns>(Ranges)))...);
+  }
+
+public:
+  iterator begin() { return begin_impl(index_sequence_for<RangeTs...>{}); }
+  iterator end() { return end_impl(index_sequence_for<RangeTs...>{}); }
+  concat_range(RangeTs &&... Ranges)
+      : Ranges(std::forward<RangeTs>(Ranges)...) {}
+};
+}
+
+/// Concatenated range across two or more ranges.
+///
+/// The desired value type must be explicitly specified.
+template <typename ValueT, typename... RangeTs>
+detail::concat_range<ValueT, RangeTs...> concat(RangeTs &&... Ranges) {
+  static_assert(sizeof...(RangeTs) > 1,
+                "Need more than one range to concatenate!");
+  return detail::concat_range<ValueT, RangeTs...>(
+      std::forward<RangeTs>(Ranges)...);
+}
+
 //===----------------------------------------------------------------------===//
 //     Extra additions to <utility>
 //===----------------------------------------------------------------------===//

include/llvm/IR/Module.h

@@ -590,69 +590,29 @@ public:
 /// @name Convenience iterators
 /// @{
 
-  template <bool IsConst> class global_object_iterator_t {
-    friend Module;
-
-    typename std::conditional<IsConst, const_iterator, iterator>::type
-        function_i,
-        function_e;
-    typename std::conditional<IsConst, const_global_iterator,
-                              global_iterator>::type global_i;
-
-    typedef
-        typename std::conditional<IsConst, const Module, Module>::type ModuleTy;
-
-    global_object_iterator_t(ModuleTy &M)
-        : function_i(M.begin()), function_e(M.end()),
-          global_i(M.global_begin()) {}
-    global_object_iterator_t(ModuleTy &M, int)
-        : function_i(M.end()), function_e(M.end()), global_i(M.global_end()) {}
-
-  public:
-    global_object_iterator_t &operator++() {
-      if (function_i != function_e)
-        ++function_i;
-      else
-        ++global_i;
-      return *this;
-    }
-
-    typename std::conditional<IsConst, const GlobalObject, GlobalObject>::type &
-    operator*() const {
-      if (function_i != function_e)
-        return *function_i;
-      else
-        return *global_i;
-    }
-
-    bool operator!=(const global_object_iterator_t &other) const {
-      return function_i != other.function_i || global_i != other.global_i;
-    }
-  };
-
-  typedef global_object_iterator_t</*IsConst=*/false> global_object_iterator;
-  typedef global_object_iterator_t</*IsConst=*/true>
+  typedef concat_iterator<GlobalObject, iterator, global_iterator>
+      global_object_iterator;
+  typedef concat_iterator<const GlobalObject, const_iterator,
+                          const_global_iterator>
       const_global_object_iterator;
 
-  global_object_iterator global_object_begin() {
-    return global_object_iterator(*this);
+  iterator_range<global_object_iterator> global_objects() {
+    return concat<GlobalObject>(functions(), globals());
   }
-  global_object_iterator global_object_end() {
-    return global_object_iterator(*this, 0);
+  iterator_range<const_global_object_iterator> global_objects() const {
+    return concat<const GlobalObject>(functions(), globals());
   }
 
-  const_global_object_iterator global_object_begin() const {
-    return const_global_object_iterator(*this);
-  }
-  const_global_object_iterator global_object_end() const {
-    return const_global_object_iterator(*this, 0);
+  global_object_iterator global_object_begin() {
+    return global_objects().begin();
   }
+  global_object_iterator global_object_end() { return global_objects().end(); }
 
-  iterator_range<global_object_iterator> global_objects() {
-    return make_range(global_object_begin(), global_object_end());
+  const_global_object_iterator global_object_begin() const {
+    return global_objects().begin();
   }
-  iterator_range<const_global_object_iterator> global_objects() const {
-    return make_range(global_object_begin(), global_object_end());
+  const_global_object_iterator global_object_end() const {
+    return global_objects().end();
   }
 
   /// @}

unittests/ADT/STLExtrasTest.cpp

@@ -10,6 +10,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "gtest/gtest.h"
 
+#include <list>
 #include <vector>
 
 using namespace llvm;
@@ -253,4 +254,26 @@ TEST(STLExtrasTest, CountAdaptor) {
   EXPECT_EQ(1, count(v, 3));
   EXPECT_EQ(1, count(v, 4));
 }
+
+TEST(STLExtrasTest, ConcatRange) {
+  std::vector<int> Expected = {1, 2, 3, 4, 5, 6, 7, 8};
+  std::vector<int> Test;
+
+  std::vector<int> V1234 = {1, 2, 3, 4};
+  std::list<int> L56 = {5, 6};
+  SmallVector<int, 2> SV78 = {7, 8};
+
+  // Use concat across different sized ranges of different types with different
+  // iterators.
+  for (int &i : concat<int>(V1234, L56, SV78))
+    Test.push_back(i);
+  EXPECT_EQ(Expected, Test);
+
+  // Use concat between a temporary, an L-value, and an R-value to make sure
+  // complex lifetimes work well.
+  Test.clear();
+  for (int &i : concat<int>(std::vector<int>(V1234), L56, std::move(SV78)))
+    Test.push_back(i);
+  EXPECT_EQ(Expected, Test);
+}
 }