1 | #pragma once |
2 | |
3 | #include <vector> |
4 | |
5 | #include <IO/WriteBuffer.h> |
6 | |
7 | |
8 | namespace DB |
9 | { |
10 | |
11 | namespace 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 | */ |
23 | template <typename VectorType> |
24 | class WriteBufferFromVector : public WriteBuffer |
25 | { |
26 | private: |
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 | |
44 | public: |
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 | |