From b6529f351179f218d66dcb3c1a2cb40ef5050bf7 Mon Sep 17 00:00:00 2001
From: JaySon <tshent@qq.com>
Date: Sun, 17 Jan 2021 09:33:13 -0600
Subject: [PATCH] Cherry pick bugfixes for PODArray (#1357)

---
 dbms/src/Common/PODArray.h                | 21 ++++---
 dbms/src/Common/tests/gtest_pod_array.cpp | 67 +++++++++++++++++++++++
 2 files changed, 81 insertions(+), 7 deletions(-)
 create mode 100644 dbms/src/Common/tests/gtest_pod_array.cpp

diff --git a/dbms/src/Common/PODArray.h b/dbms/src/Common/PODArray.h
index 4983b9989e7..2e8385a15c4 100644
--- a/dbms/src/Common/PODArray.h
+++ b/dbms/src/Common/PODArray.h
@@ -228,14 +228,20 @@ class PODArrayBase : private boost::noncopyable, private TAllocator    /// empty
         return c_start;
     }
 
-    template <typename ... TAllocatorParams>
+    template <typename... TAllocatorParams>
     void push_back_raw(const char * ptr, TAllocatorParams &&... allocator_params)
     {
-        if (unlikely(c_end >= c_end_of_storage))
-            reserveForNextSize(std::forward<TAllocatorParams>(allocator_params)...);
+        push_back_raw_many(1, ptr, std::forward<TAllocatorParams>(allocator_params)...);
+    }
 
-        memcpy(c_end, ptr, ELEMENT_SIZE);
-        c_end += byte_size(1);
+    template <typename... TAllocatorParams>
+    void push_back_raw_many(size_t number_of_items, const void * ptr, TAllocatorParams &&... allocator_params)
+    {
+        size_t items_byte_size = byte_size(number_of_items);
+        if (unlikely(c_end + items_byte_size > c_end_of_storage))
+            reserve(size() + number_of_items, std::forward<TAllocatorParams>(allocator_params)...);
+        memcpy(c_end, ptr, items_byte_size);
+        c_end += items_byte_size;
     }
 
     void protect()
@@ -434,11 +440,12 @@ class PODArray : public PODArrayBase<sizeof(T), INITIAL_SIZE, TAllocator, pad_ri
     template <typename It1, typename It2>
     void insert(iterator it, It1 from_begin, It2 from_end)
     {
-        insertPrepare(from_begin, from_end);
-
         size_t bytes_to_copy = this->byte_size(from_end - from_begin);
         size_t bytes_to_move = (end() - it) * sizeof(T);
 
+        // This may make `it` invalid, do it after calculating `bytes_to_move`
+        insertPrepare(from_begin, from_end);
+
         if (unlikely(bytes_to_move))
             memcpy(this->c_end + bytes_to_copy - bytes_to_move, this->c_end - bytes_to_move, bytes_to_move);
 
diff --git a/dbms/src/Common/tests/gtest_pod_array.cpp b/dbms/src/Common/tests/gtest_pod_array.cpp
new file mode 100644
index 00000000000..726d7bf25f9
--- /dev/null
+++ b/dbms/src/Common/tests/gtest_pod_array.cpp
@@ -0,0 +1,67 @@
+#include <Common/PODArray.h>
+#include <gtest/gtest.h>
+
+using namespace DB;
+
+TEST(Common, PODArrayInsert)
+{
+    std::string str = "test_string_abacaba";
+    PODArray<char> chars;
+    chars.insert(chars.end(), str.begin(), str.end());
+    EXPECT_EQ(str, std::string(chars.data(), chars.size()));
+
+    std::string insert_in_the_middle = "insert_in_the_middle";
+    auto pos = str.size() / 2;
+    str.insert(str.begin() + pos, insert_in_the_middle.begin(), insert_in_the_middle.end());
+    chars.insert(chars.begin() + pos, insert_in_the_middle.begin(), insert_in_the_middle.end());
+    EXPECT_EQ(str, std::string(chars.data(), chars.size()));
+
+    std::string insert_with_resize;
+    insert_with_resize.reserve(chars.capacity() * 2);
+    char cur_char = 'a';
+    while (insert_with_resize.size() < insert_with_resize.capacity())
+    {
+        insert_with_resize += cur_char;
+        if (cur_char == 'z')
+            cur_char = 'a';
+        else
+            ++cur_char;
+    }
+    str.insert(str.begin(), insert_with_resize.begin(), insert_with_resize.end());
+    chars.insert(chars.begin(), insert_with_resize.begin(), insert_with_resize.end());
+    EXPECT_EQ(str, std::string(chars.data(), chars.size()));
+}
+
+TEST(Common, PODPushBackRawMany)
+{
+    PODArray<char> chars;
+    chars.push_back_raw_many(5, "first");
+    EXPECT_EQ(std::string("first"), std::string(chars.data(), chars.size()));
+    EXPECT_EQ(5UL, chars.size());
+    EXPECT_LE(chars.capacity() - chars.size(), 10UL);
+    chars.push_back_raw_many(10, "0123456789");
+    EXPECT_EQ(15UL, chars.size());
+    EXPECT_EQ(std::string("first0123456789"), std::string(chars.data(), chars.size()));
+}
+
+TEST(Common, PODNoOverallocation)
+{
+    /// Check that PaddedPODArray allocates for smaller number of elements than the power of two due to padding.
+    /// NOTE: It's Ok to change these numbers if you will modify initial size or padding.
+
+    PaddedPODArray<char> chars;
+    std::vector<size_t> capacities;
+
+    size_t prev_capacity = 0;
+    for (size_t i = 0; i < 1000000; ++i)
+    {
+        chars.emplace_back();
+        if (chars.capacity() != prev_capacity)
+        {
+            prev_capacity = chars.capacity();
+            capacities.emplace_back(prev_capacity);
+        }
+    }
+
+    EXPECT_EQ(capacities, (std::vector<size_t>{4065, 8161, 16353, 32737, 65505, 131041, 262113, 524257, 1048545}));
+}