summaryrefslogtreecommitdiff
path: root/Userland/Libraries/LibCore/InputBitStream.h
blob: 0e844dc107fc248b2b4d7f0de7f0706566700774 (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
/*
 * 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 bool is_readable() const override { return m_stream.is_readable(); }
    virtual ErrorOr<size_t> 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 bool is_writable() const override { return m_stream.is_writable(); }
    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 {
                        const auto 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
                    const auto 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 {
                // FIXME: This returns Optional so TRY is not useable
                auto temp_buffer = ByteBuffer::create_uninitialized(1);
                if (!temp_buffer.has_value())
                    return Error::from_string_literal("Couldn't allocate temporary byte buffer"sv);

                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;
};

}