jcs.org
Serenity Operating System
1/*
2 * Copyright (c) 2021, kleines Filmröllchen <filmroellchen@serenityos.org>.
3 *
4 * SPDX-License-Identifier: BSD-2-Clause
5 */
6
7#pragma once
8
9#include <AK/ByteBuffer.h>
10#include <AK/MaybeOwned.h>
11#include <AK/OwnPtr.h>
12#include <AK/Stream.h>
13
14namespace AK {
15
16/// A stream wrapper class that allows you to read arbitrary amounts of bits
17/// in big-endian order from another stream.
18class BigEndianInputBitStream : public Stream {
19public:
20 explicit BigEndianInputBitStream(MaybeOwned<Stream> stream)
21 : m_stream(move(stream))
22 {
23 }
24
25 // ^Stream
26 virtual ErrorOr<Bytes> read_some(Bytes bytes) override
27 {
28 if (m_current_byte.has_value() && is_aligned_to_byte_boundary()) {
29 bytes[0] = m_current_byte.release_value();
30 // FIXME: This accidentally slices off the first byte of the returned span.
31 return m_stream->read_some(bytes.slice(1));
32 }
33 align_to_byte_boundary();
34 return m_stream->read_some(bytes);
35 }
36 virtual ErrorOr<size_t> write_some(ReadonlyBytes bytes) override { return m_stream->write_some(bytes); }
37 virtual bool is_eof() const override { return m_stream->is_eof() && !m_current_byte.has_value(); }
38 virtual bool is_open() const override { return m_stream->is_open(); }
39 virtual void close() override
40 {
41 m_stream->close();
42 align_to_byte_boundary();
43 }
44
45 ErrorOr<bool> read_bit()
46 {
47 return read_bits<bool>(1);
48 }
49 /// Depending on the number of bits to read, the return type can be chosen appropriately.
50 /// This avoids a bunch of static_cast<>'s for the user.
51 // TODO: Support u128, u256 etc. as well: The concepts would be quite complex.
52 template<Unsigned T = u64>
53 ErrorOr<T> read_bits(size_t count)
54 {
55 if constexpr (IsSame<bool, T>) {
56 VERIFY(count == 1);
57 }
58 T result = 0;
59
60 size_t nread = 0;
61 while (nread < count) {
62 if (m_current_byte.has_value()) {
63 if constexpr (!IsSame<bool, T> && !IsSame<u8, T>) {
64 // read as many bytes as possible directly
65 if (((count - nread) >= 8) && is_aligned_to_byte_boundary()) {
66 // shift existing data over
67 result <<= 8;
68 result |= m_current_byte.value();
69 nread += 8;
70 m_current_byte.clear();
71 } else {
72 auto const bit = (m_current_byte.value() >> (7 - m_bit_offset)) & 1;
73 result <<= 1;
74 result |= bit;
75 ++nread;
76 if (m_bit_offset++ == 7)
77 m_current_byte.clear();
78 }
79 } else {
80 // Always take this branch for booleans or u8: there's no purpose in reading more than a single bit
81 auto const bit = (m_current_byte.value() >> (7 - m_bit_offset)) & 1;
82 if constexpr (IsSame<bool, T>)
83 result = bit;
84 else {
85 result <<= 1;
86 result |= bit;
87 }
88 ++nread;
89 if (m_bit_offset++ == 7)
90 m_current_byte.clear();
91 }
92 } else {
93 m_current_byte = TRY(m_stream->read_value<u8>());
94 m_bit_offset = 0;
95 }
96 }
97
98 return result;
99 }
100
101 /// Discards any sub-byte stream positioning the input stream may be keeping track of.
102 /// Non-bitwise reads will implicitly call this.
103 void align_to_byte_boundary()
104 {
105 m_current_byte.clear();
106 m_bit_offset = 0;
107 }
108
109 /// Whether we are (accidentally or intentionally) at a byte boundary right now.
110 ALWAYS_INLINE bool is_aligned_to_byte_boundary() const { return m_bit_offset == 0; }
111
112private:
113 Optional<u8> m_current_byte;
114 size_t m_bit_offset { 0 };
115 MaybeOwned<Stream> m_stream;
116};
117
118/// A stream wrapper class that allows you to read arbitrary amounts of bits
119/// in little-endian order from another stream.
120class LittleEndianInputBitStream : public Stream {
121public:
122 explicit LittleEndianInputBitStream(MaybeOwned<Stream> stream)
123 : m_stream(move(stream))
124 {
125 }
126
127 // ^Stream
128 virtual ErrorOr<Bytes> read_some(Bytes bytes) override
129 {
130 if (m_current_byte.has_value() && is_aligned_to_byte_boundary()) {
131 bytes[0] = m_current_byte.release_value();
132 // FIXME: This accidentally slices off the first byte of the returned span.
133 return m_stream->read_some(bytes.slice(1));
134 }
135 align_to_byte_boundary();
136 return m_stream->read_some(bytes);
137 }
138 virtual ErrorOr<size_t> write_some(ReadonlyBytes bytes) override { return m_stream->write_some(bytes); }
139 virtual bool is_eof() const override { return m_stream->is_eof() && !m_current_byte.has_value(); }
140 virtual bool is_open() const override { return m_stream->is_open(); }
141 virtual void close() override
142 {
143 m_stream->close();
144 align_to_byte_boundary();
145 }
146
147 ErrorOr<bool> read_bit()
148 {
149 return read_bits<bool>(1);
150 }
151 /// Depending on the number of bits to read, the return type can be chosen appropriately.
152 /// This avoids a bunch of static_cast<>'s for the user.
153 // TODO: Support u128, u256 etc. as well: The concepts would be quite complex.
154 template<Unsigned T = u64>
155 ErrorOr<T> read_bits(size_t count)
156 {
157 if constexpr (IsSame<bool, T>) {
158 VERIFY(count == 1);
159 }
160 T result = 0;
161
162 size_t nread = 0;
163 while (nread < count) {
164 if (m_current_byte.has_value()) {
165 if constexpr (!IsSame<bool, T> && !IsSame<u8, T>) {
166 // read as many bytes as possible directly
167 if (((count - nread) >= 8) && is_aligned_to_byte_boundary()) {
168 // shift existing data over
169 result |= (m_current_byte.value() << nread);
170 nread += 8;
171 m_current_byte.clear();
172 } else {
173 auto const bit = (m_current_byte.value() >> m_bit_offset) & 1;
174 result |= (bit << nread);
175 ++nread;
176 if (m_bit_offset++ == 7)
177 m_current_byte.clear();
178 }
179 } else {
180 // Always take this branch for booleans or u8: there's no purpose in reading more than a single bit
181 auto const bit = (m_current_byte.value() >> m_bit_offset) & 1;
182 if constexpr (IsSame<bool, T>)
183 result = bit;
184 else
185 result |= (bit << nread);
186 ++nread;
187 if (m_bit_offset++ == 7)
188 m_current_byte.clear();
189 }
190 } else {
191 m_current_byte = TRY(m_stream->read_value<u8>());
192 m_bit_offset = 0;
193 }
194 }
195
196 return result;
197 }
198
199 /// Discards any sub-byte stream positioning the input stream may be keeping track of.
200 /// Non-bitwise reads will implicitly call this.
201 u8 align_to_byte_boundary()
202 {
203 u8 remaining_bits = m_current_byte.value_or(0) >> m_bit_offset;
204 m_current_byte.clear();
205 m_bit_offset = 0;
206 return remaining_bits;
207 }
208
209 /// Whether we are (accidentally or intentionally) at a byte boundary right now.
210 ALWAYS_INLINE bool is_aligned_to_byte_boundary() const { return m_bit_offset == 0; }
211
212private:
213 Optional<u8> m_current_byte;
214 size_t m_bit_offset { 0 };
215 MaybeOwned<Stream> m_stream;
216};
217
218/// A stream wrapper class that allows you to write arbitrary amounts of bits
219/// in big-endian order to another stream.
220class BigEndianOutputBitStream : public Stream {
221public:
222 explicit BigEndianOutputBitStream(MaybeOwned<Stream> stream)
223 : m_stream(move(stream))
224 {
225 }
226
227 virtual ErrorOr<Bytes> read_some(Bytes) override
228 {
229 return Error::from_errno(EBADF);
230 }
231
232 virtual ErrorOr<size_t> write_some(ReadonlyBytes bytes) override
233 {
234 VERIFY(m_bit_offset == 0);
235 return m_stream->write_some(bytes);
236 }
237
238 template<Unsigned T>
239 ErrorOr<void> write_bits(T value, size_t bit_count)
240 {
241 VERIFY(m_bit_offset <= 7);
242
243 while (bit_count > 0) {
244 u8 next_bit = (value >> (bit_count - 1)) & 1;
245 bit_count--;
246
247 m_current_byte <<= 1;
248 m_current_byte |= next_bit;
249 m_bit_offset++;
250
251 if (m_bit_offset > 7) {
252 TRY(m_stream->write_value(m_current_byte));
253 m_bit_offset = 0;
254 m_current_byte = 0;
255 }
256 }
257
258 return {};
259 }
260
261 virtual bool is_eof() const override
262 {
263 return true;
264 }
265
266 virtual bool is_open() const override
267 {
268 return m_stream->is_open();
269 }
270
271 virtual void close() override
272 {
273 }
274
275 size_t bit_offset() const
276 {
277 return m_bit_offset;
278 }
279
280 ErrorOr<void> align_to_byte_boundary()
281 {
282 if (m_bit_offset == 0)
283 return {};
284
285 TRY(write_bits(0u, 8 - m_bit_offset));
286 VERIFY(m_bit_offset == 0);
287 return {};
288 }
289
290private:
291 MaybeOwned<Stream> m_stream;
292 u8 m_current_byte { 0 };
293 size_t m_bit_offset { 0 };
294};
295
296/// A stream wrapper class that allows you to write arbitrary amounts of bits
297/// in little-endian order to another stream.
298class LittleEndianOutputBitStream : public Stream {
299public:
300 explicit LittleEndianOutputBitStream(MaybeOwned<Stream> stream)
301 : m_stream(move(stream))
302 {
303 }
304
305 virtual ErrorOr<Bytes> read_some(Bytes) override
306 {
307 return Error::from_errno(EBADF);
308 }
309
310 virtual ErrorOr<size_t> write_some(ReadonlyBytes bytes) override
311 {
312 VERIFY(m_bit_offset == 0);
313 return m_stream->write_some(bytes);
314 }
315
316 template<Unsigned T>
317 ErrorOr<void> write_bits(T value, size_t bit_count)
318 {
319 VERIFY(m_bit_offset <= 7);
320
321 size_t input_offset = 0;
322 while (input_offset < bit_count) {
323 u8 next_bit = (value >> input_offset) & 1;
324 input_offset++;
325
326 m_current_byte |= next_bit << m_bit_offset;
327 m_bit_offset++;
328
329 if (m_bit_offset > 7) {
330 TRY(m_stream->write_value(m_current_byte));
331 m_bit_offset = 0;
332 m_current_byte = 0;
333 }
334 }
335
336 return {};
337 }
338
339 virtual bool is_eof() const override
340 {
341 return true;
342 }
343
344 virtual bool is_open() const override
345 {
346 return m_stream->is_open();
347 }
348
349 virtual void close() override
350 {
351 }
352
353 size_t bit_offset() const
354 {
355 return m_bit_offset;
356 }
357
358 ErrorOr<void> align_to_byte_boundary()
359 {
360 if (m_bit_offset == 0)
361 return {};
362
363 TRY(write_bits(0u, 8 - m_bit_offset));
364 VERIFY(m_bit_offset == 0);
365 return {};
366 }
367
368private:
369 MaybeOwned<Stream> m_stream;
370 u8 m_current_byte { 0 };
371 size_t m_bit_offset { 0 };
372};
373
374}