summaryrefslogtreecommitdiff
path: root/Userland/Libraries/LibCore/InputBitStream.h
blob: 494963d2a863e7b1122a0c17276ad2cbe524b2a0 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
/*
 * Copyright (c) 2021, kleines Filmröllchen <filmroellchen@serenityos.org>.
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#pragma once

#include <AK/ByteBuffer.h>
#include <AK/Concepts.h>
#include <AK/Error.h>
#include <AK/NonnullOwnPtr.h>
#include <AK/NonnullRefPtr.h>
#include <AK/OwnPtr.h>
#include <AK/Span.h>
#include <AK/StdLibExtraDetails.h>
#include <AK/Types.h>
#include <LibCore/Stream.h>

namespace Core::Stream {

/// A stream wrapper class that allows you to read arbitrary amounts of bits
/// in big-endian order from another stream.
/// Note that this stream does not own its underlying stream, it merely takes a reference.
class BigEndianInputBitStream : public Stream {
public:
    static ErrorOr<NonnullOwnPtr<BigEndianInputBitStream>> construct(Stream& stream)
    {
        return adopt_nonnull_own_or_enomem<BigEndianInputBitStream>(new BigEndianInputBitStream(stream));
    }

    // ^Stream
    virtual ErrorOr<Bytes> read(Bytes bytes) override
    {
        if (m_current_byte.has_value() && is_aligned_to_byte_boundary()) {
            bytes[0] = m_current_byte.release_value();
            return m_stream.read(bytes.slice(1));
        }
        align_to_byte_boundary();
        return m_stream.read(bytes);
    }
    virtual ErrorOr<size_t> write(ReadonlyBytes bytes) override { return m_stream.write(bytes); }
    virtual bool write_or_error(ReadonlyBytes bytes) override { return m_stream.write_or_error(bytes); }
    virtual bool is_eof() const override { return m_stream.is_eof() && !m_current_byte.has_value(); }
    virtual bool is_open() const override { return m_stream.is_open(); }
    virtual void close() override
    {
        m_stream.close();
        align_to_byte_boundary();
    }

    ErrorOr<bool> read_bit()
    {
        return read_bits<bool>(1);
    }
    /// Depending on the number of bits to read, the return type can be chosen appropriately.
    /// This avoids a bunch of static_cast<>'s for the user.
    // TODO: Support u128, u256 etc. as well: The concepts would be quite complex.
    template<Unsigned T = u64>
    ErrorOr<T> read_bits(size_t count)
    {
        if constexpr (IsSame<bool, T>) {
            VERIFY(count == 1);
        }
        T result = 0;

        size_t nread = 0;
        while (nread < count) {
            if (m_current_byte.has_value()) {
                if constexpr (!IsSame<bool, T> && !IsSame<u8, T>) {
                    // read as many bytes as possible directly
                    if (((count - nread) >= 8) && is_aligned_to_byte_boundary()) {
                        // shift existing data over
                        result <<= 8;
                        result |= m_current_byte.value();
                        nread += 8;
                        m_current_byte.clear();
                    } else {
                        auto const bit = (m_current_byte.value() >> (7 - m_bit_offset)) & 1;
                        result <<= 1;
                        result |= bit;
                        ++nread;
                        if (m_bit_offset++ == 7)
                            m_current_byte.clear();
                    }
                } else {
                    // Always take this branch for booleans or u8: there's no purpose in reading more than a single bit
                    auto const bit = (m_current_byte.value() >> (7 - m_bit_offset)) & 1;
                    if constexpr (IsSame<bool, T>)
                        result = bit;
                    else {
                        result <<= 1;
                        result |= bit;
                    }
                    ++nread;
                    if (m_bit_offset++ == 7)
                        m_current_byte.clear();
                }
            } else {
                auto temp_buffer = TRY(ByteBuffer::create_uninitialized(1));
                TRY(m_stream.read(temp_buffer.bytes()));
                m_current_byte = temp_buffer[0];
                m_bit_offset = 0;
            }
        }

        return result;
    }

    /// Discards any sub-byte stream positioning the input stream may be keeping track of.
    /// Non-bitwise reads will implicitly call this.
    void align_to_byte_boundary()
    {
        m_current_byte.clear();
        m_bit_offset = 0;
    }

    /// Whether we are (accidentally or intentionally) at a byte boundary right now.
    ALWAYS_INLINE bool is_aligned_to_byte_boundary() const { return m_bit_offset == 0; }

private:
    BigEndianInputBitStream(Stream& stream)
        : m_stream(stream)
    {
    }

    Optional<u8> m_current_byte;
    size_t m_bit_offset { 0 };
    Stream& m_stream;
};

/// A stream wrapper class that allows you to read arbitrary amounts of bits
/// in little-endian order from another stream.
/// Note that this stream does not own its underlying stream, it merely takes a reference.
class LittleEndianInputBitStream : public Stream {
public:
    static ErrorOr<NonnullOwnPtr<LittleEndianInputBitStream>> construct(Stream& stream)
    {
        return adopt_nonnull_own_or_enomem<LittleEndianInputBitStream>(new LittleEndianInputBitStream(stream));
    }

    LittleEndianInputBitStream(Stream& stream)
        : m_stream(stream)
    {
    }

    // ^Stream
    virtual ErrorOr<Bytes> read(Bytes bytes) override
    {
        if (m_current_byte.has_value() && is_aligned_to_byte_boundary()) {
            bytes[0] = m_current_byte.release_value();
            return m_stream.read(bytes.slice(1));
        }
        align_to_byte_boundary();
        return m_stream.read(bytes);
    }
    virtual ErrorOr<size_t> write(ReadonlyBytes bytes) override { return m_stream.write(bytes); }
    virtual bool write_or_error(ReadonlyBytes bytes) override { return m_stream.write_or_error(bytes); }
    virtual bool is_eof() const override { return m_stream.is_eof() && !m_current_byte.has_value(); }
    virtual bool is_open() const override { return m_stream.is_open(); }
    virtual void close() override
    {
        m_stream.close();
        align_to_byte_boundary();
    }

    ErrorOr<bool> read_bit()
    {
        return read_bits<bool>(1);
    }
    /// Depending on the number of bits to read, the return type can be chosen appropriately.
    /// This avoids a bunch of static_cast<>'s for the user.
    // TODO: Support u128, u256 etc. as well: The concepts would be quite complex.
    template<Unsigned T = u64>
    ErrorOr<T> read_bits(size_t count)
    {
        if constexpr (IsSame<bool, T>) {
            VERIFY(count == 1);
        }
        T result = 0;

        size_t nread = 0;
        while (nread < count) {
            if (m_current_byte.has_value()) {
                if constexpr (!IsSame<bool, T> && !IsSame<u8, T>) {
                    // read as many bytes as possible directly
                    if (((count - nread) >= 8) && is_aligned_to_byte_boundary()) {
                        // shift existing data over
                        result |= (m_current_byte.value() << nread);
                        nread += 8;
                        m_current_byte.clear();
                    } else {
                        auto const bit = (m_current_byte.value() >> m_bit_offset) & 1;
                        result |= (bit << nread);
                        ++nread;
                        if (m_bit_offset++ == 7)
                            m_current_byte.clear();
                    }
                } else {
                    // Always take this branch for booleans or u8: there's no purpose in reading more than a single bit
                    auto const bit = (m_current_byte.value() >> m_bit_offset) & 1;
                    if constexpr (IsSame<bool, T>)
                        result = bit;
                    else
                        result |= (bit << nread);
                    ++nread;
                    if (m_bit_offset++ == 7)
                        m_current_byte.clear();
                }
            } else {
                auto temp_buffer = TRY(ByteBuffer::create_uninitialized(1));
                auto read_bytes = TRY(m_stream.read(temp_buffer.bytes()));
                if (read_bytes.is_empty())
                    return Error::from_string_literal("eof");
                m_current_byte = temp_buffer[0];
                m_bit_offset = 0;
            }
        }

        return result;
    }

    /// Discards any sub-byte stream positioning the input stream may be keeping track of.
    /// Non-bitwise reads will implicitly call this.
    u8 align_to_byte_boundary()
    {
        u8 remaining_bits = m_current_byte.value_or(0) >> m_bit_offset;
        m_current_byte.clear();
        m_bit_offset = 0;
        return remaining_bits;
    }

    /// Whether we are (accidentally or intentionally) at a byte boundary right now.
    ALWAYS_INLINE bool is_aligned_to_byte_boundary() const { return m_bit_offset == 0; }

private:
    Optional<u8> m_current_byte;
    size_t m_bit_offset { 0 };
    Stream& m_stream;
};

}