stride.hpp

Go to the documentation of this file.

// Copyright (c) 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, chuffa
 
#ifndef _ITERTOOLS_STRIDE_HPP
#define _ITERTOOLS_STRIDE_HPP
 
#include "./iterator_facade.hpp"
#include "./utils.hpp"
 
#include <cstddef>
#include <iterator>
#include <stdexcept>
#include <utility>
 
namespace itertools {
 
  template <typename Iter> struct stride_iter : iterator_facade<stride_iter<Iter>, typename std::iterator_traits<Iter>::value_type> {
    Iter it;
 
    std::ptrdiff_t stride{1};
 
    stride_iter() = default;
 
    stride_iter(Iter it, std::ptrdiff_t stride) : it(it), stride(stride) {
      if (stride <= 0) throw std::runtime_error("The itertools::strided range requires a positive stride");
    }
 
    void increment() { std::advance(it, stride); }
 
    [[nodiscard]] bool operator==(stride_iter const &other) const { return it == other.it; }
 
    [[nodiscard]] decltype(auto) dereference() const { return *it; }
  };
 
  template <typename R> struct strided {
    R rg;
 
    std::ptrdiff_t stride;
 
    using iterator = stride_iter<decltype(std::begin(rg))>;
 
    using const_iterator = stride_iter<decltype(std::cbegin(rg))>;
 
    [[nodiscard]] bool operator==(strided const &) const = default;
 
    private:
    // Helper function to calculate the index of the end iterator.
    [[nodiscard]] std::ptrdiff_t end_offset() const {
      auto size = distance(std::cbegin(rg), std::cend(rg));
      return (size == 0) ? 0 : ((size - 1) / stride + 1) * stride;
    }
 
    public:
    [[nodiscard]] iterator begin() noexcept { return {std::begin(rg), stride}; }
 
    [[nodiscard]] const_iterator cbegin() const noexcept { return {std::cbegin(rg), stride}; }
 
    [[nodiscard]] const_iterator begin() const noexcept { return cbegin(); }
 
    [[nodiscard]] iterator end() noexcept { return {std::next(std::begin(rg), end_offset()), stride}; }
 
    [[nodiscard]] const_iterator cend() const noexcept { return {std::next(std::cbegin(rg), end_offset()), stride}; }
 
    [[nodiscard]] const_iterator end() const noexcept { return cend(); }
  };
 
  template <typename R> [[nodiscard]] strided<R> stride(R &&rg, std::ptrdiff_t stride) { return {std::forward<R>(rg), stride}; }
 
} // namespace itertools
 
#endif // _ITERTOOLS_STRIDE_HPP