iterators.hpp

Go to the documentation of this file.

// Copyright (c) 2019-2024 Simons Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0.txt
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Thomas Hahn, Olivier Parcollet, Nils Wentzell
 
#pragma once
 
#include "./stdutil/array.hpp"
 
#include <array>
#include <cstddef>
#include <iterator>
 
namespace nda {
 
  namespace detail {
 
    // N-dimensional rectangular grid iterator in C traversal order.
    template <int Rank>
    class grid_iterator {
      // Stride or number of elements to skip when increasing the iterator in the current dimension.
      long stri = 0;
 
      // Position of the iterator in the current dimension.
      long pos = 0;
 
      // Position times the stride in the current dimension.
      long offset = 0;
 
      // Iterator to the beginning of the Rank - 1 dimensional grid (excluding the current dimension).
      grid_iterator<Rank - 1> it_begin;
 
      // Iterator to the end of the Rank - 1 dimensional grid (excluding the current dimension).
      grid_iterator<Rank - 1> it_end;
 
      // Iterator to the Rank - 1 dimensional grid (excluding the current dimension).
      grid_iterator<Rank - 1> it;
 
      public:
      using iterator_category = std::forward_iterator_tag;
      using value_type        = long;
      using difference_type   = std::ptrdiff_t;
      using pointer           = long *;
      using reference         = long &;
 
      // Default constructor.
      grid_iterator() = default;
 
      // Construct an iterator from the shape of the grid, the stride of the grid and a flag indicating if the iterator
      // is at the end.
      grid_iterator(long const *lengths, long const *strides, bool at_end)
         : stri(strides[0]),
           pos(at_end ? lengths[0] : 0),
           offset(pos * stri),
           it_begin(lengths + 1, strides + 1, false),
           it_end(lengths + 1, strides + 1, true),
           it(it_begin) {}
 
      // Get the position/multi-dimensional index of the iterator.
      [[nodiscard]] std::array<long, Rank> indices() { return stdutil::front_append(it.indices(), pos); }
 
      // Dereference operator returns the sum of the offsets of every dimension = its linear index.
      [[nodiscard]] long operator*() const { return offset + *it; }
 
      // Member access operator returns the sum of the offsets of every dimension = its linear index.
      [[nodiscard]] long operator->() const { return operator*(); }
 
      // True if the positions of the iterators are equal in every dimension, false otherwise.
      [[nodiscard]] bool operator==(grid_iterator const &rhs) const { return ((rhs.pos == pos) and (rhs.it == it)); }
 
      // True if the positions of the iterators are not equal in every dimension, false otherwise.
      [[nodiscard]] bool operator!=(grid_iterator const &rhs) const { return not operator==(rhs); }
 
      // Prefix increment operator.
      grid_iterator &operator++() {
        // increment the iterator of the subgrid
        ++it;
 
        // if the iterator of the subgrid is at the end, reset it and increment the current position and offset
        if (it == it_end) { //FIXME [[unlikely]]
          ++pos;
          offset += stri;
          it = it_begin;
        }
        return *this;
      }
 
      // Postfix increment operator.
      grid_iterator operator++(int) {
        auto c = *this;
        ++(*this);
        return c;
      }
    };
 
    // Specialization of nda::grid_iterator for 1-dimensional grids.
    template <>
    class grid_iterator<1> {
      // Stride or number of elements to skip when increasing the iterator.
      long stri = 0;
 
      // Position of the iterator.
      long pos = 0;
 
      // Position times the stride.
      long offset = 0;
 
      public:
      using iterator_category = std::forward_iterator_tag;
      using value_type        = long;
      using difference_type   = std::ptrdiff_t;
      using pointer           = long *;
      using reference         = long &;
 
      // Default constructor.
      grid_iterator() = default;
 
      // Construct an iterator from the shape of the grid, the stride of the grid and a flag indicating if the iterator
      // is at the end.
      grid_iterator(long const *lengths, long const *strides, bool at_end) : stri(strides[0]), pos(at_end ? lengths[0] : 0), offset(pos * stri) {}
 
      // Get the position/index of the iterator.
      [[nodiscard]] std::array<long, 1> indices() { return {pos}; }
 
      // Dereference operator returns the offset = its linear index.
      [[nodiscard]] long operator*() const { return offset; }
 
      // Member access operator returns the offset = its linear index.
      [[nodiscard]] long operator->() const { return operator*(); }
 
      // True if the positions of the iterators are equal, false otherwise.
      [[nodiscard]] bool operator==(grid_iterator const &rhs) const { return (rhs.pos == pos); }
 
      // True if the positions of the iterators are not equal, false otherwise.
      [[nodiscard]] bool operator!=(grid_iterator const &rhs) const { return (rhs.pos != pos); }
 
      // Prefix increment operator increments the offset by the stride and the position by one.
      grid_iterator &operator++() {
        offset += stri;
        ++pos;
        return *this;
      }
 
      // Prefix decrement operator decrements the offset by the stride and the position by one.
      grid_iterator &operator--() {
        offset -= stri;
        --pos;
        return *this;
      }
 
      // Compound assignment addition operator increments the offset by n times the stride and the position by n.
      grid_iterator &operator+=(std::ptrdiff_t n) {
        offset += n * stri;
        pos += n;
        return *this;
      }
 
      // Binary addition of a grid iterator and an integer.
      [[nodiscard]] friend grid_iterator operator+(grid_iterator it, std::ptrdiff_t n) { return it += n; }
 
