1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7    http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
22namespace tensorflow {
23
24template <typename ValueType, typename InputIteratorT, typename IndexIteratorT,
25          typename OffsetT = ptrdiff_t>
26class 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