1 | #pragma once |
2 | |
3 | #include <algorithm> |
4 | #include <cstring> |
5 | #include <memory> |
6 | #include <iostream> |
7 | |
8 | #include <Common/Exception.h> |
9 | #include <IO/BufferBase.h> |
10 | |
11 | |
12 | namespace DB |
13 | { |
14 | |
15 | namespace ErrorCodes |
16 | { |
17 | extern const int CANNOT_WRITE_AFTER_END_OF_BUFFER; |
18 | } |
19 | |
20 | |
21 | /** A simple abstract class for buffered data writing (char sequences) somewhere. |
22 | * Unlike std::ostream, it provides access to the internal buffer, |
23 | * and also allows you to manually manage the position inside the buffer. |
24 | * |
25 | * The successors must implement the nextImpl() method. |
26 | */ |
27 | class WriteBuffer : public BufferBase |
28 | { |
29 | public: |
30 | WriteBuffer(Position ptr, size_t size) : BufferBase(ptr, size, 0) {} |
31 | void set(Position ptr, size_t size) { BufferBase::set(ptr, size, 0); } |
32 | |
33 | /** write the data in the buffer (from the beginning of the buffer to the current position); |
34 | * set the position to the beginning; throw an exception, if something is wrong |
35 | */ |
36 | inline void next() |
37 | { |
38 | if (!offset() && available()) |
39 | return; |
40 | bytes += offset(); |
41 | |
42 | try |
43 | { |
44 | nextImpl(); |
45 | } |
46 | catch (...) |
47 | { |
48 | /** If the nextImpl() call was unsuccessful, move the cursor to the beginning, |
49 | * so that later (for example, when the stack was expanded) there was no second attempt to write data. |
50 | */ |
51 | pos = working_buffer.begin(); |
52 | throw; |
53 | } |
54 | |
55 | pos = working_buffer.begin(); |
56 | } |
57 | |
58 | /** it is desirable in the successors to place the next() call in the destructor, |
59 | * so that the last data is written |
60 | */ |
61 | virtual ~WriteBuffer() {} |
62 | |
63 | |
64 | inline void nextIfAtEnd() |
65 | { |
66 | if (!hasPendingData()) |
67 | next(); |
68 | } |
69 | |
70 | |
71 | void write(const char * from, size_t n) |
72 | { |
73 | size_t bytes_copied = 0; |
74 | |
75 | while (bytes_copied < n) |
76 | { |
77 | nextIfAtEnd(); |
78 | size_t bytes_to_copy = std::min(static_cast<size_t>(working_buffer.end() - pos), n - bytes_copied); |
79 | memcpy(pos, from + bytes_copied, bytes_to_copy); |
80 | pos += bytes_to_copy; |
81 | bytes_copied += bytes_to_copy; |
82 | } |
83 | } |
84 | |
85 | |
86 | inline void write(char x) |
87 | { |
88 | nextIfAtEnd(); |
89 | *pos = x; |
90 | ++pos; |
91 | } |
92 | |
93 | virtual void sync() {} |
94 | virtual void finalize() {} |
95 | |
96 | private: |
97 | /** Write the data in the buffer (from the beginning of the buffer to the current position). |
98 | * Throw an exception if something is wrong. |
99 | */ |
100 | virtual void nextImpl() { throw Exception("Cannot write after end of buffer." , ErrorCodes::CANNOT_WRITE_AFTER_END_OF_BUFFER); } |
101 | }; |
102 | |
103 | |
104 | using WriteBufferPtr = std::shared_ptr<WriteBuffer>; |
105 | |
106 | |
107 | } |
108 | |