From ab82a456ef39cb25d71501f019738ae7ed9e59a7 Mon Sep 17 00:00:00 2001
From: Zachary Turner <zturner@google.com>
Date: Thu, 29 Sep 2016 22:59:30 +0000
Subject: [PATCH] Add llvm::enumerate() to STLExtras.

enumerate allows you to iterate over a range by pairing the
iterator's value with its index in the enumeration.  This gives
you most of the benefits of using a for loop while still allowing
the range syntax.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@282804 91177308-0d34-0410-b5e6-96231b3b80d8
---
 include/llvm/ADT/STLExtras.h    | 67 +++++++++++++++++++++++++++++++++
 unittests/ADT/STLExtrasTest.cpp | 49 ++++++++++++++++++++++++
 2 files changed, 116 insertions(+)

diff --git a/include/llvm/ADT/STLExtras.h b/include/llvm/ADT/STLExtras.h
index e6215e4ae5b..b557c7743a9 100644
--- a/include/llvm/ADT/STLExtras.h
+++ b/include/llvm/ADT/STLExtras.h
@@ -626,6 +626,73 @@ template <typename T> struct deref {
   }
 };
 
+namespace detail {
+template <typename I, typename V> class enumerator_impl {
+public:
+  template <typename V> struct result_pair {
+    result_pair(std::size_t Index, V Value) : Index(Index), Value(Value) {}
+
+    const std::size_t Index;
+    V Value;
+  };
+
+  template <typename I, typename V> struct iterator {
+    iterator(I Iter, std::size_t Index) : Iter(Iter), Index(Index) {}
+
+    result_pair<const V> operator*() const {
+      return result_pair<const V>(Index, *Iter);
+    }
+    result_pair<V> operator*() { return result_pair<V>(Index, *Iter); }
+
+    iterator &operator++() {
+      ++Iter;
+      ++Index;
+      return *this;
+    }
+
+    bool operator!=(const iterator &RHS) const { return Iter != RHS.Iter; }
+
+  private:
+    I Iter;
+    std::size_t Index;
+  };
+
+  enumerator_impl(I Begin, I End)
+      : Begin(std::move(Begin)), End(std::move(End)) {}
+
+  iterator<I, V> begin() { return iterator<I, V>(Begin, 0); }
+  iterator<I, V> end() { return iterator<I, V>(End, std::size_t(-1)); }
+
+  iterator<I, V> begin() const { return iterator<I, V>(Begin, 0); }
+  iterator<I, V> end() const { return iterator<I, V>(End, std::size_t(-1)); }
+
+private:
+  I Begin;
+  I End;
+};
+}
+
+/// Given an input range, returns a new range whose values are are pair (A,B)
+/// such that A is the 0-based index of the item in the sequence, and B is
+/// the value from the original sequence.  Example:
+///
+/// std::vector<char> Items = {'A', 'B', 'C', 'D'};
+/// for (auto X : enumerate(Items)) {
+///   printf("Item %d - %c\n", X.Item, X.Value);
+/// }
+///
+/// Output:
+///   Item 0 - A
+///   Item 1 - B
+///   Item 2 - C
+///   Item 3 - D
+///
+template <typename R> auto enumerate(R &&Range) {
+  typedef decltype(std::begin(Range)) I;
+  typedef decltype(*std::begin(Range)) V;
+  return detail::enumerator_impl<I, V>(std::begin(Range), std::end(Range));
+}
+
 } // End llvm namespace
 
 #endif
diff --git a/unittests/ADT/STLExtrasTest.cpp b/unittests/ADT/STLExtrasTest.cpp
index dc62b03741c..790cf1c9b61 100644
--- a/unittests/ADT/STLExtrasTest.cpp
+++ b/unittests/ADT/STLExtrasTest.cpp
@@ -10,6 +10,8 @@
 #include "llvm/ADT/STLExtras.h"
 #include "gtest/gtest.h"
 
+#include <vector>
+
 using namespace llvm;
 
 namespace {
@@ -37,4 +39,51 @@ TEST(STLExtrasTest, Rank) {
   EXPECT_EQ(4, f(rank<6>()));
 }
 
+TEST(STLExtrasTest, Enumerate) {
+  std::vector<char> foo = {'a', 'b', 'c'};
+
+  std::vector<std::pair<std::size_t, char>> results;
+
+  for (auto X : llvm::enumerate(foo)) {
+    results.push_back(std::make_pair(X.Index, X.Value));
+  }
+  ASSERT_EQ(3, results.size());
+  EXPECT_EQ(0, results[0].first);
+  EXPECT_EQ('a', results[0].second);
+  EXPECT_EQ(1, results[1].first);
+  EXPECT_EQ('b', results[1].second);
+  EXPECT_EQ(2, results[2].first);
+  EXPECT_EQ('c', results[2].second);
+
+  results.clear();
+  const std::vector<int> bar = {'1', '2', '3'};
+  for (auto X : llvm::enumerate(bar)) {
+    results.push_back(std::make_pair(X.Index, X.Value));
+  }
+  EXPECT_EQ(0, results[0].first);
+  EXPECT_EQ('1', results[0].second);
+  EXPECT_EQ(1, results[1].first);
+  EXPECT_EQ('2', results[1].second);
+  EXPECT_EQ(2, results[2].first);
+  EXPECT_EQ('3', results[2].second);
+
+  results.clear();
+  const std::vector<int> baz;
+  for (auto X : llvm::enumerate(baz)) {
+    results.push_back(std::make_pair(X.Index, X.Value));
+  }
+  EXPECT_TRUE(baz.empty());
+}
+
+TEST(STLExtrasTest, EnumerateModify) {
+  std::vector<char> foo = {'a', 'b', 'c'};
+
+  for (auto X : llvm::enumerate(foo)) {
+    ++X.Value;
+  }
+
+  EXPECT_EQ('b', foo[0]);
+  EXPECT_EQ('c', foo[1]);
+  EXPECT_EQ('d', foo[2]);
+}
 }
-- 
2.50.1