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>
#include <vector>
 
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 an 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 Iter>
  struct prod_iter_vec : iterator_facade<prod_iter_vec<Iter>, std::vector<typename std::iterator_traits<Iter>::value_type>> {
    using value_type = typename std::iterator_traits<Iter>::value_type;

    std::vector<Iter> its_begin;

    std::vector<Iter> its_end;

    std::vector<Iter> its;

    bool done = false;

    prod_iter_vec() = default;

    prod_iter_vec(std::vector<Iter> its_begin, std::vector<Iter> its_end)
       : its_begin(std::move(its_begin)), its_end(std::move(its_end)), its(this->its_begin) {}

    void increment() {
      if (its.empty()) {
        done = true;
        return;
      }
      for (std::size_t i = its.size(); i-- > 1;) {
        ++its[i];
        if (its[i] != its_end[i]) return;
        its[i] = its_begin[i];
      }
      ++its[0];
    }

    [[nodiscard]] bool operator==(prod_iter_vec const &other) const { return its == other.its && done == other.done; }

    template <typename SentinelIter> [[nodiscard]] bool operator==(sentinel_t<SentinelIter> const &s) const {
      if (its.empty()) return done;
      return s.it == its[0];
    }

    [[nodiscard]] std::vector<value_type> dereference() const {
      std::vector<value_type> result;
      result.reserve(its.size());
      for (auto const &it : its) result.push_back(*it);
      return result;
    }
  };

  template <typename R> struct multiplied_vec {
    std::vector<R> tu;

    using iterator = prod_iter_vec<decltype(std::begin(std::declval<R &>()))>;

    using const_iterator = prod_iter_vec<decltype(std::cbegin(std::declval<R &>()))>;

    explicit multiplied_vec(std::vector<R> rgs) : tu(std::move(rgs)) {}

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

    [[nodiscard]] iterator begin() noexcept {
      std::vector<decltype(std::begin(std::declval<R &>()))> b, e;
      b.reserve(tu.size());
      e.reserve(tu.size());
      for (auto &r : tu) {
        b.push_back(std::begin(r));
        e.push_back(std::end(r));
      }
      return iterator{std::move(b), std::move(e)};
    }

    [[nodiscard]] const_iterator cbegin() const noexcept {
      std::vector<decltype(std::cbegin(std::declval<R const &>()))> b, e;
      b.reserve(tu.size());
      e.reserve(tu.size());
      for (auto const &r : tu) {
        b.push_back(std::cbegin(r));
        e.push_back(std::cend(r));
      }
      return const_iterator{std::move(b), std::move(e)};
    }

    [[nodiscard]] const_iterator begin() const noexcept { return cbegin(); }

    [[nodiscard]] auto end() noexcept {
      using SentinelIter = decltype(std::end(std::declval<R &>()));
      if (tu.empty()) return make_sentinel(SentinelIter{});
      return make_sentinel(std::end(tu[0]));
    }

    [[nodiscard]] auto cend() const noexcept {
      using SentinelIter = decltype(std::cend(std::declval<R const &>()));
      if (tu.empty()) return make_sentinel(SentinelIter{});
      return make_sentinel(std::cend(tu[0]));
    }

    [[nodiscard]] auto end() const noexcept { return cend(); }
  };


  template <typename... Rs> [[nodiscard]] itertools::multiplied<Rs...> product(Rs &&...rgs) { return {std::forward<Rs>(rgs)...}; }

  template <typename R> [[nodiscard]] itertools::multiplied_vec<R> product_vec(std::vector<R> rgs) { return multiplied_vec<R>{std::move(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