1 files changed, 112 insertions, 0 deletions
diff --git a/Userland/Libraries/LibCore/InputBitStream.h b/Userland/Libraries/LibCore/InputBitStream.h
index ea2bb1c9ca..e5340cb96f 100644
--- a/Userland/Libraries/LibCore/InputBitStream.h
+++ b/Userland/Libraries/LibCore/InputBitStream.h
@@ -131,4 +131,116 @@ private:
     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 bool is_readable() const override { return m_stream.is_readable(); }
+    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 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 |= (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;
+};
+
 }