      // Binary subtraction of two grid iterators.
      [[nodiscard]] friend std::ptrdiff_t operator-(grid_iterator const &lhs, grid_iterator const &rhs) { return lhs.pos - rhs.pos; }
 
      // True if the position of the left hand side iterator is less than the position of the right hand side iterator,
      // false otherwise.
      [[nodiscard]] friend bool operator<(grid_iterator const &lhs, grid_iterator const &rhs) { return lhs.pos < rhs.pos; }
 
      // True if the position of the left hand side iterator is greater than the position of the right hand side
      // iterator, false otherwise.
      [[nodiscard]] friend bool operator>(grid_iterator const &lhs, grid_iterator const &rhs) { return lhs.pos > rhs.pos; }
    };
 
  } // namespace detail
 
  template <int Rank, typename T, typename Pointer>
  class array_iterator {
    // Pointer to the data (to the first element).
    T *data = nullptr;
 
    // Shape of the array.
    std::array<long, Rank> len;
 
    // Strides of the array.
    std::array<long, Rank> stri;
 
    // Grid iterator.
    detail::grid_iterator<Rank> iter;
 
    public:
    using iterator_category = std::forward_iterator_tag;
 
    using value_type = T;
 
    using difference_type = std::ptrdiff_t;
 
    using pointer = T *;
 
    using reference = T &;
 
    array_iterator() = default;
 
    array_iterator(std::array<long, Rank> const &lengths, std::array<long, Rank> const &strides, T *start, bool at_end)
       : data(start), len(lengths), stri(strides), iter(len.data(), stri.data(), at_end) {}
 
    [[nodiscard]] auto indices() { return iter.indices(); }
 
    [[nodiscard]] value_type &operator*() const { return ((Pointer)data)[*iter]; }
 
    [[nodiscard]] value_type &operator->() const { return operator*(); }
 
    array_iterator &operator++() {
      ++iter;
      return *this;
    }
 
    array_iterator operator++(int) {
      auto c = *this;
      ++iter;
      return c;
    }
 
    [[nodiscard]] bool operator==(array_iterator const &rhs) const { return (rhs.iter == iter); }
 
    [[nodiscard]] bool operator!=(array_iterator const &rhs) const { return (!operator==(rhs)); }
  };
 
  template <typename T, typename Pointer>
  class array_iterator<1, T, Pointer> {
    // Pointer to the data.
    T *data = nullptr;
 
    // Shape of the array.
    std::array<long, 1> len{};
 
    // Strides of the array.
    std::array<long, 1> stri{};
 
    // Grid iterator.
    detail::grid_iterator<1> iter;
 
    public:
    using iterator_category = std::random_access_iterator_tag;
 
    using value_type = T;
 
    using difference_type = std::ptrdiff_t;
 
    using pointer = T *;
 
    using reference = T &;
 
    array_iterator() = default;
 
    array_iterator(std::array<long, 1> const &lengths, std::array<long, 1> const &strides, T *start, bool at_end)
       : data(start), len(lengths), stri(strides), iter(len.data(), stri.data(), at_end) {}
 
    [[nodiscard]] auto indices() { return iter.indices(); }
 
    [[nodiscard]] T &operator*() const { return ((Pointer)data)[*iter]; }
 
    T &operator->() const { return operator*(); }
 
    array_iterator &operator++() {
      ++iter;
      return *this;
    }
 
    array_iterator operator++(int) {
      auto c = *this;
      ++iter;
      return c;
    }
 
    array_iterator &operator--() {
      --iter;
      return *this;
    }
 
    array_iterator operator--(int) {
      auto c = *this;
      --iter;
      return c;
    }
 
    [[nodiscard]] bool operator==(array_iterator const &rhs) const { return (rhs.iter == iter); }
 
    [[nodiscard]] bool operator!=(array_iterator const &rhs) const { return (!operator==(rhs)); }
 
    array_iterator &operator+=(std::ptrdiff_t n) {
      iter += n;
      return *this;
    }
 
    array_iterator &operator-=(std::ptrdiff_t n) {
      iter += (-n);
      return *this;
    }
 
    [[nodiscard]] friend array_iterator operator+(std::ptrdiff_t n, array_iterator it) { return it += n; }
 
    [[nodiscard]] friend array_iterator operator+(array_iterator it, std::ptrdiff_t n) { return it += n; }
 
    [[nodiscard]] friend array_iterator operator-(array_iterator it, std::ptrdiff_t n) { return it -= n; }
 
    [[nodiscard]] friend std::ptrdiff_t operator-(array_iterator const &lhs, array_iterator const &rhs) { return lhs.iter - rhs.iter; }
 
    [[nodiscard]] T &operator[](std::ptrdiff_t n) { return ((Pointer)data)[*(iter + n)]; }
 
    // FIXME C++20 ? with <=> operator
    [[nodiscard]] friend bool operator<(array_iterator const &lhs, array_iterator const &rhs) { return lhs.iter < rhs.iter; }
 
    [[nodiscard]] friend bool operator>(array_iterator const &lhs, array_iterator const &rhs) { return lhs.iter > rhs.iter; }
 
    [[nodiscard]] friend bool operator<=(array_iterator const &lhs, array_iterator const &rhs) { return not(lhs.iter > rhs.iter); }
 
    [[nodiscard]] friend bool operator>=(array_iterator const &lhs, array_iterator const &rhs) { return not(lhs.iter < rhs.iter); }
  };
 
} // namespace nda