update bit reader to support signed reads

streaming/bitreader.go

@@ -55,14 +55,7 @@ func (r *BitReader) ReadInt64(count int) (int64, error) {
 		return 0, err
 	}
 
-	if count > 0 {
+	return twosComplement(value, count), nil
-		valueMasked := value &^ (1 << uint(count-1))
-		if valueMasked != value {
-			return -int64(valueMasked), nil
-		}
-	}
-
-	return int64(value), nil
 }
 
 func (r *BitReader) ReadUint64(count int) (uint64, error) {
@@ -71,26 +64,7 @@ func (r *BitReader) ReadUint64(count int) (uint64, error) {
 		return 0, err
 	}
 
-	var result uint64
+	return readBits(buffer, bitOffset, count), nil
-
-	remainder := count
-	for byteOffset := 0; remainder > 0; byteOffset++ {
-		bitsRead := 8 - bitOffset
-		if bitsRead > remainder {
-			bitsRead = remainder
-		}
-
-		bufferByte := buffer[byteOffset]
-		bufferByte >>= uint(bitOffset)
-		bufferByte &= ^(0xff << uint(bitsRead))
-
-		result |= (uint64(bufferByte) << uint(count-remainder))
-
-		remainder -= bitsRead
-		bitOffset = 0
-	}
-
-	return result, nil
 }
 
 func (r *BitReader) readBytes(count int) ([]byte, int, error) {
@@ -128,3 +102,36 @@ func (r *BitReader) readBytes(count int) ([]byte, int, error) {
 
 	return buffer, bitOffsetInByte, nil
 }
+
+func readBits(buffer []byte, bitOffset, count int) uint64 {
+	var result uint64
+
+	remainder := count
+	for byteOffset := 0; remainder > 0; byteOffset++ {
+		bitsRead := 8 - bitOffset
+		if bitsRead > remainder {
+			bitsRead = remainder
+		}
+
+		bufferByte := buffer[byteOffset]
+		bufferByte >>= uint(bitOffset)
+		bufferByte &= ^(0xff << uint(bitsRead))
+
+		result |= (uint64(bufferByte) << uint(count-remainder))
+
+		remainder -= bitsRead
+		bitOffset = 0
+	}
+
+	return result
+}
+
+func twosComplement(value uint64, bits int) int64 {
+	signMask := uint64(1 << uint(bits-1))
+	if value&signMask == 0 {
+		return int64(value &^ signMask)
+	} else {
+		valueMask := ^(^uint64(0) << uint(bits-1))
+		return -int64(valueMask & (^value + 1))
+	}
+}

streaming/streaming_test.go

@@ -17,50 +17,26 @@ var data = []byte{
 	0xEF, // 11101111
 }
 
-func testRead(test *testing.T, reader *BitReader, count int, expected uint64) {
+func TestBitReader(t *testing.T) {
-	value, err := reader.ReadUint64(count)
+	r := NewBitReader(bytes.NewReader(data))
+
+	test := func(count int, expected uint64) {
+		value, err := r.ReadUint64(count)
 		if err != nil {
 			log.Printf("error: %s\n", err.Error())
-		test.Fail()
+			t.Fail()
 		} else if value != expected {
 			log.Printf("value: %.16x, expected: %.16x\n", value, expected)
-		test.Fail()
+			t.Fail()
 		}
 	}
 
-func TestUnsigned(t *testing.T) {
+	test(0, 0x00)
-	r := NewBitReader(bytes.NewReader(data))
+	test(8, 0x01)
-
+	test(16, 0x4523)
-	testRead(t, r, 0, 0x00)
+	test(3, 0x67&0x07)
-	testRead(t, r, 8, 0x01)
+	test(13, 0x8967>>3)
-	testRead(t, r, 16, 0x4523)
+	test(13, 0xcdab&0x1fff)
-	testRead(t, r, 3, 0x67&0x07)
+	test(2, (0xcdab>>13)&3)
-	testRead(t, r, 13, 0x8967>>3)
+	test(9, 0xefcdab>>15)
-	testRead(t, r, 13, 0xcdab&0x1fff)
-	testRead(t, r, 2, (0xcdab>>13)&3)
-	testRead(t, r, 9, 0xefcdab>>15)
-}
-
-func TestUnsignedSmall(t *testing.T) {
-	r := NewBitReader(bytes.NewReader(data))
-
-	testRead(t, r, 0, 0x00)
-	testRead(t, r, 8, 0x01)
-	testRead(t, r, 8, 0x23)
-	testRead(t, r, 8, 0x45)
-	testRead(t, r, 16, 0x8967)
-}
-
-func TestUnsignedTiny(t *testing.T) {
-	r := NewBitReader(bytes.NewReader(data))
-
-	testRead(t, r, 0, 0x00)
-	testRead(t, r, 8, 0x01)
-	testRead(t, r, 16, 0x4523)
-	testRead(t, r, 3, 0x67&0x07)
-	testRead(t, r, 13, 0x8967>>3)
-	testRead(t, r, 13, 0xCDAB&0x1fff)
-	testRead(t, r, 2, (0xcdab>>13)&0x03)
-	testRead(t, r, 5, (0xefcdab>>15)&0x1f)
-	testRead(t, r, 4, (0xefcdab>>20)&0x0f)
 }