summaryrefslogtreecommitdiffstats
path: root/library/cpp/iterator/zip.h
blob: 6ec0993240d3016d9fd93574fb286e983bee7817 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
#pragma once 
 
#include <util/generic/store_policy.h> 
 
#include <algorithm> 
#include <tuple> 
 
 
namespace NPrivate { 
 
    template <typename TContainer, typename TIteratorCategory = typename std::iterator_traits<decltype(std::begin(std::declval<TContainer>()))>::iterator_category> 
    static constexpr bool HasRandomAccessIterator(int32_t) { 
        return std::is_same_v<TIteratorCategory, std::random_access_iterator_tag>; 
    } 
 
    template <typename TContainer> 
    static constexpr bool HasRandomAccessIterator(uint32_t) { 
        return false; 
    } 
 
    template <typename... TContainers> 
    struct TZipper { 
        template <std::size_t... I> 
        struct TZipperWithIndex { 
        private: 
            using THolders = std::tuple<TAutoEmbedOrPtrPolicy<TContainers>...>; 
            using TValue = std::tuple<decltype(*std::begin(std::declval<TContainers&>()))...>; 
            using TIteratorState = std::tuple<decltype(std::begin(std::declval<TContainers&>()))...>; 
            using TSentinelState = std::tuple<decltype(std::end(std::declval<TContainers&>()))...>; 
 
            static constexpr bool TrivialSentinel = std::is_same_v<TIteratorState, TSentinelState>; 
 
            struct TIterator; 
            struct TSentinelCandidate { 
                TSentinelState Iterators_; 
            }; 
            using TSentinel = std::conditional_t<TrivialSentinel, TIterator, TSentinelCandidate>; 
 
#ifndef _MSC_VER 
            // windows compiler crashes here 
            static constexpr bool LimitByFirstContainer = TrivialSentinel && 
                ((HasRandomAccessIterator<TContainers>(0)) && ...); 
#else 
            static constexpr bool LimitByFirstContainer = false; 
#endif 
 
            struct TIterator { 
                using difference_type = std::ptrdiff_t; 
                using value_type = TValue; 
                using pointer = TValue*; 
                using reference = TValue&; 
                using const_reference = const TValue&;
                using iterator_category = std::input_iterator_tag; 
 
                TValue operator*() { 
                    return {*std::get<I>(Iterators_)...}; 
                } 
                TValue operator*() const { 
                    return {*std::get<I>(Iterators_)...}; 
                } 

                TIterator& operator++() {
                    (++std::get<I>(Iterators_), ...); 
                    return *this;
                } 

                TIterator operator++(int) {
                    return TIterator{TIteratorState{std::get<I>(Iterators_)++...}};
                }

                bool operator!=(const TSentinel& other) const { 
                    if constexpr (LimitByFirstContainer) { 
                        return std::get<0>(Iterators_) != std::get<0>(other.Iterators_); 
                    } else { 
                        // yes, for all correct iterators but end() it is a correct way to compare 
                        return ((std::get<I>(Iterators_) != std::get<I>(other.Iterators_)) && ...); 
                    } 
                } 
                bool operator==(const TSentinel& other) const { 
                    return !(*this != other); 
                } 
 
                TIteratorState Iterators_; 
            }; 
        public: 
            using iterator = TIterator; 
            using const_iterator = TIterator; 
            using value_type = typename TIterator::value_type;
            using reference = typename TIterator::reference;
            using const_reference = typename TIterator::const_reference;
 
            TIterator begin() const { 
                return {TIteratorState{std::begin(*std::get<I>(Holders_).Ptr())...}}; 
            } 
 
            TSentinel end() const { 
                if constexpr (LimitByFirstContainer) { 
                    auto endOfFirst = std::begin(*std::get<0>(Holders_).Ptr()) + std::min({ 
                        std::end(*std::get<I>(Holders_).Ptr()) - std::begin(*std::get<I>(Holders_).Ptr())...}); 
                    TIterator iter{TSentinelState{std::end(*std::get<I>(Holders_).Ptr())...}}; 
                    std::get<0>(iter.Iterators_) = endOfFirst; 
                    return iter; 
                } else { 
                    return {TSentinelState{std::end(*std::get<I>(Holders_).Ptr())...}}; 
                } 
            } 
 
            mutable THolders Holders_; 
        }; 
 
        template <std::size_t... I> 
        static auto Zip(TContainers&&... containers, std::index_sequence<I...>) { 
            return TZipperWithIndex<I...>{{std::forward<TContainers>(containers)...}}; 
        } 
    }; 
 
} 
 
 
//! Acts as pythonic zip, BUT result length is equal to shortest length of input containers 
//! Usage: for (auto [ai, bi, ci] : Zip(a, b, c)) {...} 
template <typename... TContainers> 
auto Zip(TContainers&&... containers) { 
    return ::NPrivate::TZipper<TContainers...>::Zip(
        std::forward<TContainers>(containers)..., 
        std::make_index_sequence<sizeof...(TContainers)>{} 
    ); 
}