folly/container/Enumerate.h

   1 /*
   2  * Copyright 2017-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
  47 namespace folly {
  48
  49 namespace detail {
  50
  51 template <class T>
  52 struct MakeConst {
  53   using type = const T;
  54 };
  55 template <class T>
  56 struct MakeConst<T&> {
  57   using type = const T&;
  58 };
  59 template <class T>
  60 struct 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).
  67 template <class Iterator>
  68 FOLLY_ALWAYS_INLINE auto getPointer(const Iterator& it, long)
  69     -> decltype(std::addressof(*it)) {
  70   return std::addressof(*it);
  71 }
  72 template <class Iterator>
  73 FOLLY_ALWAYS_INLINE auto getPointer(const Iterator& it, int)
  74     -> decltype(it.operator->()) {
  75   return it.operator->();
  76 }
  77
  78 template <class Iterator>
  79 class 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
 147 template <class Range>
 148 class 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
 166 template <class Range>
 167 detail::RangeEnumerator<Range> enumerate(Range&& r) {
 168   return detail::RangeEnumerator<Range>(std::forward<Range>(r));
 169 }
 170
 171 } // namespace folly