aboutsummaryrefslogtreecommitdiffstats
path: root/library/cpp/iterator/enumerate.h
blob: 4b6a43dad7872caea2ab57a83c0d3333a9e33372 (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
#pragma once

#include <util/generic/store_policy.h>

#include <limits>
#include <tuple>


namespace NPrivate {

    template <typename TContainer, typename TSize>
    struct TEnumerator {
    private:
        using TStorage = TAutoEmbedOrPtrPolicy<TContainer>;
        using TValue = std::tuple<const TSize, decltype(*std::begin(std::declval<TContainer&>()))>;
        using TIteratorState = decltype(std::begin(std::declval<TContainer&>()));
        using TSentinelState = decltype(std::end(std::declval<TContainer&>()));

        static constexpr bool TrivialSentinel = std::is_same_v<TIteratorState, TSentinelState>;

        struct TIterator;
        struct TSentinelCandidate {
            TSentinelState Iterator_;
        };
        using TSentinel = std::conditional_t<TrivialSentinel, TIterator, TSentinelCandidate>;

        struct TIterator {
            using difference_type = std::ptrdiff_t;
            using value_type = TValue;
            using pointer = void;
            using reference = value_type;
            using iterator_category = std::input_iterator_tag;

            reference operator*() const {
                return {Index_, *Iterator_};
            }

            TIterator& operator++() {
                ++Index_;
                ++Iterator_;
                return *this;
            }

            TIterator operator++(int) {
                TIterator result = *this;
                ++(*this);
                return result;
            }

            bool operator!=(const TSentinel& other) const {
                return Iterator_ != other.Iterator_;
            }

            bool operator==(const TSentinel& other) const {
                return Iterator_ == other.Iterator_;
            }

            TSize Index_;
            TIteratorState Iterator_;
        };

    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::reference;

        TIterator begin() const {
            return {0, std::begin(*Storage_.Ptr())};
        }

        TSentinel end() const {
            if constexpr (TrivialSentinel) {
                return TIterator{std::numeric_limits<TSize>::max(), std::end(*Storage_.Ptr())};
            } else {
                return TSentinel{std::end(*Storage_.Ptr())};
            }
        }

        mutable TStorage Storage_;
    };

}

//! Usage: for (auto [i, x] : Enumerate(container)) {...}
template <typename TContainerOrRef>
auto Enumerate(TContainerOrRef&& container) {
    return NPrivate::TEnumerator<TContainerOrRef, std::size_t>{std::forward<TContainerOrRef>(container)};
}

//! Usage: for (auto [i, x] : SEnumerate(container)) {...}
// The index is signed for codebases that prefer signed numerics (such as YTsaurus).
template <typename TContainerOrRef>
auto SEnumerate(TContainerOrRef&& container) {
    return NPrivate::TEnumerator<TContainerOrRef, std::ptrdiff_t>{std::forward<TContainerOrRef>(container)};
}