| 33 | |
| 34 | template <typename T> |
| 35 | struct StridedIterator { |
| 36 | using iterator_category = std::random_access_iterator_tag; |
| 37 | using difference_type = int32_t; |
| 38 | using value_type = T; |
| 39 | using reference = value_type&; |
| 40 | using pointer = value_type*; |
| 41 | |
| 42 | // Constructors |
| 43 | StridedIterator() = default; |
| 44 | |
| 45 | explicit StridedIterator(T* ptr, int64_t stride, difference_type offset = 0) |
| 46 | : stride_(stride), ptr_(ptr + offset * stride) {} |
| 47 | |
| 48 | explicit StridedIterator(array& arr, int axis, difference_type offset = 0) |
| 49 | : StridedIterator(arr.data<T>(), arr.strides()[axis], offset) {} |
| 50 | |
| 51 | // Accessors |
| 52 | reference operator*() const { |
| 53 | return ptr_[0]; |
| 54 | } |
| 55 | |
| 56 | reference operator[](difference_type idx) const { |
| 57 | return ptr_[idx * stride_]; |
| 58 | } |
| 59 | |
| 60 | // Comparisons |
| 61 | bool operator==(const StridedIterator& other) const { |
| 62 | return ptr_ == other.ptr_ && stride_ == other.stride_; |
| 63 | } |
| 64 | |
| 65 | bool operator!=(const StridedIterator& other) const { |
| 66 | return ptr_ != other.ptr_; |
| 67 | } |
| 68 | |
| 69 | bool operator<(const StridedIterator& other) const { |
| 70 | return ptr_ < other.ptr_; |
| 71 | } |
| 72 | |
| 73 | bool operator>(const StridedIterator& other) const { |
| 74 | return ptr_ > other.ptr_; |
| 75 | } |
| 76 | |
| 77 | bool operator<=(const StridedIterator& other) const { |
| 78 | return ptr_ <= other.ptr_; |
| 79 | } |
| 80 | |
| 81 | bool operator>=(const StridedIterator& other) const { |
| 82 | return ptr_ >= other.ptr_; |
| 83 | } |
| 84 | |
| 85 | difference_type operator-(const StridedIterator& other) const { |
| 86 | return (ptr_ - other.ptr_) / stride_; |
| 87 | } |
| 88 | |
| 89 | // Moving |
| 90 | StridedIterator& operator++() { |
| 91 | ptr_ += stride_; |
| 92 | return *this; |