1 | #pragma once |
2 | |
3 | #include <Common/Arena.h> |
4 | #include <common/StringRef.h> |
5 | #include <IO/WriteBuffer.h> |
6 | |
7 | |
8 | namespace DB |
9 | { |
10 | |
11 | /** Writes data contiguously into Arena. |
12 | * As it will be located in contiguous memory segment, it can be read back with ReadBufferFromMemory. |
13 | * |
14 | * While using this object, no other allocations in arena are possible. |
15 | */ |
16 | class WriteBufferFromArena : public WriteBuffer |
17 | { |
18 | private: |
19 | Arena & arena; |
20 | const char *& begin; |
21 | |
22 | void nextImpl() override |
23 | { |
24 | /// Allocate more memory. At least same size as used before (this gives 2x growth ratio), |
25 | /// and at most grab all remaining size in current chunk of arena. |
26 | size_t continuation_size = std::max(count(), arena.remainingSpaceInCurrentChunk()); |
27 | |
28 | /// allocContinue method will possibly move memory region to new place and modify "begin" pointer. |
29 | |
30 | char * continuation = arena.allocContinue(continuation_size, begin); |
31 | char * end = continuation + continuation_size; |
32 | |
33 | /// internal buffer points to whole memory segment and working buffer - to free space for writing. |
34 | internalBuffer() = Buffer(const_cast<char *>(begin), end); |
35 | buffer() = Buffer(continuation, end); |
36 | } |
37 | |
38 | public: |
39 | /// begin_ - start of previously used contiguous memory segment or nullptr (see Arena::allocContinue method). |
40 | WriteBufferFromArena(Arena & arena_, const char *& begin_) |
41 | : WriteBuffer(nullptr, 0), arena(arena_), begin(begin_) |
42 | { |
43 | nextImpl(); |
44 | pos = working_buffer.begin(); |
45 | } |
46 | |
47 | StringRef finish() |
48 | { |
49 | /// Return over-allocated memory back into arena. |
50 | arena.rollback(buffer().end() - position()); |
51 | /// Reference to written data. |
52 | return { position() - count(), count() }; |
53 | } |
54 | }; |
55 | |
56 | } |
57 | |
58 | |