1#pragma once
2
3#include <Common/Arena.h>
4#include <common/StringRef.h>
5#include <IO/WriteBuffer.h>
6
7
8namespace 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 */
16class WriteBufferFromArena : public WriteBuffer
17{
18private:
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
38public:
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