From d2549d035ac7b7f333f1e349ba20f3ebe8bd4609 Mon Sep 17 00:00:00 2001 From: Giuseppe Ottaviano Date: Wed, 6 Jul 2016 16:24:04 -0700 Subject: [PATCH] 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 | 145 +++++++++++++++++++++++++++++++++++ folly/Makefile.am | 1 + folly/test/EnumerateTest.cpp | 127 ++++++++++++++++++++++++++++++ 3 files changed, 273 insertions(+) create mode 100644 folly/Enumerate.h create mode 100644 folly/test/EnumerateTest.cpp diff --git a/folly/Enumerate.h b/folly/Enumerate.h new file mode 100644 index 00000000..7d4ae4d0 --- /dev/null +++ b/folly/Enumerate.h @@ -0,0 +1,145 @@ +/* + * 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 +#include + +/* + * 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 + */ + +namespace folly { + +namespace detail { + +template +struct MakeConst { + using type = const T; +}; +template +struct MakeConst { + using type = const T&; +}; +template +struct MakeConst { + using type = const T*; +}; + +template +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::value_type; + using reference = typename std::iterator_traits::reference; + using pointer = typename std::iterator_traits::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::type operator*() const { + return *it_; + } + typename MakeConst::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 RangeEnumerator { + Range r_; + using Iterator = decltype(r_.begin()); + + public: + explicit RangeEnumerator(Range&& r) : r_(std::forward(r)) {} + + Enumerator begin() { + return Enumerator(r_.begin()); + } + Enumerator end() { + return Enumerator(r_.end()); + } +}; + +} // namespace detail + +template +detail::RangeEnumerator enumerate(Range&& r) { + return detail::RangeEnumerator(std::forward(r)); +} + +} // namespace folly diff --git a/folly/Makefile.am b/folly/Makefile.am index 9df86b5b..381bd1bb 100644 --- a/folly/Makefile.am +++ b/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 \ diff --git a/folly/test/EnumerateTest.cpp b/folly/test/EnumerateTest.cpp new file mode 100644 index 00000000..1b6fe223 --- /dev/null +++ b/folly/test/EnumerateTest.cpp @@ -0,0 +1,127 @@ +/* + * 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 +#include + +#include +#include +#include + +TEST(Enumerate, Basic) { + std::vector 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 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 +struct IsConstReference { + constexpr static bool value = false; +}; +template +struct IsConstReference { + constexpr static bool value = true; +}; + +TEST(Enumerate, BasicConstArg) { + const std::vector v = {"abc", "a", "ab"}; + size_t i = 0; + for (auto it : folly::enumerate(v)) { + static_assert( + IsConstReference::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 v = {"abc", "a", "ab"}; + size_t i = 0; + for (const auto it : folly::enumerate(v)) { + static_assert(IsConstReference::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 v = {"abc", "a", "ab"}; + size_t i = 0; + for (const auto it : folly::enumerate(decltype(v)(v))) { // Copy v. + static_assert(IsConstReference::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 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 v; + for (auto it : folly::enumerate(v)) { + (void)it; // Silence warnings. + EXPECT_TRUE(false); + } +} -- 2.34.1