product.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_PRODUCT_HPP
#define _ITERTOOLS_PRODUCT_HPP
 
#include "./iterator_facade.hpp"
#include "./sentinel.hpp"
 
#include <cstddef>
#include <iterator>
#include <tuple>
#include <utility>
 
namespace itertools {
 
  template <typename EndIters, typename... Iters>
  struct prod_iter : iterator_facade<prod_iter<EndIters, Iters...>, std::tuple<typename std::iterator_traits<Iters>::value_type...>> {
    std::tuple<Iters...> its_begin;
 
    EndIters its_end;
 
    std::tuple<Iters...> its = its_begin;
 
    static constexpr long Rank = sizeof...(Iters);
 
    prod_iter() = default;
 
    prod_iter(std::tuple<Iters...> its_begin, EndIters its_end) : its_begin(std::move(its_begin)), its_end(std::move(its_end)) {}
 
    private:
    // Helper function to recursively increment the current iterators.
    template <int N> void _increment() {
      // increment Nth iterator
      ++std::get<N>(its);
      // recursively increment previous iterators if necessary
      if constexpr (N > 0) {
        // if Nth iterator is at its end, reset it to its begin iterator and increment N-1st iterator
        if (std::get<N>(its) == std::get<N>(its_end)) {
          std::get<N>(its) = std::get<N>(its_begin);
          _increment<N - 1>();
        }
      }
    }
 
    public:
    void increment() { _increment<Rank - 1>(); }
 
    [[nodiscard]] bool operator==(prod_iter const &other) const { return its == other.its; }
 
    template <typename SentinelIter> [[nodiscard]] bool operator==(sentinel_t<SentinelIter> const &s) const { return (s.it == std::get<0>(its)); }
 
    private:
    // Helper function to dereference all original iterators.
    template <size_t... Is> [[gnu::always_inline]] [[nodiscard]] auto tuple_map_impl(std::index_sequence<Is...>) const {
      return std::tuple<decltype(*std::get<Is>(its))...>(*std::get<Is>(its)...);
    }
 
    public:
    [[nodiscard]] decltype(auto) dereference() const { return tuple_map_impl(std::index_sequence_for<Iters...>{}); }
  };
 
  template <typename... Rs> struct multiplied {
    std::tuple<Rs...> tu;
 
    using iterator = prod_iter<std::tuple<decltype(std::end(std::declval<Rs &>()))...>, decltype(std::begin(std::declval<Rs &>()))...>;
 
    using const_iterator = prod_iter<std::tuple<decltype(std::cend(std::declval<Rs &>()))...>, decltype(std::cbegin(std::declval<Rs &>()))...>;
 
    template <typename... Us> multiplied(Us &&...rgs) : tu{std::forward<Us>(rgs)...} {}
 
    [[nodiscard]] bool operator==(multiplied const &) const = default;
 
    private:
    // Helper function to create a itertools::prod_iter representing the beginning of the product range.
    template <size_t... Is> [[gnu::always_inline]] auto _begin(std::index_sequence<Is...>) {
      return iterator{std::make_tuple(std::begin(std::get<Is>(tu))...), std::make_tuple(std::end(std::get<Is>(tu))...)};
    }
 
    // Const version of _begin(std::index_sequence<Is...>).
    template <size_t... Is> [[gnu::always_inline]] auto _cbegin(std::index_sequence<Is...>) const {
      return const_iterator{std::make_tuple(std::cbegin(std::get<Is>(tu))...), std::make_tuple(std::cend(std::get<Is>(tu))...)};
    }
 
    public:
    [[nodiscard]] iterator begin() noexcept { return _begin(std::index_sequence_for<Rs...>{}); }
 
    [[nodiscard]] const_iterator cbegin() const noexcept { return _cbegin(std::index_sequence_for<Rs...>{}); }
 
    [[nodiscard]] const_iterator begin() const noexcept { return cbegin(); }
 
    [[nodiscard]] auto end() noexcept { return make_sentinel(std::end(std::get<0>(tu))); }
 
    [[nodiscard]] auto cend() const noexcept { return make_sentinel(std::cend(std::get<0>(tu))); }
 
    [[nodiscard]] auto end() const noexcept { return cend(); }
  };
 
  // Class template argument deduction guide.
  template <typename... Rs> multiplied(Rs &&...) -> multiplied<std::decay_t<Rs>...>;
 
  template <typename... Rs> [[nodiscard]] itertools::multiplied<Rs...> product(Rs &&...rgs) { return {std::forward<Rs>(rgs)...}; }
 
  namespace detail {
 
    // Helper function to create a product range from a container of ranges.
    template <typename C, size_t... Is> [[gnu::always_inline]] [[nodiscard]] auto make_product_impl(C &cont, std::index_sequence<Is...>) {
      return product(cont[Is]...);
    }
 
  } // namespace detail
 
  template <typename R, size_t N> [[nodiscard]] auto make_product(std::array<R, N> &arr) {
    return detail::make_product_impl(arr, std::make_index_sequence<N>{});
  }
 
  template <typename R, size_t N> [[nodiscard]] auto make_product(std::array<R, N> const &arr) {
    return detail::make_product_impl(arr, std::make_index_sequence<N>{});
  }
 
} // namespace itertools
 
#endif // _ITERTOOLS_PRODUCT_HPP