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