zip.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_ZIP_HPP
#define _ITERTOOLS_ZIP_HPP
 
#include "./iterator_facade.hpp"
#include "./sentinel.hpp"
 
#include <cstddef>
#include <iterator>
#include <tuple>
#include <utility>
 
namespace itertools {
 
  template <typename... Iters>
  struct zip_iter : iterator_facade<zip_iter<Iters...>, std::tuple<typename std::iterator_traits<Iters>::value_type...>> {
    std::tuple<Iters...> its;
 
    zip_iter() = default;
 
    zip_iter(std::tuple<Iters...> its) : its(std::move(its)) {}
 
    private:
    // Helper function which increments all original iterators.
    template <size_t... Is> [[gnu::always_inline]] void increment_all(std::index_sequence<Is...>) { ((void)(++std::get<Is>(its)), ...); }
 
    public:
    void increment() { increment_all(std::index_sequence_for<Iters...>{}); }
 
    [[nodiscard]] bool operator==(zip_iter const &other) const { return its == other.its; }
 
    template <typename SentinelIter> [[nodiscard]] bool operator==(sentinel_t<SentinelIter> const &s) const {
      return [&]<size_t... Is>(std::index_sequence<Is...>) {
        return ((std::get<Is>(its) == std::get<Is>(s.it)) || ...);
      }(std::index_sequence_for<Iters...>{});
    }
 
    private:
    // Helper function to dereference all original iterators.
    template <size_t... Is> [[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 zipped {
    std::tuple<Rs...> tu;
 
    using seq_t = std::index_sequence_for<Rs...>;
 
    using iterator = zip_iter<decltype(std::begin(std::declval<Rs &>()))...>;
 
    using const_iterator = zip_iter<decltype(std::cbegin(std::declval<Rs &>()))...>;
 
    template <typename... Us> zipped(Us &&...rgs) : tu{std::forward<Us>(rgs)...} {}
 
    [[nodiscard]] bool operator==(zipped const &) const = default;
 
    private:
    // Helper function that applies a given callable to each range in the stored tuple and returns a new tuple with the results.
    template <typename F, size_t... Is> [[gnu::always_inline]] auto tuple_map(F &&f, std::index_sequence<Is...>) {
      return std::make_tuple(std::forward<F>(f)(std::get<Is>(tu))...);
    }
 
    // Const overload of tuple_map(F &&, std::index_sequence<Is...>).
    template <typename F, size_t... Is> [[gnu::always_inline]] auto tuple_map(F &&f, std::index_sequence<Is...>) const {
      return std::make_tuple(std::forward<F>(f)(std::get<Is>(tu))...);
    }
 
    public:
    [[nodiscard]] iterator begin() noexcept {
      return tuple_map([](auto &&rg) { return std::begin(rg); }, seq_t{});
    }
 
    [[nodiscard]] const_iterator cbegin() const noexcept {
      return tuple_map([](auto &&rg) { return std::cbegin(rg); }, seq_t{});
    }
 
    [[nodiscard]] const_iterator begin() const noexcept { return cbegin(); }
 
    [[nodiscard]] auto end() noexcept {
      return make_sentinel(tuple_map([](auto &&rg) { return std::end(rg); }, seq_t{}));
    }
 
    [[nodiscard]] auto cend() const noexcept {
      return make_sentinel(tuple_map([](auto &&rg) { return std::cend(rg); }, seq_t{}));
    }
 
    [[nodiscard]] auto end() const noexcept { return cend(); }
  };
 
  template <typename... Rs> [[nodiscard]] zipped<Rs...> zip(Rs &&...rgs) { return {std::forward<Rs>(rgs)...}; }
 
} // namespace itertools
 
#endif // _ITERTOOLS_ZIP_HPP