Home | History | Annotate | Download | only in util 
      1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
      2 
      3 Licensed under the Apache License, Version 2.0 (the "License");
      4 you may not use this file except in compliance with the License.
      5 You may obtain a copy of the License at
      6 
      7     http://www.apache.org/licenses/LICENSE-2.0
      8 
      9 Unless required by applicable law or agreed to in writing, software
     10 distributed under the License is distributed on an "AS IS" BASIS,
     11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 See the License for the specific language governing permissions and
     13 limitations under the License.
     14 ==============================================================================*/
     15 
     16 #ifndef TENSORFLOW_UTIL_PERMUTATION_INPUT_ITERATOR_H_
     17 #define TENSORFLOW_UTIL_PERMUTATION_INPUT_ITERATOR_H_
     18 
     19 #include <iostream>
     20 #include <iterator>
     21 
     22 namespace tensorflow {
     23 
     24 template <typename ValueType, typename InputIteratorT, typename IndexIteratorT,
     25           typename OffsetT = ptrdiff_t>
     26 class PermutationInputIterator {
     27  public:
     28   // Required iterator traits
     29   typedef PermutationInputIterator self_type;  ///< My own type
     30   typedef OffsetT difference_type;  ///< Type to express the result of
     31                                     ///< subtracting one iterator from another
     32   typedef ValueType
     33       value_type;  ///< The type of the element the iterator can point to
     34   typedef ValueType* pointer;   ///< The type of a pointer to an element the
     35                                 ///< iterator can point to
     36   typedef ValueType reference;  ///< The type of a reference to an element the
     37                                 ///< iterator can point to
     38 
     39   typedef std::random_access_iterator_tag
     40       iterator_category;  ///< The iterator category
     41 
     42  private:
     43   InputIteratorT input_itr;
     44   IndexIteratorT index_itr;
     45 
     46  public:
     47   /// Constructor
     48   __host__ __device__ __forceinline__ PermutationInputIterator(
     49       InputIteratorT input_itr,  ///< Input iterator to wrap
     50       IndexIteratorT index_itr)  ///< Conversion functor to wrap
     51       : input_itr(input_itr), index_itr(index_itr) {}
     52 
     53   /// Postfix increment
     54   __host__ __device__ __forceinline__ self_type operator++(int) {
     55     self_type retval = *this;
     56     index_itr++;
     57     return retval;
     58   }
     59 
     60   /// Prefix increment
     61   __host__ __device__ __forceinline__ self_type operator++() {
     62     index_itr++;
     63     return *this;
     64   }
     65 
     66   /// Indirection
     67   __host__ __device__ __forceinline__ reference operator*() const {
     68     return input_itr[*index_itr];
     69   }
     70 
     71   /// Addition
     72   template <typename Distance>
     73   __host__ __device__ __forceinline__ self_type operator+(Distance n) const {
     74     self_type retval(input_itr, index_itr + n);
     75     return retval;
     76   }
     77 
     78   /// Addition assignment
     79   template <typename Distance>
     80   __host__ __device__ __forceinline__ self_type& operator+=(Distance n) {
     81     index_itr += n;
     82     return *this;
     83   }
     84 
     85   /// Subtraction
     86   template <typename Distance>
     87   __host__ __device__ __forceinline__ self_type operator-(Distance n) const {
     88     self_type retval(input_itr, index_itr - n);
     89     return retval;
     90   }
     91 
     92   /// Subtraction assignment
     93   template <typename Distance>
     94   __host__ __device__ __forceinline__ self_type& operator-=(Distance n) {
     95     index_itr -= n;
     96     return *this;
     97   }
     98 
     99   /// Distance
    100   __host__ __device__ __forceinline__ difference_type
    101   operator-(self_type other) const {
    102     return index_itr - other.index_itr;
    103   }
    104 
    105   /// Array subscript
    106   template <typename Distance>
    107   __host__ __device__ __forceinline__ reference operator[](Distance n) const {
    108     return input_itr[index_itr[n]];
    109   }
    110 
    111   /// Structure dereference
    112   __host__ __device__ __forceinline__ pointer operator->() {
    113     return input_itr + *index_itr;
    114   }
    115 
    116   /// Equal to
    117   __host__ __device__ __forceinline__ bool operator==(const self_type& rhs) {
    118     return (index_itr == rhs.index_itr && input_itr == rhs.input_itr);
    119   }
    120 
    121   /// Not equal to
    122   __host__ __device__ __forceinline__ bool operator!=(const self_type& rhs) {
    123     return !(*this == rhs);
    124   }
    125 
    126   /// ostream operator
    127   friend std::ostream& operator<<(std::ostream& os, const self_type& itr) {
    128     return os;
    129   }
    130 };
    131 
    132 }  // end namespace tensorflow
    133 
    134 #endif  // TENSORFLOW_UTIL_PERMUTATION_INPUT_ITERATOR_H_
    135