From 172ede6dfd29f86a90c47a99b9cfe8a8d738ba8c Mon Sep 17 00:00:00 2001
From: Nico Weber <thakis@chromium.org>
Date: Mon, 9 Sep 2024 10:04:22 -0400
Subject: [PATCH] AK+LibCompress: Allow peeking more bits than available

...and then fail only when discard_previously_peeked_bits() more
bits than we have.

`unzip ~/Downloads/enwik8.zip`: 1.005s -> 1.050s
4.9% slowdown :/

Fixes #25005
---
 AK/BitStream.h                             | 18 +++++++-----------
 Userland/Libraries/LibCompress/Deflate.cpp |  2 +-
 2 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/AK/BitStream.h b/AK/BitStream.h
index 94b429ed3de2fa..8bdf0fcc40e019 100644
--- a/AK/BitStream.h
+++ b/AK/BitStream.h
@@ -208,7 +208,7 @@ class LittleEndianInputBitStream : public LittleEndianBitStream {
     ErrorOr<T> read_bits(size_t count)
     {
         auto result = TRY(peek_bits<T>(count));
-        discard_previously_peeked_bits(count);
+        TRY(discard_previously_peeked_bits(count));
 
         return result;
     }
@@ -222,15 +222,14 @@ class LittleEndianInputBitStream : public LittleEndianBitStream {
         return m_bit_buffer & lsb_mask<T>(min(count, m_bit_count));
     }
 
-    ALWAYS_INLINE void discard_previously_peeked_bits(u8 count)
+    ALWAYS_INLINE ErrorOr<void> discard_previously_peeked_bits(u8 count)
     {
-        // We allow "retrieving" more bits than we can provide, but we need to make sure that we don't underflow the current bit counter.
-        // This only affects certain "modes", but all the relevant checks have been handled in the respective `peek_bits` call.
         if (count > m_bit_count)
-            count = m_bit_count;
+            return Error::from_string_literal("Reached end-of-stream without collecting the required number of bits");
 
         m_bit_buffer >>= count;
         m_bit_count -= count;
+        return {};
     }
 
     /// Discards any sub-byte stream positioning the input stream may be keeping track of.
@@ -241,7 +240,7 @@ class LittleEndianInputBitStream : public LittleEndianBitStream {
 
         if (auto offset = m_bit_count % bits_per_byte; offset != 0) {
             remaining_bits = m_bit_buffer & lsb_mask<u8>(offset);
-            discard_previously_peeked_bits(offset);
+            MUST(discard_previously_peeked_bits(offset));
         }
 
         return remaining_bits;
@@ -252,12 +251,9 @@ class LittleEndianInputBitStream : public LittleEndianBitStream {
     {
         while (requested_bit_count > m_bit_count) [[likely]] {
             if (m_stream->is_eof()) [[unlikely]] {
-                if (m_unsatisfiable_read_behavior == UnsatisfiableReadBehavior::FillWithZero) {
+                if (m_unsatisfiable_read_behavior == UnsatisfiableReadBehavior::FillWithZero)
                     m_bit_count = requested_bit_count;
-                    return {};
-                }
-
-                return Error::from_string_literal("Reached end-of-stream without collecting the required number of bits");
+                return {};
             }
 
             size_t bits_to_read = bit_buffer_size - m_bit_count;
diff --git a/Userland/Libraries/LibCompress/Deflate.cpp b/Userland/Libraries/LibCompress/Deflate.cpp
index a952361443a256..b12cf2b2ff5505 100644
--- a/Userland/Libraries/LibCompress/Deflate.cpp
+++ b/Userland/Libraries/LibCompress/Deflate.cpp
@@ -148,7 +148,7 @@ ErrorOr<u32> CanonicalCode::read_symbol(LittleEndianInputBitStream& stream) cons
     auto prefix = TRY(stream.peek_bits<size_t>(m_max_prefixed_code_length));
 
     if (auto [symbol_value, code_length] = m_prefix_table[prefix]; code_length != 0) {
-        stream.discard_previously_peeked_bits(code_length);
+        TRY(stream.discard_previously_peeked_bits(code_length));
         return symbol_value;
     }