1#pragma once
2
3#include <vector>
4
5#include <IO/WriteBuffer.h>
6
7
8namespace DB
9{
10
11namespace ErrorCodes
12{
13 extern const int CANNOT_WRITE_AFTER_END_OF_BUFFER;
14}
15
16/** Writes data to existing std::vector or similar type. When not enough space, it doubles vector size.
17 *
18 * In destructor, vector is cutted to the size of written data.
19 * You can call to 'finish' to resize earlier.
20 *
21 * The vector should live until this object is destroyed or until the 'finish' method is called.
22 */
23template <typename VectorType>
24class WriteBufferFromVector : public WriteBuffer
25{
26private:
27 VectorType & vector;
28 bool is_finished = false;
29
30 static constexpr size_t initial_size = 32;
31 static constexpr size_t size_multiplier = 2;
32
33 void nextImpl() override
34 {
35 if (is_finished)
36 throw Exception("WriteBufferFromVector is finished", ErrorCodes::CANNOT_WRITE_AFTER_END_OF_BUFFER);
37
38 size_t old_size = vector.size();
39 vector.resize(old_size * size_multiplier);
40 internal_buffer = Buffer(reinterpret_cast<Position>(vector.data() + old_size), reinterpret_cast<Position>(vector.data() + vector.size()));
41 working_buffer = internal_buffer;
42 }
43
44public:
45 WriteBufferFromVector(VectorType & vector_)
46 : WriteBuffer(reinterpret_cast<Position>(vector_.data()), vector_.size()), vector(vector_)
47 {
48 if (vector.empty())
49 {
50 vector.resize(initial_size);
51 set(reinterpret_cast<Position>(vector.data()), vector.size());
52 }
53 }
54
55 /// Append to vector instead of rewrite.
56 struct AppendModeTag {};
57 WriteBufferFromVector(VectorType & vector_, AppendModeTag)
58 : WriteBuffer(nullptr, 0), vector(vector_)
59 {
60 size_t old_size = vector.size();
61 size_t size = (old_size < initial_size) ? initial_size
62 : ((old_size < vector.capacity()) ? vector.capacity()
63 : vector.capacity() * size_multiplier);
64 vector.resize(size);
65 set(reinterpret_cast<Position>(vector.data() + old_size), (size - old_size) * sizeof(typename VectorType::value_type));
66 }
67
68 void finish()
69 {
70 if (is_finished)
71 return;
72 is_finished = true;
73 vector.resize(
74 ((position() - reinterpret_cast<Position>(vector.data()))
75 + sizeof(typename VectorType::value_type) - 1) /// Align up.
76 / sizeof(typename VectorType::value_type));
77
78 /// Prevent further writes.
79 set(nullptr, 0);
80 }
81
82 bool isFinished() const { return is_finished; }
83
84 void restart()
85 {
86 set(reinterpret_cast<Position>(vector.data()), vector.size());
87 is_finished = false;
88 }
89
90 ~WriteBufferFromVector() override
91 {
92 try
93 {
94 finish();
95 }
96 catch (...)
97 {
98 tryLogCurrentException(__PRETTY_FUNCTION__);
99 }
100 }
101};
102
103}
104