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
12namespace DB
13{
14
15namespace 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 */
27class WriteBuffer : public BufferBase
28{
29public:
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
96private:
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
104using WriteBufferPtr = std::shared_ptr<WriteBuffer>;
105
106
107}
108