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 <iostream>
#include <iterator>
#include <stdexcept>
#include <utility>
 
namespace itertools {
 
  class range {
    // First value of the range.
    long first_ = 0;
 
    // Last value of the range (excluded).
    long last_ = -1;
 
    // 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;
 
    [[deprecated("range default construction deprecated. Use range::all for full range in slicing operation")]] range() = default;
 
    range(long first, long last) noexcept : first_(first), last_(last) {}
 
    range(long first, long last, long 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(long 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 last;
 
      long step;
 
      using value_type = long;
 
      using iterator_category = std::forward_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 c = *this;
        pos += step;
        return c;
      }
 
      [[nodiscard]] bool atEnd() const noexcept { return step > 0 ? pos >= last : pos <= last; }
 
      [[nodiscard]] bool operator==(const_iterator const &other) const noexcept {
        return (other.pos == this->pos) || (other.atEnd() && this->atEnd());
      }
 
      [[nodiscard]] bool operator!=(const_iterator const &other) const noexcept { return (!operator==(other)); }
 
      [[nodiscard]] long operator*() const noexcept { return pos; }
 
      [[nodiscard]] long operator->() const noexcept { return operator*(); }
    };
 
    [[nodiscard]] const_iterator cbegin() const noexcept { return {first_, last_, step_}; }
 
    [[nodiscard]] const_iterator begin() const noexcept { return {first_, last_, step_}; }
 
    [[nodiscard]] const_iterator cend() const noexcept { return {last_, last_, step_}; }
 
    [[nodiscard]] const_iterator end() const noexcept { return {last_, last_, step_}; }
  };
 
  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