range.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_RANGE_HPP
#define _ITERTOOLS_RANGE_HPP
 
#include "./product.hpp"
 
#include <algorithm>
#include <array>
#include <cassert>
#include <compare>
#include <concepts>
#include <iostream>
#include <iterator>
#include <stdexcept>
#include <utility>
 
namespace itertools {


  class range {
    // First value of the range.
    long first_;
 
    // Last value of the range (excluded).
    long last_;
 
    // Number of integers between two elements of the range.
    long step_ = 1;
 
    public:
    struct all_t {};

    static inline constexpr all_t all = {};

    using index_t = long;

    range(std::integral auto first, std::integral auto last) noexcept : first_(first), last_(last) {}

    range(std::integral auto first, std::integral auto last, std::integral auto step) : first_(first), last_(last), step_(step) {
      if (step_ == 0) throw std::runtime_error("Step-size cannot be zero in construction of integer range");
    }

    explicit range(std::integral auto last) : range(0, last, 1) {}

    [[nodiscard]] bool operator==(range const &) const = default;

    [[nodiscard]] long first() const { return first_; }

    [[nodiscard]] long last() const { return last_; }

    [[nodiscard]] long step() const { return step_; }

    [[nodiscard]] long size() const { return std::max(0l, (last_ + step_ - (step_ > 0 ? 1 : -1) - first_) / step_); }

    [[nodiscard]] range operator+(long shift) const { return {first_ + shift, last_ + shift, step_}; }

    friend inline std::ostream &operator<<(std::ostream &os, const range &rg) {
      os << "range(" << rg.first() << "," << rg.last() << "," << rg.step() << ")";
      return os;
    }

    struct const_iterator {
      long pos;

      long step;

      using value_type = long;

      using iterator_category = std::random_access_iterator_tag;

      using pointer = value_type *;

      using difference_type = std::ptrdiff_t;

      using reference = value_type const &;

      const_iterator &operator++() noexcept {
        pos += step;
        return *this;
      }

      const_iterator operator++(int) noexcept {
        const_iterator tmp = *this;
        pos += step;
        return tmp;
      }

      const_iterator &operator--() noexcept {
        pos -= step;
        return *this;
      }

      const_iterator operator--(int) noexcept {
        const_iterator tmp = *this;
        pos -= step;
        return tmp;
      }

      [[nodiscard]] std::strong_ordering operator<=>(const_iterator const &rhs) const noexcept {
        return (step > 0 ? pos <=> rhs.pos : -pos <=> -rhs.pos);
      }

      [[nodiscard]] bool operator==(const_iterator const &other) const noexcept { return pos == other.pos; }

      [[nodiscard]] long operator*() const noexcept { return pos; }

      [[nodiscard]] long operator->() const noexcept { return operator*(); }

      const_iterator &operator+=(difference_type n) noexcept {
        pos += n * step;
        return *this;
      }

      [[nodiscard]] const_iterator operator+(difference_type n) const noexcept { return {.pos = pos + n * step, .step = step}; }

      [[nodiscard]] friend const_iterator operator+(difference_type n, const_iterator it) noexcept { return it + n; }

      const_iterator &operator-=(difference_type n) noexcept {
        pos -= n * step;
        return *this;
      }

      [[nodiscard]] const_iterator operator-(difference_type n) const noexcept { return {.pos = pos - n * step, .step = step}; }

      [[nodiscard]] difference_type operator-(const_iterator const &rhs) const noexcept { return (pos - rhs.pos) / step; }

      [[nodiscard]] value_type operator[](difference_type n) const noexcept { return pos + n * step; }
 
    }; // end struct const_iterator

    using const_reverse_iterator = std::reverse_iterator<range::const_iterator>;

    [[nodiscard]] const_iterator cbegin() const noexcept { return {.pos = first_, .step = step_}; }

    [[nodiscard]] const_iterator begin() const noexcept { return {.pos = first_, .step = step_}; }

    [[nodiscard]] const_reverse_iterator crbegin() const noexcept { return const_reverse_iterator{end()}; }

    [[nodiscard]] const_reverse_iterator rbegin() const noexcept { return const_reverse_iterator{end()}; }

    [[nodiscard]] const_iterator cend() const noexcept { return {.pos = first_ + step_ * size(), .step = step_}; }

    [[nodiscard]] const_iterator end() const noexcept { return {.pos = first_ + step_ * size(), .step = step_}; }

    [[nodiscard]] const_reverse_iterator crend() const noexcept { return const_reverse_iterator{begin()}; }

    [[nodiscard]] const_reverse_iterator rend() const noexcept { return const_reverse_iterator{begin()}; }
  };

  template <typename... Is, typename EnableIf = std::enable_if_t<(std::is_integral_v<Is> and ...), int>> [[nodiscard]] auto product_range(Is... is) {
    return product(range(is)...);
  }
 
  namespace detail {
 
    // Helper function to create a product range of integer ranges from a tuple or an array.
    template <typename T, size_t... Is> [[gnu::always_inline]] auto product_range_impl(T const &idxs, std::index_sequence<Is...>) {
      return product_range(std::get<Is>(idxs)...);
    }
 
  } // namespace detail

  template <typename... Is, typename EnableIf = std::enable_if_t<(std::is_integral_v<Is> and ...), int>>
  [[nodiscard]] auto product_range(std::tuple<Is...> const &idx_tpl) {
    return detail::product_range_impl(idx_tpl, std::make_index_sequence<sizeof...(Is)>{});
  }

  template <typename I, size_t N, typename EnableIf = std::enable_if_t<std::is_integral_v<I>, int>>
  [[nodiscard]] auto product_range(std::array<I, N> const &idx_arr) {
    return detail::product_range_impl(idx_arr, std::make_index_sequence<N>{});
  }

  template <typename F> void foreach (range const &rg, F && f) {
    auto i = rg.first(), last = rg.last(), step = rg.step();
    for (; i < last; i += step) std::forward<F>(f)(i);
  }

 
} // namespace itertools
 
#endif // _ITERTOOLS_RANGE_HPP