From e9e89d7e4edefa5a4c5a2762a0ae710b092773f8 Mon Sep 17 00:00:00 2001
From: Tim Schumacher <timschumi@gmx.de>
Date: Tue, 28 Nov 2023 13:06:19 +0100
Subject: [PATCH] AK: Optimize BitStream refilling a bit further

This tries to optimize the refill code by making it easier to digest for
the branch predictor. This includes not looping as much across function
calls and marking our EOF case to be unlikely.

Co-Authored-By: Lucas Chollet <lucas.chollet@free.fr>
---
 AK/BitStream.h | 39 ++++++++++++++++++++-------------------
 1 file changed, 20 insertions(+), 19 deletions(-)

diff --git a/AK/BitStream.h b/AK/BitStream.h
index 4d7882ab13..574f667c0b 100644
--- a/AK/BitStream.h
+++ b/AK/BitStream.h
@@ -215,15 +215,8 @@ public:
     template<Unsigned T = u64>
     ErrorOr<T> peek_bits(size_t count)
     {
-        while (count > m_bit_count) {
-            if (TRY(refill_buffer_from_stream()))
-                continue;
-
-            if (m_unsatisfiable_read_behavior == UnsatisfiableReadBehavior::Reject)
-                return Error::from_string_literal("Reached end-of-stream without collecting the required number of bits");
-
-            break;
-        }
+        if (count > m_bit_count)
+            TRY(refill_buffer_from_stream(count));
 
         return m_bit_buffer & lsb_mask<T>(min(count, m_bit_count));
     }
@@ -254,21 +247,29 @@ public:
     }
 
 private:
-    ErrorOr<bool> refill_buffer_from_stream()
+    ErrorOr<void> refill_buffer_from_stream(size_t requested_bit_count)
     {
-        if (m_stream->is_eof())
-            return false;
+        while (requested_bit_count > m_bit_count) [[likely]] {
+            if (m_stream->is_eof()) [[unlikely]] {
+                if (m_unsatisfiable_read_behavior == UnsatisfiableReadBehavior::FillWithZero) {
+                    m_bit_count = requested_bit_count;
+                    return {};
+                }
 
-        size_t bits_to_read = bit_buffer_size - m_bit_count;
-        size_t bytes_to_read = bits_to_read / bits_per_byte;
+                return Error::from_string_literal("Reached end-of-stream without collecting the required number of bits");
+            }
 
-        BufferType buffer = 0;
-        auto bytes = TRY(m_stream->read_some({ &buffer, bytes_to_read }));
+            size_t bits_to_read = bit_buffer_size - m_bit_count;
+            size_t bytes_to_read = bits_to_read / bits_per_byte;
 
-        m_bit_buffer |= (buffer << m_bit_count);
-        m_bit_count += bytes.size() * bits_per_byte;
+            BufferType buffer = 0;
+            auto bytes = TRY(m_stream->read_some({ &buffer, bytes_to_read }));
 
-        return true;
+            m_bit_buffer |= (buffer << m_bit_count);
+            m_bit_count += bytes.size() * bits_per_byte;
+        }
+
+        return {};
     }
 
     UnsatisfiableReadBehavior m_unsatisfiable_read_behavior;