intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

5 files changed, 379 insertions, 0 deletions

diff --git a/library/cpp/containers/stack_vector/stack_vec.cpp b/library/cpp/containers/stack_vector/stack_vec.cpp
new file mode 100644
index 0000000000..21c0ab3f11
--- /dev/null
+++ b/library/cpp/containers/stack_vector/stack_vec.cpp
@@ -0,0 +1 @@
+#include "stack_vec.h"
diff --git a/library/cpp/containers/stack_vector/stack_vec.h b/library/cpp/containers/stack_vector/stack_vec.h
new file mode 100644
index 0000000000..fcc5d9a2a5
--- /dev/null
+++ b/library/cpp/containers/stack_vector/stack_vec.h
@@ -0,0 +1,212 @@
+#pragma once
+
+#include <util/generic/vector.h>
+#include <util/ysaveload.h>
+
+#include <type_traits>
+
+// A vector preallocated on the stack.
+// After exceeding the preconfigured stack space falls back to the heap.
+// Publicly inherits TVector, but disallows swap (and hence shrink_to_fit, also operator= is reimplemented via copying).
+//
+// Inspired by: http://qt-project.org/doc/qt-4.8/qvarlengtharray.html#details
+
+template <typename T, size_t CountOnStack = 256, bool UseFallbackAlloc = true, class Alloc = std::allocator<T>>
+class TStackVec;
+
+template <typename T, class Alloc = std::allocator<T>>
+using TSmallVec = TStackVec<T, 16, true, Alloc>;
+
+template <typename T, size_t CountOnStack = 256>
+using TStackOnlyVec = TStackVec<T, CountOnStack, false>;
+
+namespace NPrivate {
+    template <class Alloc, class StackAlloc, typename T, typename U>
+    struct TRebind {
+        typedef TReboundAllocator<Alloc, U> other;
+    };
+
+    template <class Alloc, class StackAlloc, typename T>
+    struct TRebind<Alloc, StackAlloc, T, T> {
+        typedef StackAlloc other;
+    };
+
+    template <typename T, size_t CountOnStack, bool UseFallbackAlloc, class Alloc = std::allocator<T>>
+    class TStackBasedAllocator: public Alloc {
+    public:
+        typedef TStackBasedAllocator<T, CountOnStack, UseFallbackAlloc, Alloc> TSelf;
+
+        using typename Alloc::difference_type;
+        using typename Alloc::size_type;
+        using typename Alloc::value_type;
+
+        template <class U>
+        struct rebind: public ::NPrivate::TRebind<Alloc, TSelf, T, U> {
+        };
+
+    public:
+        TStackBasedAllocator() = default;
+
+        template <
+            typename... TArgs,
+            typename = std::enable_if_t<
+                std::is_constructible_v<Alloc, TArgs...>
+            >
+        >
+        TStackBasedAllocator(TArgs&&... args)
+            : Alloc(std::forward<TArgs>(args)...)
+        {}
+
+        T* allocate(size_type n) {
+            if (!IsStorageUsed && CountOnStack >= n) {
+                IsStorageUsed = true;
+                return reinterpret_cast<T*>(&StackBasedStorage[0]);
+            } else {
+                if constexpr (!UseFallbackAlloc) {
+                    Y_FAIL(
+                            "Stack storage overflow. Capacity: %d, requested: %d", (int)CountOnStack, int(n));
+                }
+                return FallbackAllocator().allocate(n);
+            }
+        }
+
+        void deallocate(T* p, size_type n) {
+            if (p >= reinterpret_cast<T*>(&StackBasedStorage[0]) &&
+                    p < reinterpret_cast<T*>(&StackBasedStorage[CountOnStack])) {
+                Y_VERIFY(IsStorageUsed);
+                IsStorageUsed = false;
+            } else {
+                FallbackAllocator().deallocate(p, n);
+            }
+        }
+
+    private:
+        std::aligned_storage_t<sizeof(T), alignof(T)> StackBasedStorage[CountOnStack];
+        bool IsStorageUsed = false;
+
+    private:
+        Alloc& FallbackAllocator() noexcept {
+            return static_cast<Alloc&>(*this);
+        }
+    };
+}
+
+template <typename T, size_t CountOnStack, bool UseFallbackAlloc, class Alloc>
+class TStackVec: public TVector<T, ::NPrivate::TStackBasedAllocator<T, CountOnStack, UseFallbackAlloc, TReboundAllocator<Alloc, T>>> {
+private:
+    using TBase = TVector<T, ::NPrivate::TStackBasedAllocator<T, CountOnStack, UseFallbackAlloc, TReboundAllocator<Alloc, T>>>;
+    using TAllocator = typename TBase::allocator_type;
+
+public:
+    using typename TBase::const_iterator;
+    using typename TBase::const_reverse_iterator;
+    using typename TBase::iterator;
+    using typename TBase::reverse_iterator;
+    using typename TBase::size_type;
+    using typename TBase::value_type;
+
+public:
+    TStackVec(const TAllocator& alloc = TAllocator())
+        : TBase(alloc)
+    {
+        TBase::reserve(CountOnStack);
+    }
+
+    explicit TStackVec(size_type count, const TAllocator& alloc = TAllocator())
+        : TBase(alloc)
+    {
+        if (count <= CountOnStack) {
+            TBase::reserve(CountOnStack);
+        }
+        TBase::resize(count);
+    }
+
+    TStackVec(size_type count, const T& val, const TAllocator& alloc = TAllocator())
+        : TBase(alloc)
+    {
+        if (count <= CountOnStack) {
+            TBase::reserve(CountOnStack);
+        }
+        TBase::assign(count, val);
+    }
+
+    TStackVec(const TStackVec& src)
+        : TStackVec(src.begin(), src.end())
+    {
+    }
+
+    template <class A>
+    TStackVec(const TVector<T, A>& src)
+        : TStackVec(src.begin(), src.end())
+    {
+    }
+
+    TStackVec(std::initializer_list<T> il, const TAllocator& alloc = TAllocator())
+        : TStackVec(il.begin(), il.end(), alloc)
+    {
+    }
+
+    template <class TIter>
+    TStackVec(TIter first, TIter last, const TAllocator& alloc = TAllocator())
+        : TBase(alloc)
+    {
+        // NB(eeight) Since we want to call 'reserve' here, we cannot just delegate to TVector ctor.
+        // The best way to insert values afterwards is to call TVector::insert. However there is a caveat.
+        // In order to call this ctor of TVector, T needs to be just move-constructible. Insert however
+        // requires T to be move-assignable.
+        TBase::reserve(CountOnStack);
+        if constexpr (std::is_move_assignable_v<T>) {
+            // Fast path
+            TBase::insert(TBase::end(), first, last);
+        } else {
+            // Slow path.
+            for (; first != last; ++first) {
+                TBase::push_back(*first);
+            }
+        }
+    }
+
+public:
+    void swap(TStackVec&) = delete;
+    void shrink_to_fit() = delete;
+
+    TStackVec& operator=(const TStackVec& src) {
+        TBase::assign(src.begin(), src.end());
+        return *this;
+    }
+
+    template <class A>
+    TStackVec& operator=(const TVector<T, A>& src) {
+        TBase::assign(src.begin(), src.end());
+        return *this;
+    }
+
+    TStackVec& operator=(std::initializer_list<T> il) {
+        TBase::assign(il.begin(), il.end());
+        return *this;
+    }
+};
+
+template <typename T, size_t CountOnStack, class Alloc>
+class TSerializer<TStackVec<T, CountOnStack, true, Alloc>>: public TVectorSerializer<TStackVec<T, CountOnStack, true, Alloc>> {
+};
+
+template <typename T, size_t CountOnStack, class Alloc>
+class TSerializer<TStackVec<T, CountOnStack, false, Alloc>> {
+public:
+    static void Save(IOutputStream* rh, const TStackVec<T, CountOnStack, false, Alloc>& v) {
+        if constexpr (CountOnStack < 256) {
+            ::Save(rh, (ui8)v.size());
+        } else {
+            ::Save(rh, v.size());
+        }
+        ::SaveArray(rh, v.data(), v.size());
+    }
+
+    static void Load(IInputStream* rh, TStackVec<T, CountOnStack, false, Alloc>& v) {
+        std::conditional_t<CountOnStack < 256, ui8, size_t> size;
+        ::Load(rh, size);
+        v.resize(size);
+        ::LoadPodArray(rh, v.data(), v.size());
+    }
+};
diff --git a/library/cpp/containers/stack_vector/stack_vec_ut.cpp b/library/cpp/containers/stack_vector/stack_vec_ut.cpp
new file mode 100644
index 0000000000..19f9677781
--- /dev/null
+++ b/library/cpp/containers/stack_vector/stack_vec_ut.cpp
@@ -0,0 +1,144 @@
+#include "stack_vec.h"
+
+#include <library/cpp/testing/unittest/registar.h>
+
+namespace {
+    struct TNotCopyAssignable {
+        const int Value;
+    };
+
+    static_assert(std::is_copy_constructible_v<TNotCopyAssignable>);
+    static_assert(!std::is_copy_assignable_v<TNotCopyAssignable>);
+
+    template <class T, size_t JunkPayloadSize>
+    struct TThickAlloc: public std::allocator<T> {
+        template <class U>
+        struct rebind {
+            using other = TThickAlloc<U, JunkPayloadSize>;
+        };
+
+        char Junk[JunkPayloadSize]{sizeof(T)};
+    };
+
+    template <class T>
+    struct TStatefulAlloc: public std::allocator<T> {
+        using TBase = std::allocator<T>;
+
+        template <class U>
+        struct rebind {
+            using other = TStatefulAlloc<U>;
+        };
+
+        TStatefulAlloc(size_t* allocCount)
+            : AllocCount(allocCount)
+        {}
+
+        size_t* AllocCount;
+
+        T* allocate(size_t n)
+        {
+            *AllocCount += 1;
+            return TBase::allocate(n);
+        }
+    };
+}
+
+Y_UNIT_TEST_SUITE(TStackBasedVectorTest) {
+    Y_UNIT_TEST(TestCreateEmpty) {
+        TStackVec<int> ints;
+        UNIT_ASSERT_EQUAL(ints.size(), 0);
+    }
+
+    Y_UNIT_TEST(TestCreateNonEmpty) {
+        TStackVec<int> ints(5);
+        UNIT_ASSERT_EQUAL(ints.size(), 5);
+
+        for (size_t i = 0; i < ints.size(); ++i) {
+            UNIT_ASSERT_EQUAL(ints[i], 0);
+        }
+    }
+
+    Y_UNIT_TEST(TestReallyOnStack) {
+        const TStackVec<int> vec(5);
+
+        UNIT_ASSERT(
+            (const char*)&vec <= (const char*)&vec[0] &&
+            (const char*)&vec[0] <= (const char*)&vec + sizeof(vec)
+        );
+    }
+
+    Y_UNIT_TEST(TestFallback) {
+        TSmallVec<int> ints;
+        for (int i = 0; i < 14; ++i) {
+            ints.push_back(i);
+        }
+
+        for (size_t i = 0; i < ints.size(); ++i) {
+            UNIT_ASSERT_EQUAL(ints[i], (int)i);
+        }
+
+        for (int i = 14; i < 20; ++i) {
+            ints.push_back(i);
+        }
+
+        for (size_t i = 0; i < ints.size(); ++i) {
+            UNIT_ASSERT_EQUAL(ints[i], (int)i);
+        }
+
+        TSmallVec<int> ints2 = ints;
+
+        for (size_t i = 0; i < ints2.size(); ++i) {
+            UNIT_ASSERT_EQUAL(ints2[i], (int)i);
+        }
+
+        TSmallVec<int> ints3;
+        ints3 = ints2;
+
+        for (size_t i = 0; i < ints3.size(); ++i) {
+            UNIT_ASSERT_EQUAL(ints3[i], (int)i);
+        }
+    }
+
+    Y_UNIT_TEST(TestCappedSize) {
+        TStackVec<int, 8, false> ints;
+        ints.push_back(1);
+        ints.push_back(2);
+
+        auto intsCopy = ints;
+        UNIT_ASSERT_VALUES_EQUAL(intsCopy.capacity(), 8);
+
+        for (int i = 2; i != 8; ++i) {
+            intsCopy.push_back(i);
+        }
+        // Just verify that the program did not crash.
+    }
+
+    Y_UNIT_TEST(TestCappedSizeWithNotCopyAssignable) {
+        TStackVec<TNotCopyAssignable, 8, false> values;
+        values.push_back({1});
+        values.push_back({2});
+
+        auto valuesCopy = values;
+        UNIT_ASSERT_VALUES_EQUAL(valuesCopy.capacity(), 8);
+
+        for (int i = 2; i != 8; ++i) {
+            valuesCopy.push_back({i});
+        }
+        // Just verify that the program did not crash.
+    }
+
+    Y_UNIT_TEST(TestCustomAllocSize) {
+        constexpr size_t n = 16384;
+        using TVec = TStackVec<size_t, 1, true, TThickAlloc<size_t, n>>;
+        UNIT_ASSERT_LT(sizeof(TVec), 1.5 * n);
+    }
+
+    Y_UNIT_TEST(TestStatefulAlloc) {
+        size_t count = 0;
+        TStackVec<size_t, 1, true, TStatefulAlloc<size_t>> vec{{ &count }};
+        for (size_t i = 0; i < 5; ++i) {
+            vec.push_back(1);
+        }
+        UNIT_ASSERT_VALUES_EQUAL(count, 3);
+    }
+}
diff --git a/library/cpp/containers/stack_vector/ut/ya.make b/library/cpp/containers/stack_vector/ut/ya.make
new file mode 100644
index 0000000000..1d70496954
--- /dev/null
+++ b/library/cpp/containers/stack_vector/ut/ya.make
@@ -0,0 +1,11 @@
+UNITTEST()
+
+OWNER(ilnurkh)
+
+SRCDIR(library/cpp/containers/stack_vector)
+
+SRCS(
+    stack_vec_ut.cpp
+)
+
+END()
diff --git a/library/cpp/containers/stack_vector/ya.make b/library/cpp/containers/stack_vector/ya.make
new file mode 100644
index 0000000000..cfb63295ec
--- /dev/null
+++ b/library/cpp/containers/stack_vector/ya.make
@@ -0,0 +1,11 @@
+LIBRARY()
+
+OWNER(ilnurkh)
+
+SRCS(
+    stack_vec.cpp
+)
+
+END()
+
+RECURSE_FOR_TESTS(ut)