Python-like enumerate()

Summary:
Range-based for cannot be used if the element index is needed along
with the element. In these situations, it is often necessary to fall
back to the standard for loop, which is easy to get wrong, or maintain
an extra count variable, which is error-prone when control flow is
nontrivial (for example in the presence of `continue`).

This diff introduces a simple implementation of Python's
`enumerate()`, with the same signature. Since in C++ tuple-unpacking
is verbose, the iteration variable returned is a proxy object `it`
where the iteration index can be retrieved with `it.idx`, and the
value with `*it` or `it->...`, like a normal iterator.

Differential Revision: D3477877

fbshipit-source-id: 376af7f559e8b60f02a3f81f0c026a901e23ddcf

folly/Enumerate.h

+/*
+ * Copyright 2016 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <iterator>
+#include <memory>
+
+/*
+ * Similar to Python's enumerate(), folly::enumerate() can be used to
+ * iterate a range with a for-range loop, and it also allows to
+ * retrieve the count of iterations so far.
+ *
+ * For example:
+ *
+ * for (auto it : folly::enumerate(vec)) {
+ *   // *it is a reference to the current element. Const if vec is const.
+ *   // it->member can be used as well.
+ *   // it.index contains the iteration count.
+ * }
+ *
+ * If the iteration variable is const, the reference is too.
+ *
+ * for (const auto it : folly::enumerate(vec)) {
+ *   // *it is always a const reference.
+ * }
+ *
+ * @author Giuseppe Ottaviano <ott@fb.com>
+ */
+
+namespace folly {
+
+namespace detail {
+
+template <class T>
+struct MakeConst {
+  using type = const T;
+};
+template <class T>
+struct MakeConst<T&> {
+  using type = const T&;
+};
+template <class T>
+struct MakeConst<T*> {
+  using type = const T*;
+};
+
+template <class Iterator>
+class Enumerator {
+ public:
+  explicit Enumerator(Iterator it) : it_(std::move(it)) {}
+
+  class Proxy {
+   public:
+    using difference_type = ssize_t;
+    using value_type = typename std::iterator_traits<Iterator>::value_type;
+    using reference = typename std::iterator_traits<Iterator>::reference;
+    using pointer = typename std::iterator_traits<Iterator>::pointer;
+    using iterator_category = std::input_iterator_tag;
+
+    explicit Proxy(const Enumerator* e) : it_(e->it_), index(e->idx_) {}
+
+    // Non-const Proxy: Forward constness from Iterator.
+    reference operator*() {
+      return *it_;
+    }
+    pointer operator->() {
+      return std::addressof(**this);
+    }
+
+    // Const Proxy: Force const references.
+    typename MakeConst<reference>::type operator*() const {
+      return *it_;
+    }
+    typename MakeConst<pointer>::type operator->() const {
+      return std::addressof(**this);
+    }
+
+   private:
+    const Iterator& it_;
+
+   public:
+    const size_t index;
+  };
+
+  Proxy operator*() const {
+    return Proxy(this);
+  }
+
+  Enumerator& operator++() {
+    ++it_;
+    ++idx_;
+    return *this;
+  }
+
+  bool operator==(const Enumerator& rhs) {
+    return it_ == rhs.it_;
+  }
+
+  bool operator!=(const Enumerator& rhs) {
+    return !(*this == rhs);
+  }
+
+ private:
+  Iterator it_;
+  size_t idx_ = 0;
+};
+
+template <class Range>
+class RangeEnumerator {
+  Range r_;
+  using Iterator = decltype(r_.begin());
+
+ public:
+  explicit RangeEnumerator(Range&& r) : r_(std::forward<Range>(r)) {}
+
+  Enumerator<Iterator> begin() {
+    return Enumerator<Iterator>(r_.begin());
+  }
+  Enumerator<Iterator> end() {
+    return Enumerator<Iterator>(r_.end());
+  }
+};
+
+} // namespace detail
+
+template <class Range>
+detail::RangeEnumerator<Range> enumerate(Range&& r) {
+  return detail::RangeEnumerator<Range>(std::forward<Range>(r));
+}
+
+} // namespace folly

folly/Makefile.am

