Refactor bitbuffer

common/src/main/java/com/pedro/common/BitBuffer.kt

@@ -16,26 +16,114 @@
 
 package com.pedro.common
 
+import java.nio.ByteBuffer
+import kotlin.math.ceil
+
 /**
  * Created by pedro on 9/12/23.
  *
  */
-class BitBuffer(private val buffer: ByteArray) {
+class BitBuffer(val buffer: ByteBuffer) {
+  var bufferPosition = buffer.position() * Byte.SIZE_BITS
+    private set
+  private val bufferEnd = buffer.limit() * Byte.SIZE_BITS
+
+  val bitRemaining: Int
+    get() = bufferEnd - bufferPosition
 
-  fun getBits(offset: Int, size: Int): Int {
-    val startIndex = offset / 8
-    val startBit = offset % 8
-    var result = 0
-    var bitNum = startBit
+  val remaining: Int
+    get() = ceil(bitRemaining.toDouble() / Byte.SIZE_BITS).toInt()
 
+  private fun getBits(size: Int): Long {
+    var result = 0L
     for (i in 0 until size) {
-      val nextByte = (startBit + i) % 8 < startBit
-      val currentIndex = startIndex + if (nextByte) 1 else 0
-      val currentBit = (startBit + i) % 8
+      val bitPosition = bufferPosition + i
+      val currentIndex = bitPosition / 8
+      val currentBit = bitPosition % 8
       val bitValue = (buffer[currentIndex].toInt() ushr 7 - currentBit) and 0x01
-      result = (result shl 1) or bitValue
-      bitNum++
+      result = (result shl 1) or bitValue.toLong()
     }
+    bufferPosition += size
     return result
   }
+
+  fun getBool() = getBits(1) == 1L
+
+  fun getInt(i: Int) = getBits(i).toInt()
+
+  fun getLong(i: Int) = getBits(i)
+
+  fun skip(i: Int) {
+    bufferPosition += i
+  }
+
+  fun skipBool() = skip(1)
+
+  fun readUE(): Int {
+    var leadingZeroBits = 0
+    while (!getBool()) {
+      leadingZeroBits++
+    }
+    return if (leadingZeroBits > 0) {
+      (1 shl leadingZeroBits) - 1 + getInt(leadingZeroBits)
+    } else {
+      0
+    }
+  }
+
+  fun readUVLC(): Int {
+    var leadingZeros = 0
+    var value = 0
+    var currentIndex = bufferPosition / 8
+    var currentBit = 7 - bufferPosition % 8
+
+    while (buffer[currentIndex].toInt() and (1 shl currentBit) == 0) {
+      leadingZeros++
+      if (currentBit == 0) {
+        currentIndex++
+        currentBit = 7
+      } else {
+        currentBit--
+      }
+    }
+
+    repeat((0 until leadingZeros + 1).count()) {
+      if (currentBit == 0) {
+        currentIndex++
+        currentBit = 7
+      } else {
+        currentBit--
+      }
+
+      value = (value shl 1) or ((buffer[currentIndex].toInt() ushr currentBit) and 1)
+    }
+    bufferPosition += 2 * leadingZeros + 1
+    return value
+  }
+
+  fun resetPosition() {
+    bufferPosition = buffer.position() * Byte.SIZE_BITS
+  }
+
+  companion object {
+
+    fun extractRbsp(buffer: ByteBuffer, headerLength: Int): ByteBuffer {
+      val rbsp = ByteBuffer.allocateDirect(buffer.remaining())
+
+      val indices = buffer.indicesOf(byteArrayOf(0x00, 0x00, 0x03))
+
+      rbsp.put(buffer, 0, headerLength)
+
+      var previous = buffer.position()
+      indices.forEach {
+        rbsp.put(buffer, previous, it + 2 - previous)
+        previous = it + 3 // skip emulation_prevention_three_byte
+      }
+      rbsp.put(buffer, previous, buffer.limit() - previous)
+
+      rbsp.limit(rbsp.position())
+      rbsp.rewind()
+      return rbsp
+    }
+  }
 }

common/src/main/java/com/pedro/common/Extensions.kt

@@ -200,7 +200,7 @@ fun Int.toUInt32(): ByteArray {
 }
 
 fun Long.toUInt64(): ByteArray {
-  return byteArrayOf((this ushr 56).toByte(), (this ushr 48).toByte(), (this ushr 48).toByte(), (this ushr 32).toByte(), (this ushr 24).toByte(), (this ushr 16).toByte(), (this ushr 8).toByte(), this.toByte())
+  return byteArrayOf((this ushr 56).toByte(), (this ushr 48).toByte(), (this ushr 40).toByte(), (this ushr 32).toByte(), (this ushr 24).toByte(), (this ushr 16).toByte(), (this ushr 8).toByte(), this.toByte())
 }
 
 fun Int.toUInt32LittleEndian(): ByteArray = Integer.reverseBytes(this).toUInt32()
@@ -310,4 +310,29 @@ fun List<ByteArray>.combine(): ByteArray {
 
 fun SecureRandom.nextBytes(size: Int): ByteArray {
   return ByteArray(size).apply { nextBytes(this) }
+}
+
+fun Boolean.toInt() = if (this) 1 else 0
+
+fun ByteBuffer.indicesOf(prefix: ByteArray): List<Int> {
+  if (prefix.isEmpty()) return emptyList()
+  val indices = mutableListOf<Int>()
+
+  outer@ for (i in 0 until this.limit() - prefix.size + 1) {
+    for (j in prefix.indices) {
+      if (this.get(i + j) != prefix[j]) {
+        continue@outer
+      }
+    }
+    indices.add(i)
+  }
+  return indices
+}
+
+fun ByteBuffer.put(buffer: ByteBuffer, offset: Int, length: Int) {
+  val limit = buffer.limit()
+  buffer.position(offset)
+  buffer.limit(offset + length)
+  this.put(buffer)
+  buffer.limit(limit)
 }

common/src/test/java/com/pedro/common/BitBufferTest.kt

@@ -16,23 +16,223 @@
 
 package com.pedro.common
 
-import junit.framework.TestCase.assertEquals
+import org.junit.Assert.assertArrayEquals
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertTrue
 import org.junit.Test
+import java.nio.ByteBuffer
 
 /**
  * Created by pedro on 10/12/23.
  */
 class BitBufferTest {
 
+  private fun bitBufferOf(vararg bytes: Int): BitBuffer {
+    return BitBuffer(ByteBuffer.wrap(ByteArray(bytes.size) { bytes[it].toByte() }))
+  }
+
   @Test
   fun checkGetBits() {
-    val buffer = byteArrayOf(0x00, 0x00, 0x00, 0x24, 0x4f, 0x7e, 0x7f, 0x00, 0x68.toByte(), 0x83.toByte(), 0x00, 0x83.toByte(), 0x02)
+    val buffer = ByteBuffer.wrap(byteArrayOf(0x00, 0x00, 0x00, 0x24, 0x4f, 0x7e, 0x7f, 0x00, 0x68.toByte(), 0x83.toByte(), 0x00, 0x83.toByte(), 0x02))
     val bitBuffer = BitBuffer(buffer)
-    val result = bitBuffer.getBits(24, 5)
-    val result2 = bitBuffer.getBits(29, 4)
-    val result3 = bitBuffer.getBits(29 + 4, 4)
+    bitBuffer.skip(24)
+    val result = bitBuffer.getInt(5)
+    val result2 = bitBuffer.getInt(4)
+    val result3 = bitBuffer.getLong(4)
     assertEquals(0x04, result)
     assertEquals(8, result2)
-    assertEquals(9, result3)
+    assertEquals(9L, result3)
+  }
+
+  @Test
+  fun `GIVEN a single byte WHEN getInt 8 THEN read the whole byte`() {
+    val bitBuffer = bitBufferOf(0xAB)
+    assertEquals(0xAB, bitBuffer.getInt(8))
+    assertEquals(8, bitBuffer.bufferPosition)
+  }
+
+  @Test
+  fun `GIVEN two bytes WHEN getInt across the byte boundary THEN combine both bytes`() {
+    val bitBuffer = bitBufferOf(0xAB, 0xCD)
+    assertEquals(0xABCD, bitBuffer.getInt(16))
+  }
+
+  @Test
+  fun `GIVEN two bytes WHEN getInt of a size that is not byte aligned THEN read the high bits`() {
+    val bitBuffer = bitBufferOf(0xAB, 0xCD)
+    assertEquals(0xABC, bitBuffer.getInt(12))
+    assertEquals(12, bitBuffer.bufferPosition)
+  }
+
+  @Test
+  fun `GIVEN an unaligned offset WHEN getInt crosses a byte boundary THEN read the right bits`() {
+    val bitBuffer = bitBufferOf(0xAB, 0xCD)
+    bitBuffer.skip(4)
+    assertEquals(0xBC, bitBuffer.getInt(8))
+    assertEquals(12, bitBuffer.bufferPosition)
+  }
+
+  @Test
+  fun `GIVEN a byte WHEN getBool repeatedly THEN read every bit as a boolean`() {
+    val bitBuffer = bitBufferOf(0xA0) // 1010 0000
+    assertTrue(bitBuffer.getBool())
+    assertFalse(bitBuffer.getBool())
+    assertTrue(bitBuffer.getBool())
+    assertFalse(bitBuffer.getBool())
+    assertEquals(4, bitBuffer.bufferPosition)
+  }
+
+  @Test
+  fun `GIVEN six bytes WHEN getLong 48 THEN read a value bigger than an Int`() {
+    val bitBuffer = bitBufferOf(0x01, 0x02, 0x03, 0x04, 0x05, 0x06)
+    assertEquals(0x010203040506L, bitBuffer.getLong(48))
+  }
+
+  @Test
+  fun `GIVEN all bits set WHEN getLong 63 THEN return Long MAX_VALUE`() {
+    val bitBuffer = bitBufferOf(0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF)
+    assertEquals(Long.MAX_VALUE, bitBuffer.getLong(63))
+  }
+
+  @Test
+  fun `GIVEN all bits set WHEN getLong 64 THEN keep the whole bit pattern`() {
+    val bitBuffer = bitBufferOf(0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF)
+    assertEquals(-1L, bitBuffer.getLong(64))
+  }
+
+  @Test
+  fun `GIVEN getInt 0 WHEN read THEN return 0 and do not move`() {
+    val bitBuffer = bitBufferOf(0xAB)
+    assertEquals(0, bitBuffer.getInt(0))
+    assertEquals(0, bitBuffer.bufferPosition)
+  }
+
+  @Test
+  fun `GIVEN a buffer WHEN skip THEN advance position and keep reads aligned`() {
+    val bitBuffer = bitBufferOf(0x12, 0x34, 0x56)
+    assertEquals(0, bitBuffer.bufferPosition)
+    bitBuffer.skip(8)
+    assertEquals(8, bitBuffer.bufferPosition)
+    assertEquals(0x34, bitBuffer.getInt(8))
+    assertEquals(16, bitBuffer.bufferPosition)
+  }
+
+  @Test
+  fun `GIVEN a buffer WHEN skipBool THEN advance a single bit`() {
+    val bitBuffer = bitBufferOf(0x40) // 0100 0000
+    bitBuffer.skipBool()
+    assertEquals(1, bitBuffer.bufferPosition)
+    assertTrue(bitBuffer.getBool())
+  }
+
+  @Test
+  fun `GIVEN a buffer WHEN reading THEN bitRemaining and remaining are updated`() {
+    val bitBuffer = bitBufferOf(0x00, 0x00, 0x00)
+    assertEquals(24, bitBuffer.bitRemaining)
+    assertEquals(3, bitBuffer.remaining)
+    bitBuffer.skip(4)
+    assertEquals(20, bitBuffer.bitRemaining)
+    assertEquals(3, bitBuffer.remaining) // ceil(20 / 8)
+    bitBuffer.getInt(4)
+    assertEquals(16, bitBuffer.bitRemaining)
+    assertEquals(2, bitBuffer.remaining)
+    bitBuffer.skip(16)
+    assertEquals(0, bitBuffer.bitRemaining)
+    assertEquals(0, bitBuffer.remaining)
+  }
+
+  @Test
+  fun `GIVEN a consumed buffer WHEN resetPosition THEN read again from the start`() {
+    val bitBuffer = bitBufferOf(0x5A, 0x3C)
+    assertEquals(0x5A, bitBuffer.getInt(8))
+    assertEquals(8, bitBuffer.bufferPosition)
+    bitBuffer.resetPosition()
+    assertEquals(0, bitBuffer.bufferPosition)
+    assertEquals(0x5A, bitBuffer.getInt(8))
+  }
+
+  @Test
+  fun `GIVEN a ByteBuffer with a non zero position WHEN create BitBuffer THEN start at that position`() {
+    val buffer = ByteBuffer.wrap(byteArrayOf(0x11, 0x22, 0x33))
+    buffer.position(1)
+    val bitBuffer = BitBuffer(buffer)
+    assertEquals(8, bitBuffer.bufferPosition)
+    assertEquals(16, bitBuffer.bitRemaining)
+    assertEquals(0x22, bitBuffer.getInt(8))
+    bitBuffer.resetPosition()
+    assertEquals(8, bitBuffer.bufferPosition)
+  }
+
+  @Test
+  fun `GIVEN an exp golomb stream WHEN readUE THEN decode each code number`() {
+    val bitBuffer = bitBufferOf(0xA6) // 1 010 011 0 -> 0, 1, 2
+    assertEquals(0, bitBuffer.readUE())
+    assertEquals(1, bitBuffer.readUE())
+    assertEquals(2, bitBuffer.readUE())
+  }
+
+  @Test
+  fun `GIVEN a longer exp golomb code WHEN readUE THEN decode the value`() {
+    val bitBuffer = bitBufferOf(0x20) // 00100 000 -> 3
+    assertEquals(3, bitBuffer.readUE())
+    assertEquals(5, bitBuffer.bufferPosition)
+  }
+
+  @Test
+  fun `GIVEN a uvlc with no leading zeros WHEN readUVLC THEN return 0 and consume one bit`() {
+    val bitBuffer = bitBufferOf(0x80) // 1000 0000
+    assertEquals(0, bitBuffer.readUVLC())
+    assertEquals(1, bitBuffer.bufferPosition)
+  }
+
+  @Test
+  fun `GIVEN a uvlc with leading zeros WHEN readUVLC THEN decode value and consume 2n+1 bits`() {
+    val bitBuffer = bitBufferOf(0x60) // 0110 0000
+    assertEquals(2, bitBuffer.readUVLC())
+    assertEquals(3, bitBuffer.bufferPosition)
+  }
+
+  @Test
+  fun `GIVEN a uvlc whose value crosses a byte boundary WHEN readUVLC THEN decode it`() {
+    // 0000 1 101 | 10... -> 4 leading zeros, terminator, value bits span into the second byte
+    val bitBuffer = bitBufferOf(0x0D, 0x80)
+    assertEquals(22, bitBuffer.readUVLC())
+    assertEquals(9, bitBuffer.bufferPosition)
+  }
+
+  @Test
+  fun `GIVEN a uvlc whose leading zeros cross a byte boundary WHEN readUVLC THEN decode it`() {
+    // 8 leading zeros (whole first byte), terminator on the second byte, value spans into the third
+    val bitBuffer = bitBufferOf(0x00, 0xD5, 0x40)
+    assertEquals(341, bitBuffer.readUVLC())
+    assertEquals(17, bitBuffer.bufferPosition)
+  }
+
+  @Test
+  fun `GIVEN an unaligned position WHEN readUVLC THEN decode from that bit offset`() {
+    // skip 4 bits, then 0 1 11 -> 1 leading zero, terminator, value = 3
+    val bitBuffer = bitBufferOf(0x07)
+    bitBuffer.skip(4)
+    assertEquals(3, bitBuffer.readUVLC())
+    assertEquals(7, bitBuffer.bufferPosition)
+  }
+
+  @Test
+  fun `GIVEN a stream with emulation prevention bytes WHEN extractRbsp THEN remove them`() {
+    val source = ByteBuffer.wrap(byteArrayOf(0xAA.toByte(), 0xBB.toByte(), 0x00, 0x00, 0x03, 0x01))
+    val rbsp = BitBuffer.extractRbsp(source, 1)
+    val out = ByteArray(rbsp.remaining())
+    rbsp.get(out)
+    assertArrayEquals(byteArrayOf(0xAA.toByte(), 0xBB.toByte(), 0x00, 0x00, 0x01), out)
+  }
+
+  @Test
+  fun `GIVEN a stream without emulation bytes WHEN extractRbsp THEN keep it untouched`() {
+    val source = ByteBuffer.wrap(byteArrayOf(0x01, 0x02, 0x03))
+    val rbsp = BitBuffer.extractRbsp(source, 0)
+    val out = ByteArray(rbsp.remaining())
+    rbsp.get(out)
+    assertArrayEquals(byteArrayOf(0x01, 0x02, 0x03), out)
   }
-}
+}