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