Create bit-level buffer for reading and writing individual bits

lib/common/src/me/devnatan/kompress/BitBuffer.kt

@@ -0,0 +1,74 @@
+package me.devnatan.kompress
+
+/** Internally manages a 32-bit buffer and tracks how many bits are currently stored. */
+internal class BitBuffer {
+    private var buffer: UInt = 0u
+    private var bitsInBuffer: Int = 0
+
+    fun writeBits(value: UInt, numBits: Int) {
+        require(numBits in 1..32) { "`numBits` must be between 1 and 32. Given $numBits" }
+
+        val masked = if (numBits == 32) {
+            value
+        } else {
+            value and ((1u shl numBits) - 1u)
+        }
+
+        if (bitsInBuffer + numBits <= 32) {
+            buffer = buffer or (masked shl bitsInBuffer)
+            bitsInBuffer += numBits
+        } else {
+            // would overflow 32-bit buffer (this shouldn't happen?)
+            // but we handle it by only taking what fits
+            val bitsToWrite = 32 - bitsInBuffer
+            val partialMask = (1u shl bitsToWrite) - 1u
+            buffer = buffer or ((masked and partialMask) shl bitsInBuffer)
+            bitsInBuffer = 32
+        }
+    }
+
+    fun readBits(numBits: Int): UInt {
+        require(numBits in 1..32) { "`numBits` must be between 1 and 32. Given: $numBits" }
+        require(bitsInBuffer >= numBits) { "Not enough bits in buffer. Expected: $numBits, given: $bitsInBuffer" }
+
+        val value = if (numBits == 32)
+            buffer
+        else
+            buffer and ((1u shl numBits) - 1u)
+
+        buffer = if (numBits == 32)
+            0u
+        else
+            buffer shr numBits
+
+        bitsInBuffer -= numBits
+
+        return value
+    }
+
+    fun hasBits(numBits: Int): Boolean = bitsInBuffer >= numBits
+
+    fun alignToByte() {
+        val bitsToSkip = bitsInBuffer % 8
+        if (bitsToSkip <= 0) return
+
+        buffer = buffer shr bitsToSkip
+        bitsInBuffer -= bitsToSkip
+    }
+
+    fun reset() {
+        buffer = 0u
+        bitsInBuffer = 0
+    }
+
+    fun availableBytes(): Int = bitsInBuffer / 8
+
+    fun readByte(): UByte {
+        require(bitsInBuffer >= 8) { "Not enough bits for a byte. Given: $bitsInBuffer" }
+        return readBits(8).toUByte()
+    }
+
+    fun writeByte(byte: UByte) {
+        writeBits(byte.toUInt(), 8)
+    }
+}

lib/common/test/me/devnatan/kompress/BitBufferTest.kt

@@ -0,0 +1,67 @@
+package me.devnatan.kompress
+
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+class BitBufferTest {
+
+    @Test
+    fun `write and read zero`() {
+        val buffer = BitBuffer()
+        buffer.writeBits(0u, 1)
+        assertEquals(
+            expected = 0u,
+            actual = buffer.readBits(1)
+        )
+    }
+
+    @Test
+    fun `write and read single`() {
+        val buffer = BitBuffer()
+        buffer.writeBits(1u, 1)
+        assertEquals(
+            expected = 1u,
+            actual = buffer.readBits(1)
+        )
+    }
+
+    @Test
+    fun `write and read full`() {
+        val buffer = BitBuffer()
+        buffer.writeBits(0xFFu, 8)
+        assertEquals(
+            expected = 0xFFu,
+            actual = buffer.readBits(8)
+        )
+    }
+
+    @Test
+    fun `write and read 16`() {
+        val buffer = BitBuffer()
+        buffer.writeBits(0xABCDu, 16)
+        assertEquals(
+            expected = 0xABCDu,
+            actual = buffer.readBits(16)
+        )
+    }
+
+    @Test
+    fun `write and read 32`() {
+        val buffer = BitBuffer()
+        buffer.writeBits(0x12345678u, 32)
+        assertEquals(
+            expected = 0x12345678u,
+            actual = buffer.readBits(32)
+        )
+    }
+
+    @Test
+    fun `write and read 32 max`() {
+        val buffer = BitBuffer()
+        buffer.writeBits(0xFFFFFFFFu, 32)
+        assertEquals(
+            expected = 0xFFFFFFFFu,
+            actual = buffer.readBits(32)
+        )
+    }
+}