Add a wrapper for RTL_BITMAP (#6)

* Add a wrapper for RTL_BITMAP

* Fix Bitmap and BitmapRangeIterator, add basic tests

---------

Co-authored-by: Serhii Malyi <maliy.sergey@apriorit.com>
Co-authored-by: Sergey Podobry <sergius@apriorit.com>

Author: 3 people

include/kf/Bitmap.h

@@ -0,0 +1,122 @@
+#pragma once
+#include <kf/stl/new>
+#include <kf/BitmapRangeIterator.h>
+
+namespace kf
+{
+    //////////////////////////////////////////////////////////////////////////
+    // Bitmap - effectively works with set of bits
+
+    template<POOL_TYPE poolType>
+    class Bitmap
+    {
+    public:
+        Bitmap() noexcept = default;
+
+        ~Bitmap() noexcept
+        {
+            deinitialize();
+        }
+
+        // Non-copyable
+        Bitmap(const Bitmap&) = delete;
+        Bitmap& operator=(const Bitmap&) = delete;
+
+        // Movable
+        Bitmap(Bitmap&& other) noexcept : m_header(other.m_header)
+        {
+            other.m_header = {};
+        }
+
+        Bitmap& operator=(Bitmap&& other) noexcept
+        {
+            if (&other != this)
+            {
+                m_header = other.m_header;
+                other.m_header = {};
+            }
+
+            return *this;
+        }
+
+        // IRQL <= APC_LEVEL for PagedPool and on Windows 7 and earlier
+        [[nodiscard]] NTSTATUS initialize(ULONG size) noexcept
+        {
+            deinitialize();
+
+            // Convert bits to bytes and align up to sizeof(ULONG) bytes
+            const ULONG bufferSize = ALIGN_UP_BY(size, sizeof(ULONG) * CHAR_BIT) / CHAR_BIT;
+
+            auto buffer = operator new(bufferSize, poolType);
+            if (!buffer)
+            {
+                return STATUS_INSUFFICIENT_RESOURCES;
+            }
+
+            RtlZeroMemory(buffer, bufferSize);
+            RtlInitializeBitMap(&m_header, static_cast<PULONG>(buffer), size);
+
+            return STATUS_SUCCESS;
+        }
+
+        ULONG size() const
+        {
+            return m_header.SizeOfBitMap;
+        }
+
+        void setBits(ULONG startingIndex, ULONG numberToSet) noexcept
+        {
+            RtlSetBits(&m_header, startingIndex, numberToSet);
+        }
+
+        void clearBits(ULONG startingIndex, ULONG numberToSet) noexcept
+        {
+            RtlClearBits(&m_header, startingIndex, numberToSet);
+        }
+
+        void setAll() noexcept
+        {
+            RtlSetAllBits(&m_header);
+        }
+
+        void clearAll() noexcept
+        {
+            RtlClearAllBits(&m_header);
+        }
+
+        bool areBitsSet(ULONG startingIndex, ULONG size) noexcept
+        {
+            return RtlAreBitsSet(&m_header, startingIndex, size);
+        }
+
+        bool areBitsClear(ULONG startingIndex, ULONG size) noexcept
+        {
+            return RtlAreBitsClear(&m_header, startingIndex, size);
+        }
+
+        ULONG numberOfSetBits() noexcept
+        {
+            return RtlNumberOfSetBits(&m_header);
+        }
+
+        ULONG numberOfClearBits() noexcept
+        {
+            return RtlNumberOfClearBits(&m_header);
+        }
+
+        BitmapRangeIterator rangeIterator(ULONG startingIndex = 0) noexcept
+        {
+            return BitmapRangeIterator(&m_header, startingIndex);
+        }
+
+    private:
+        void deinitialize() noexcept
+        {
+            operator delete(m_header.Buffer);
+            m_header = {};
+        }
+
+    private:
+        RTL_BITMAP m_header = {};
+    };
+}

include/kf/BitmapRangeIterator.h

@@ -0,0 +1,85 @@
+#pragma once
+#include <optional>
+#include <kf/stl/cassert>
+
+namespace kf
+{
+    //////////////////////////////////////////////////////////////////////////
+    // BitmapRangeIterator - enumerates ranges of bits set to 1
+
+    class BitmapRangeIterator
+    {
+    public:
+        BitmapRangeIterator(PRTL_BITMAP header, ULONG startingIndex = 0)
+            : m_header(header)
+            , m_index(startingIndex)
+        {
+        }
+
+        //
+        // Return range of set bits as a pair { startIndex, lenght }, nullopt means the end of iteration
+        //
+
+        std::optional<std::pair<ULONG, ULONG>> next()
+        {
+            const auto range = internalNext();
+
+            return range.second ? std::make_optional(range) : std::nullopt;
+        }
+
+    private:
+        std::pair<ULONG, ULONG> internalNext()
+        {
+            if (m_index >= m_header->SizeOfBitMap)
+            {
+                //
+                // We're beyond the end, there is nothing to report
+                //
+
+                return {};
+            }
+
+            ULONG clearBitsStart = ULONG_MAX;
+            ULONG clearBitsCount = RtlFindNextForwardRunClear(m_header, m_index, &clearBitsStart);
+
+            if (!clearBitsCount)
+            {
+                //
+                // No clear bits are found, then we have the range at the end of the bitmap to report.
+                // Set `clearBitsStart` beyond the last bit.
+                //
+
+                clearBitsStart = m_header->SizeOfBitMap;
+            }
+
+            //
+            // Range of set bits is from the current index to the first clear bit
+            //
+
+            std::pair<ULONG, ULONG> range = { m_index, clearBitsStart - m_index };
+
+            //
+            // Update the current index to the next bit beyond the last clear bit we've found
+            //
+
+            m_index = clearBitsStart + clearBitsCount;
+
+            if (!range.second)
+            {
+                //
+                // We've got an empty range! The bitmap starts from clear bits so we need to run internalNext again
+                //
+
+                assert(!range.first); // We expect this case only at the beginning of the bitmap
+
+                range = internalNext();
+            }
+
+            return range;
+        }
+
+    private:
+        PRTL_BITMAP m_header;
+        ULONG       m_index = 0;
+    };
+}

test/Bitmap.cpp

@@ -0,0 +1,35 @@
+#include "pch.h"
+#include <kf/Bitmap.h>
+
+SCENARIO("Bitmap")
+{
+    GIVEN("initialized bitmap with size 10")
+    {
+        kf::Bitmap<PagedPool> bitmap;
+        REQUIRE_NT_SUCCESS(bitmap.initialize(10));
+
+        WHEN("do nothing")
+        {
+            THEN("bitmap size is 10")
+            {
+                REQUIRE(bitmap.size() == 10);
+            }
+
+            THEN("each bit is clear")
+            {
+                for (int i = 0; i < 10; ++i)
+                {
+                    REQUIRE(bitmap.areBitsClear(i, 1));
+                }
+            }
+
+            THEN("no bit is set")
+            {
+                for (int i = 0; i < 10; ++i)
+                {
+                    REQUIRE(!bitmap.areBitsSet(i, 1));
+                }
+            }
+        }
+    }
+}

test/BitmapRangeIterator.cpp

@@ -0,0 +1,43 @@
+#include "pch.h"
+#include <kf/Bitmap.h>
+
+SCENARIO("BitmapRangeIterator")
+{
+    GIVEN("initialized Bitmap with size 10")
+    {
+        kf::Bitmap<PagedPool> bitmap;
+        REQUIRE_NT_SUCCESS(bitmap.initialize(10));
+
+        auto iterator = bitmap.rangeIterator();
+
+        WHEN("do nothing, bitmap is [0,0,0,0,0,0,0,0,0,0]")
+        {
+            THEN("iterator returs no ranges")
+            {
+                REQUIRE(!iterator.next());
+            }
+        }
+
+        WHEN("set bitmap to [1,0,0,0,0,0,0,0,0,0]")
+        {
+            bitmap.setBits(0, 1);
+
+            THEN("iterator returs ranges: {0,1}")
+            {
+                REQUIRE(*iterator.next() == std::make_pair(0UL, 1UL));
+                REQUIRE(!iterator.next());
+            }
+        }
+
+        WHEN("set bitmap to [1,1,1,0,0,0,0,0,0,0]")
+        {
+            bitmap.setBits(0, 3);
+
+            THEN("iterator returs ranges: {0,3}")
+            {
+                REQUIRE(*iterator.next() == std::make_pair(0UL, 3UL));
+                REQUIRE(!iterator.next());
+            }
+        }
+    }
+}

test/CMakeLists.txt

@@ -39,6 +39,8 @@ wdk_add_driver(kf-test WINVER NTDDI_WIN10 STL
     pch.h
     pch.cpp
     AdjacentView.cpp
+    Bitmap.cpp
+    BitmapRangeIterator.cpp
     HexTest.cpp
     Vector.cpp
 )