folly/Enumerate.h

   1 /*
   2  * Copyright 2017 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/portability/SysTypes.h>
  23
  24 /**
  25  * Similar to Python's enumerate(), folly::enumerate() can be used to
  26  * iterate a range with a for-range loop, and it also allows to
  27  * retrieve the count of iterations so far.
  28  *
  29  * For example:
  30  *
  31  * for (auto it : folly::enumerate(vec)) {
  32  *   // *it is a reference to the current element. Const if vec is const.
  33  *   // it->member can be used as well.
  34  *   // it.index contains the iteration count.
  35  * }
  36  *
  37  * If the iteration variable is const, the reference is too.
  38  *
  39  * for (const auto it : folly::enumerate(vec)) {
  40  *   // *it is always a const reference.
  41  * }
  42  *
  43  * @author Giuseppe Ottaviano <ott@fb.com>
  44  */
  45
  46 namespace folly {
  47
  48 namespace detail {
  49
  50 template <class T>
  51 struct MakeConst {
  52   using type = const T;
  53 };
  54 template <class T>
  55 struct MakeConst<T&> {
  56   using type = const T&;
  57 };
  58 template <class T>
  59 struct MakeConst<T*> {
  60   using type = const T*;
  61 };
  62
  63 // Raw pointers don't have an operator->() member function, so the
  64 // second overload will be SFINAEd out in that case. Otherwise, the
  65 // second is preferred in the partial order for getPointer(_, 0).
  66 template <class Iterator>
  67 auto getPointer(const Iterator& it, long) -> decltype(std::addressof(*it)) {
  68   return std::addressof(*it);
  69 }
  70 template <class Iterator>
  71 auto getPointer(const Iterator& it, int) -> decltype(it.operator->()) {
  72   return it.operator->();
  73 }
  74
  75 template <class Iterator>
  76 class Enumerator {
  77  public:
  78   explicit Enumerator(Iterator it) : it_(std::move(it)) {}
  79
  80   class Proxy {
  81    public:
  82     using difference_type = ssize_t;
  83     using value_type = typename std::iterator_traits<Iterator>::value_type;
  84     using reference = typename std::iterator_traits<Iterator>::reference;
  85     using pointer = typename std::iterator_traits<Iterator>::pointer;
  86     using iterator_category = std::input_iterator_tag;
  87
  88     explicit Proxy(const Enumerator* e) : it_(e->it_), index(e->idx_) {}
  89
  90     // Non-const Proxy: Forward constness from Iterator.
  91     reference operator*() {
  92       return *it_;
  93     }
  94     pointer operator->() {
  95       return getPointer(it_, 0);
  96     }
  97
  98     // Const Proxy: Force const references.
  99     typename MakeConst<reference>::type operator*() const {
 100       return *it_;
 101     }
 102     typename MakeConst<pointer>::type operator->() const {
 103       return getPointer(it_, 0);
 104     }
 105
 106    private:
 107     const Iterator& it_;
 108
 109    public:
 110     const size_t index;
 111   };
 112
 113   Proxy operator*() const {
 114     return Proxy(this);
 115   }
 116
 117   Enumerator& operator++() {
 118     ++it_;
 119     ++idx_;
 120     return *this;
 121   }
 122
 123   template <typename OtherIterator>
 124   bool operator==(const Enumerator<OtherIterator>& rhs) {
 125     return it_ == rhs.it_;
 126   }
 127
 128   template <typename OtherIterator>
 129   bool operator!=(const Enumerator<OtherIterator>& rhs) {
 130     return !(*this == rhs);
 131   }
 132
 133  private:
 134   template <typename OtherIterator>
 135   friend class Enumerator;
 136
 137   Iterator it_;
 138   size_t idx_ = 0;
 139 };
 140
 141 template <class Range>
 142 class RangeEnumerator {
 143   Range r_;
 144   using BeginIteratorType = decltype(std::declval<Range>().begin());
 145   using EndIteratorType = decltype(std::declval<Range>().end());
 146
 147  public:
 148   explicit RangeEnumerator(Range&& r) : r_(std::forward<Range>(r)) {}
 149
 150   Enumerator<BeginIteratorType> begin() {
 151     return Enumerator<BeginIteratorType>(r_.begin());
 152   }
 153   Enumerator<EndIteratorType> end() {
 154     return Enumerator<EndIteratorType>(r_.end());
 155   }
 156 };
 157
 158 } // namespace detail
 159
 160 template <class Range>
 161 detail::RangeEnumerator<Range> enumerate(Range&& r) {
 162   return detail::RangeEnumerator<Range>(std::forward<Range>(r));
 163 }
 164
 165 } // namespace folly