folly/Enumerate.h

   1 /*
   2  * Copyright 2016 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 template <class Iterator>
  64 class Enumerator {
  65  public:
  66   explicit Enumerator(Iterator it) : it_(std::move(it)) {}
  67
  68   class Proxy {
  69    public:
  70     using difference_type = ssize_t;
  71     using value_type = typename std::iterator_traits<Iterator>::value_type;
  72     using reference = typename std::iterator_traits<Iterator>::reference;
  73     using pointer = typename std::iterator_traits<Iterator>::pointer;
  74     using iterator_category = std::input_iterator_tag;
  75
  76     explicit Proxy(const Enumerator* e) : it_(e->it_), index(e->idx_) {}
  77
  78     // Non-const Proxy: Forward constness from Iterator.
  79     reference operator*() {
  80       return *it_;
  81     }
  82     pointer operator->() {
  83       return std::addressof(**this);
  84     }
  85
  86     // Const Proxy: Force const references.
  87     typename MakeConst<reference>::type operator*() const {
  88       return *it_;
  89     }
  90     typename MakeConst<pointer>::type operator->() const {
  91       return std::addressof(**this);
  92     }
  93
  94    private:
  95     const Iterator& it_;
  96
  97    public:
  98     const size_t index;
  99   };
 100
 101   Proxy operator*() const {
 102     return Proxy(this);
 103   }
 104
 105   Enumerator& operator++() {
 106     ++it_;
 107     ++idx_;
 108     return *this;
 109   }
 110
 111   template <typename OtherIterator>
 112   bool operator==(const Enumerator<OtherIterator>& rhs) {
 113     return it_ == rhs.it_;
 114   }
 115
 116   template <typename OtherIterator>
 117   bool operator!=(const Enumerator<OtherIterator>& rhs) {
 118     return !(*this == rhs);
 119   }
 120
 121  private:
 122   template <typename OtherIterator>
 123   friend class Enumerator;
 124
 125   Iterator it_;
 126   size_t idx_ = 0;
 127 };
 128
 129 template <class Range>
 130 class RangeEnumerator {
 131   Range r_;
 132   using BeginIteratorType = decltype(std::declval<Range>().begin());
 133   using EndIteratorType = decltype(std::declval<Range>().end());
 134
 135  public:
 136   explicit RangeEnumerator(Range&& r) : r_(std::forward<Range>(r)) {}
 137
 138   Enumerator<BeginIteratorType> begin() {
 139     return Enumerator<BeginIteratorType>(r_.begin());
 140   }
 141   Enumerator<EndIteratorType> end() {
 142     return Enumerator<EndIteratorType>(r_.end());
 143   }
 144 };
 145
 146 } // namespace detail
 147
 148 template <class Range>
 149 detail::RangeEnumerator<Range> enumerate(Range&& r) {
 150   return detail::RangeEnumerator<Range>(std::forward<Range>(r));
 151 }
 152
 153 } // namespace folly