From e77031ce67b427e5722ac0a2f9302c3701b64964 Mon Sep 17 00:00:00 2001
From: asynts <asynts@gmail.com>
Date: Fri, 1 Jan 2021 21:32:59 +0100
Subject: [PATCH] AK: Deal with unsigned integers in binary search.

---
 AK/BinarySearch.h                 | 22 +++++++++++++++++-----
 AK/Tests/TestBinarySearch.cpp     | 12 ++++++++++++
 Libraries/LibCompress/Deflate.cpp |  4 +++-
 3 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/AK/BinarySearch.h b/AK/BinarySearch.h
index 14e630089f..1c998ce0a6 100644
--- a/AK/BinarySearch.h
+++ b/AK/BinarySearch.h
@@ -32,11 +32,21 @@
 
 namespace AK {
 
-struct IntegralComparator {
-    constexpr auto operator()(auto& lhs, auto& rhs) { return lhs - rhs; }
+struct DefaultComparator {
+    template<typename T, typename S>
+    constexpr int operator()(T& lhs, S& rhs)
+    {
+        if (lhs > rhs)
+            return 1;
+
+        if (lhs < rhs)
+            return -1;
+
+        return 0;
+    }
 };
 
-template<typename Container, typename Needle, typename Comparator = IntegralComparator>
+template<typename Container, typename Needle, typename Comparator = DefaultComparator>
 constexpr auto binary_search(
     Container&& haystack,
     Needle&& needle,
@@ -52,8 +62,10 @@ constexpr auto binary_search(
     size_t low = 0;
     size_t high = haystack.size() - 1;
     while (low <= high) {
-        size_t middle = low + ((high - low) / 2);
-        auto comparison = comparator(needle, haystack[middle]);
+        size_t middle = low + (high - low) / 2;
+
+        int comparison = comparator(needle, haystack[middle]);
+
         if (comparison < 0)
             if (middle != 0)
                 high = middle - 1;
diff --git a/AK/Tests/TestBinarySearch.cpp b/AK/Tests/TestBinarySearch.cpp
index d39d536b53..56a1ca617f 100644
--- a/AK/Tests/TestBinarySearch.cpp
+++ b/AK/Tests/TestBinarySearch.cpp
@@ -125,4 +125,16 @@ TEST_CASE(constexpr_array_search)
     static_assert(binary_search(array, 3) == nullptr);
 }
 
+TEST_CASE(unsigned_to_signed_regression)
+{
+    const Array<u32, 5> input { 0, 1, 2, 3, 4 };
+
+    // The algorithm computes 1 - input[2] = -1, and if this is (incorrectly) cast
+    // to an unsigned then it will look in the wrong direction and miss the 1.
+
+    size_t nearby_index = 1;
+    EXPECT_EQ(binary_search(input, 1u, &nearby_index), &input[1]);
+    EXPECT_EQ(nearby_index, 1u);
+}
+
 TEST_MAIN(BinarySearch)
diff --git a/Libraries/LibCompress/Deflate.cpp b/Libraries/LibCompress/Deflate.cpp
index 840cb6d3b7..0f67b7035e 100644
--- a/Libraries/LibCompress/Deflate.cpp
+++ b/Libraries/LibCompress/Deflate.cpp
@@ -109,8 +109,10 @@ u32 CanonicalCode::read_symbol(InputBitStream& stream) const
 
     for (;;) {
         code_bits = code_bits << 1 | stream.read_bits(1);
+        ASSERT(code_bits < (1 << 16));
 
-        // FIXME: This seems really inefficient, this could be an index into an array instead.
+        // FIXME: This is very inefficent and could greatly be improved by implementing this
+        //        algorithm: https://www.hanshq.net/zip.html#huffdec
         size_t index;
         if (AK::binary_search(m_symbol_codes.span(), code_bits, &index))
             return m_symbol_values[index];