@@ -82,6 +82,7 @@ nobase_follyinclude_HEADERS = \
 	DynamicConverter.h \
 	dynamic.h \
 	dynamic-inl.h \
+	Enumerate.h \
 	Exception.h \
 	ExceptionString.h \
 	ExceptionWrapper.h \

folly/test/EnumerateTest.cpp

+/*
+ * Copyright 2016 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <string>
+#include <vector>
+
+#include <folly/Enumerate.h>
+#include <folly/Range.h>
+#include <gtest/gtest.h>
+
+TEST(Enumerate, Basic) {
+  std::vector<std::string> v = {"abc", "a", "ab"};
+  size_t i = 0;
+  for (auto it : folly::enumerate(v)) {
+    EXPECT_EQ(it.index, i);
+    EXPECT_EQ(*it, v[i]);
+    EXPECT_EQ(it->size(), v[i].size());
+
+    // Test mutability.
+    std::string newValue = "x";
+    *it = newValue;
+    EXPECT_EQ(newValue, v[i]);
+
+    ++i;
+  }
+
+  EXPECT_EQ(i, v.size());
+}
+
+TEST(Enumerate, Temporary) {
+  std::vector<std::string> v = {"abc", "a", "ab"};
+  size_t i = 0;
+  for (auto it : folly::enumerate(decltype(v)(v))) { // Copy v.
+    EXPECT_EQ(it.index, i);
+    EXPECT_EQ(*it, v[i]);
+    EXPECT_EQ(it->size(), v[i].size());
+    ++i;
+  }
+
+  EXPECT_EQ(i, v.size());
+};
+
+template <class T>
+struct IsConstReference {
+  constexpr static bool value = false;
+};
+template <class T>
+struct IsConstReference<const T&> {
+  constexpr static bool value = true;
+};
+
+TEST(Enumerate, BasicConstArg) {
+  const std::vector<std::string> v = {"abc", "a", "ab"};
+  size_t i = 0;
+  for (auto it : folly::enumerate(v)) {
+    static_assert(
+        IsConstReference<decltype(*it)>::value, "Enumerating a const vector");
+    EXPECT_EQ(it.index, i);
+    EXPECT_EQ(*it, v[i]);
+    EXPECT_EQ(it->size(), v[i].size());
+    ++i;
+  }
+
+  EXPECT_EQ(i, v.size());
+}
+
+TEST(Enumerate, BasicConstEnumerate) {
+  std::vector<std::string> v = {"abc", "a", "ab"};
+  size_t i = 0;
+  for (const auto it : folly::enumerate(v)) {
+    static_assert(IsConstReference<decltype(*it)>::value, "Const enumeration");
+    EXPECT_EQ(it.index, i);
+    EXPECT_EQ(*it, v[i]);
+    EXPECT_EQ(it->size(), v[i].size());
+    ++i;
+  }
+
+  EXPECT_EQ(i, v.size());
+}
+
+TEST(Enumerate, TemporaryConstEnumerate) {
+  std::vector<std::string> v = {"abc", "a", "ab"};
+  size_t i = 0;
+  for (const auto it : folly::enumerate(decltype(v)(v))) { // Copy v.
+    static_assert(IsConstReference<decltype(*it)>::value, "Const enumeration");
+    EXPECT_EQ(it.index, i);
+    EXPECT_EQ(*it, v[i]);
+    EXPECT_EQ(it->size(), v[i].size());
+    ++i;
+  }
+
+  EXPECT_EQ(i, v.size());
+}
+
+TEST(Enumerate, RangeSupport) {
+  std::vector<std::string> v = {"abc", "a", "ab"};
+  size_t i = 0;
+  for (const auto it : folly::enumerate(folly::range(v))) {
+    EXPECT_EQ(it.index, i);
+    EXPECT_EQ(*it, v[i]);
+    EXPECT_EQ(it->size(), v[i].size());
+    ++i;
+  }
+
+  EXPECT_EQ(i, v.size());
+}
+
+TEST(Enumerate, EmptyRange) {
+  std::vector<std::string> v;
+  for (auto it : folly::enumerate(v)) {
+    (void)it; // Silence warnings.
+    EXPECT_TRUE(false);
+  }
+}