1 /* Copyright 2018 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_CORE_UTIL_PERMUTATION_OUTPUT_ITERATOR_H_ 17 #define TENSORFLOW_CORE_UTIL_PERMUTATION_OUTPUT_ITERATOR_H_ 18 19 #include <iostream> 20 #include <iterator> 21 22 namespace tensorflow { 23 24 template <typename ValueType, typename OutputIteratorT, typename IndexIteratorT, 25 typename OffsetT = ptrdiff_t> 26 class PermutationOutputIterator { 27 public: 28 // Required iterator traits 29 typedef PermutationOutputIterator 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 OutputIteratorT output_itr; 44 IndexIteratorT index_itr; 45 46 public: 47 /// Constructor PermutationOutputIterator(OutputIteratorT output_itr,IndexIteratorT index_itr)48 __host__ __device__ __forceinline__ PermutationOutputIterator( 49 OutputIteratorT output_itr, ///< Input iterator to wrap 50 IndexIteratorT index_itr) ///< Conversion functor to wrap 51 : output_itr(output_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 output_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(output_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(output_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 output_itr[index_itr[n]]; 109 } 110 111 /// Equal to 112 __host__ __device__ __forceinline__ bool operator==(const self_type& rhs) { 113 return (index_itr == rhs.index_itr && output_itr == rhs.output_itr); 114 } 115 116 /// Not equal to 117 __host__ __device__ __forceinline__ bool operator!=(const self_type& rhs) { 118 return !(*this == rhs); 119 } 120 121 /// ostream operator 122 friend std::ostream& operator<<(std::ostream& os, const self_type& itr) { 123 return os; 124 } 125 }; 126 127 } // end namespace tensorflow 128 129 #endif // TENSORFLOW_CORE_UTIL_PERMUTATION_OUTPUT_ITERATOR_H_ 130