1/*
2 * Copyright 2016-present Facebook, Inc.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#pragma once
18
19#include <iterator>
20#include <memory>
21
22#include <folly/CPortability.h>
23#include <folly/portability/SysTypes.h>
24
25/**
26 * Similar to Python's enumerate(), folly::enumerate() can be used to
27 * iterate a range with a for-range loop, and it also allows to
28 * retrieve the count of iterations so far.
29 *
30 * For example:
31 *
32 * for (auto&& it : folly::enumerate(vec)) {
33 * // *it is a reference to the current element. Const if vec is const.
34 * // it->member can be used as well.
35 * // it.index contains the iteration count.
36 * }
37 *
38 * If the iteration variable is const, the reference is too.
39 *
40 * for (const auto&& it : folly::enumerate(vec)) {
41 * // *it is always a const reference.
42 * }
43 *
44 * @author Giuseppe Ottaviano <ott@fb.com>
45 */
46
47namespace folly {
48
49namespace detail {
50
51template <class T>
52struct MakeConst {
53 using type = const T;
54};
55template <class T>
56struct MakeConst<T&> {
57 using type = const T&;
58};
59template <class T>
60struct MakeConst<T*> {
61 using type = const T*;
62};
63
64// Raw pointers don't have an operator->() member function, so the
65// second overload will be SFINAEd out in that case. Otherwise, the
66// second is preferred in the partial order for getPointer(_, 0).
67template <class Iterator>
68FOLLY_ALWAYS_INLINE auto getPointer(const Iterator& it, long)
69 -> decltype(std::addressof(*it)) {
70 return std::addressof(*it);
71}
72template <class Iterator>
73FOLLY_ALWAYS_INLINE auto getPointer(const Iterator& it, int)
74 -> decltype(it.operator->()) {
75 return it.operator->();
76}
77
78template <class Iterator>
79class Enumerator {
80 public:
81 explicit Enumerator(Iterator it) : it_(std::move(it)) {}
82
83 class Proxy {
84 public:
85 using difference_type = ssize_t;
86 using value_type = typename std::iterator_traits<Iterator>::value_type;
87 using reference = typename std::iterator_traits<Iterator>::reference;
88 using pointer = typename std::iterator_traits<Iterator>::pointer;
89 using iterator_category = std::input_iterator_tag;
90
91 FOLLY_ALWAYS_INLINE explicit Proxy(const Enumerator& e)
92 : it_(e.it_), index(e.idx_) {}
93
94 // Non-const Proxy: Forward constness from Iterator.
95 FOLLY_ALWAYS_INLINE reference operator*() {
96 return *it_;
97 }
98 FOLLY_ALWAYS_INLINE pointer operator->() {
99 return getPointer(it_, 0);
100 }
101
102 // Const Proxy: Force const references.
103 FOLLY_ALWAYS_INLINE typename MakeConst<reference>::type operator*() const {
104 return *it_;
105 }
106 FOLLY_ALWAYS_INLINE typename MakeConst<pointer>::type operator->() const {
107 return getPointer(it_, 0);
108 }
109
110 private:
111 const Iterator& it_;
112
113 public:
114 const size_t index;
115 };
116
117 FOLLY_ALWAYS_INLINE Proxy operator*() const {
118 return Proxy(*this);
119 }
120
121 FOLLY_ALWAYS_INLINE Enumerator& operator++() {
122 ++it_;
123 ++idx_;
124 return *this;
125 }
126
127 template <typename OtherIterator>
128 FOLLY_ALWAYS_INLINE bool operator==(
129 const Enumerator<OtherIterator>& rhs) const {
130 return it_ == rhs.it_;
131 }
132
133 template <typename OtherIterator>
134 FOLLY_ALWAYS_INLINE bool operator!=(
135 const Enumerator<OtherIterator>& rhs) const {
136 return !(it_ == rhs.it_);
137 }
138
139 private:
140 template <typename OtherIterator>
141 friend class Enumerator;
142
143 Iterator it_;
144 size_t idx_ = 0;
145};
146
147template <class Range>
148class RangeEnumerator {
149 Range r_;
150 using BeginIteratorType = decltype(std::declval<Range>().begin());
151 using EndIteratorType = decltype(std::declval<Range>().end());
152
153 public:
154 explicit RangeEnumerator(Range&& r) : r_(std::forward<Range>(r)) {}
155
156 Enumerator<BeginIteratorType> begin() {
157 return Enumerator<BeginIteratorType>(r_.begin());
158 }
159 Enumerator<EndIteratorType> end() {
160 return Enumerator<EndIteratorType>(r_.end());
161 }
162};
163
164} // namespace detail
165
166template <class Range>
167detail::RangeEnumerator<Range> enumerate(Range&& r) {
168 return detail::RangeEnumerator<Range>(std::forward<Range>(r));
169}
170
171} // namespace folly
172