intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

1 files changed, 128 insertions, 0 deletions

diff --git a/library/cpp/iterator/zip.h b/library/cpp/iterator/zip.h
new file mode 100644
index 0000000000..ac12ed35fe
--- /dev/null
+++ b/library/cpp/iterator/zip.h
@@ -0,0 +1,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)>{}
+    );
+}