AK: Add Bitmap::count_in_range and Bitmap::fill_range

AK/Bitmap.h

@@ -103,6 +103,57 @@ public:
         }
     }
 
+    size_t count_slow(bool value) const
+    {
+        return count_in_range(0, m_size, value);
+    }
+
+    size_t count_in_range(size_t start, size_t len, bool value) const
+    {
+        ASSERT(start < m_size);
+        ASSERT(start + len <= m_size);
+        if (len == 0)
+            return 0;
+
+        static const u8 bitmask_first_byte[8] = { 0xFF, 0xFE, 0xFC, 0xF8, 0xF0, 0xE0, 0xC0, 0x80 };
+        static const u8 bitmask_last_byte[8] = { 0x0, 0x1, 0x3, 0x7, 0xF, 0x1F, 0x3F, 0x7F };
+
+        size_t count;
+        const u8* first = &m_data[start / 8];
+        const u8* last = &m_data[(start + len) / 8];
+        u8 byte = *first;
+        byte &= bitmask_first_byte[start % 8];
+        if (first == last) {
+            byte &= bitmask_last_byte[(start + len) % 8];
+            count = __builtin_popcount(byte);
+        } else {
+            count = __builtin_popcount(byte);
+            byte = *last;
+            byte &= bitmask_last_byte[(start + len) % 8];
+            count += __builtin_popcount(byte);
+            if (++first < last) {
+                const u32* ptr32 = (const u32*)(((FlatPtr)first + sizeof(u32) - 1) & ~(sizeof(u32) - 1));
+                if ((const u8*)ptr32 > last)
+                    ptr32 = (const u32*)last;
+                while (first < (const u8*)ptr32) {
+                    count += __builtin_popcount(*first);
+                    first++;
+                }
+                const u32* last32 = (const u32*)((FlatPtr)last & ~(sizeof(u32) - 1));
+                while (ptr32 < last32) {
+                    count += __builtin_popcountl(*ptr32);
+                    ptr32++;
+                }
+                for (first = (const u8*)ptr32; first < last; first++)
+                    count += __builtin_popcount(*first);
+            }
+        }
+
+        if (!value)
+            count = len - count;
+        return count;
+    }
+
     u8* data() { return m_data; }
     const u8* data() const { return m_data; }
 
@@ -127,6 +178,53 @@ public:
         }
     }
 
+    template<bool VALUE>
+    void fill_range(size_t start, size_t len)
+    {
+        ASSERT(start < m_size);
+        ASSERT(start + len <= m_size);
+        if (len == 0)
+            return;
+
+        static const u8 bitmask_first_byte[8] = { 0xFF, 0xFE, 0xFC, 0xF8, 0xF0, 0xE0, 0xC0, 0x80 };
+        static const u8 bitmask_last_byte[8] = { 0x0, 0x1, 0x3, 0x7, 0xF, 0x1F, 0x3F, 0x7F };
+
+        u8* first = &m_data[start / 8];
+        u8* last = &m_data[(start + len) / 8];
+        u8 byte_mask = bitmask_first_byte[start % 8];
+        if (first == last) {
+            byte_mask &= bitmask_last_byte[(start + len) % 8];
+            if constexpr (VALUE)
+                *first |= byte_mask;
+            else
+                *first &= ~byte_mask;
+        } else {
+            if constexpr (VALUE)
+                *first |= byte_mask;
+            else
+                *first &= ~byte_mask;
+            byte_mask = bitmask_last_byte[(start + len) % 8];
+            if constexpr (VALUE)
+                *last |= byte_mask;
+            else
+                *last &= ~byte_mask;
+            if (++first < last) {
+                if constexpr (VALUE)
+                    __builtin_memset(first, 0xFF, last - first);
+                else
+                    __builtin_memset(first, 0x0, last - first);
+            }
+        }
+    }
+
+    void fill_range(size_t start, size_t len, bool value)
+    {
+        if (value)
+            fill_range<true>(start, len);
+        else
+            fill_range<false>(start, len);
+    }
+
     void fill(bool value)
     {
         __builtin_memset(m_data, value ? 0xff : 0x00, size_in_bytes());

AK/Tests/TestBitmap.cpp

@@ -136,4 +136,67 @@ TEST_CASE(find_longest_range_of_unset_bits_edge)
     EXPECT_EQ(result.value(), 32u);
 }
 
+TEST_CASE(count_in_range)
+{
+    Bitmap bitmap(256, false);
+    bitmap.set(14, true);
+    bitmap.set(17, true);
+    bitmap.set(19, true);
+    bitmap.set(20, true);
+    for (size_t i = 34; i < 250; i++) {
+        if (i < 130 || i > 183)
+            bitmap.set(i, true);
+    }
+
+    auto count_bits_slow = [](const Bitmap& b, size_t start, size_t len, bool value) -> size_t {
+        size_t count = 0;
+        for (size_t i = start; i < start + len; i++) {
+            if (b.get(i) == value)
+                count++;
+        }
+        return count;
+    };
+    auto test_with_value = [&](bool value) {
+        auto do_test = [&](size_t start, size_t len) {
+            EXPECT_EQ(bitmap.count_in_range(start, len, value), count_bits_slow(bitmap, start, len, value));
+        };
+        do_test(16, 2);
+        do_test(16, 3);
+        do_test(16, 4);
+
+        for (size_t start = 8; start < 24; start++) {
+            for (size_t end = 9; end < 25; end++) {
+                if (start >= end)
+                    continue;
+                do_test(start, end - start);
+            }
+        }
+
+        for (size_t start = 1; start <= 9; start++) {
+            for (size_t i = start + 1; i < bitmap.size() - start + 1; i++)
+                do_test(start, i - start);
+        }
+    };
+    test_with_value(true);
+    test_with_value(false);
+}
+
+TEST_CASE(fill_range)
+{
+    Bitmap bitmap(288, false);
+    bitmap.fill_range(48, 32, true);
+    bitmap.fill_range(94, 39, true);
+    bitmap.fill_range(190, 71, true);
+    bitmap.fill_range(190 + 71 - 7, 21, false); // slighly overlapping clear
+
+    for (size_t i = 0; i < bitmap.size(); i++) {
+        bool should_be_set = (i >= 48 && i < 48 + 32)
+            || (i >= 94 && i < 94 + 39)
+            || ((i >= 190 && i < 190 + 71) && !(i >= 190 + 71 - 7 && i < 190 + 71 - 7 + 21));
+        EXPECT_EQ(bitmap.get(i), should_be_set);
+    }
+
+    EXPECT_EQ(bitmap.count_slow(true), 32u + 39u + 71u - 7u);
+}
+
 TEST_MAIN(Bitmap)