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 | |
17 | namespace 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 | |
35 | class 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 | |
43 | public: |
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 | |
103 | private: |
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 | |
114 | class 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 | |
123 | public: |
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 | |
175 | private: |
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 | |