3 files changed, 184 insertions, 55 deletions

diff --git a/llvm/include/llvm/ADT/STLExtras.h b/llvm/include/llvm/ADT/STLExtras.h
index 4435b46f462..6b18e3bee03 100644
--- a/llvm/include/llvm/ADT/STLExtras.h
+++ b/llvm/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.
+    decltype(&concat_iterator::incrementHelper<0>) 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.
+    decltype(&concat_iterator::getHelper<0>) GetHelperFns[] = {
+        &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>
 //===----------------------------------------------------------------------===//
diff --git a/llvm/include/llvm/IR/Module.h b/llvm/include/llvm/IR/Module.h
index 507a5ea4174..0f7a3ed47d4 100644
--- a/llvm/include/llvm/IR/Module.h
+++ b/llvm/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();
   }
 
   /// @}
diff --git a/llvm/unittests/ADT/STLExtrasTest.cpp b/llvm/unittests/ADT/STLExtrasTest.cpp
index 76881296db6..d3bef6a2e05 100644
--- a/llvm/unittests/ADT/STLExtrasTest.cpp
+++ b/llvm/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);
+}
 }