1#pragma once
2
3#include <IO/ReadBuffer.h>
4#include <IO/WriteBuffer.h>
5#include <Core/Types.h>
6#include <Common/BitHelpers.h>
7
8#if defined(__OpenBSD__) || defined(__FreeBSD__)
9# include <sys/endian.h>
10#elif defined(__APPLE__)
11# include <libkern/OSByteOrder.h>
12
13# define htobe64(x) OSSwapHostToBigInt64(x)
14# define be64toh(x) OSSwapBigToHostInt64(x)
15#endif
16
17namespace DB
18{
19
20/** Reads data from underlying ReadBuffer bit by bit, max 64 bits at once.
21 *
22 * reads MSB bits first, imagine that you have a data:
23 * 11110000 10101010 00100100 11111110
24 *
25 * Given that r is BitReader created with a ReadBuffer that reads from data above:
26 * r.readBits(3) => 0b111
27 * r.readBit() => 0b1
28 * r.readBits(8) => 0b1010 // 4 leading zero-bits are not shown
29 * r.readBit() => 0b1
30 * r.readBit() => 0b0
31 * r.readBits(15) => 0b10001001001111111
32 * r.readBit() => 0b0
33**/
34
35class BitReader
36{
37 ReadBuffer & buf;
38
39 UInt64 bits_buffer;
40 UInt8 bits_count;
41 static constexpr UInt8 BIT_BUFFER_SIZE = sizeof(bits_buffer) * 8;
42
43public:
44 BitReader(ReadBuffer & buf_)
45 : buf(buf_),
46 bits_buffer(0),
47 bits_count(0)
48 {}
49
50 ~BitReader()
51 {}
52
53 inline UInt64 readBits(UInt8 bits)
54 {
55 UInt64 result = 0;
56 bits = std::min(static_cast<UInt8>(sizeof(result) * 8), bits);
57
58 while (bits != 0)
59 {
60 if (bits_count == 0)
61 {
62 fillBuffer();
63 if (bits_count == 0)
64 {
65 // EOF.
66 break;
67 }
68 }
69
70 const auto to_read = std::min(bits, bits_count);
71
72 const UInt64 v = bits_buffer >> (bits_count - to_read);
73 const UInt64 mask = maskLowBits<UInt64>(to_read);
74 const UInt64 value = v & mask;
75 result |= value;
76
77 // unset bits that were read
78 bits_buffer &= ~(mask << (bits_count - to_read));
79 bits_count -= to_read;
80 bits -= to_read;
81
82 result <<= std::min(bits, BIT_BUFFER_SIZE);
83 }
84
85 return result;
86 }
87
88 inline UInt64 peekBits(UInt8 /*bits*/)
89 {
90 return 0;
91 }
92
93 inline UInt8 readBit()
94 {
95 return static_cast<UInt8>(readBits(1));
96 }
97
98 inline bool eof() const
99 {
100 return bits_count == 0 && buf.eof();
101 }
102
103private:
104 void fillBuffer()
105 {
106 auto read = buf.read(reinterpret_cast<char *>(&bits_buffer), BIT_BUFFER_SIZE / 8);
107 bits_buffer = be64toh(bits_buffer);
108 bits_buffer >>= BIT_BUFFER_SIZE - read * 8;
109
110 bits_count = static_cast<UInt8>(read) * 8;
111 }
112};
113
114class BitWriter
115{
116 WriteBuffer & buf;
117
118 UInt64 bits_buffer;
119 UInt8 bits_count;
120
121 static constexpr UInt8 BIT_BUFFER_SIZE = sizeof(bits_buffer) * 8;
122
123public:
124 BitWriter(WriteBuffer & buf_)
125 : buf(buf_),
126 bits_buffer(0),
127 bits_count(0)
128 {}
129
130 ~BitWriter()
131 {
132 flush();
133 }
134
135 inline void writeBits(UInt8 bits, UInt64 value)
136 {
137 bits = std::min(static_cast<UInt8>(sizeof(value) * 8), bits);
138
139 while (bits > 0)
140 {
141 auto v = value;
142 auto to_write = bits;
143
144 const UInt8 capacity = BIT_BUFFER_SIZE - bits_count;
145 if (capacity < bits)
146 {
147 v >>= bits - capacity;
148 to_write = capacity;
149 }
150
151 const UInt64 mask = maskLowBits<UInt64>(to_write);
152 v &= mask;
153
154 bits_buffer <<= to_write;
155 bits_buffer |= v;
156 bits_count += to_write;
157
158 if (bits_count < BIT_BUFFER_SIZE)
159 break;
160
161 doFlush();
162 bits -= to_write;
163 }
164 }
165
166 inline void flush()
167 {
168 if (bits_count != 0)
169 {
170 bits_buffer <<= (BIT_BUFFER_SIZE - bits_count);
171 doFlush();
172 }
173 }
174
175private:
176 void doFlush()
177 {
178 bits_buffer = htobe64(bits_buffer);
179 buf.write(reinterpret_cast<const char *>(&bits_buffer), (bits_count + 7) / 8);
180
181 bits_count = 0;
182 bits_buffer = 0;
183 }
184};
185
186}
187