update bit reader to support signed reads

This commit is contained in:
Alex Yatskov 2019-02-09 19:12:04 -08:00
parent 08d71f5bc4
commit c9c5ed09ac
2 changed files with 56 additions and 73 deletions

View File

@ -55,14 +55,7 @@ func (r *BitReader) ReadInt64(count int) (int64, error) {
return 0, err
}
if count > 0 {
valueMasked := value &^ (1 << uint(count-1))
if valueMasked != value {
return -int64(valueMasked), nil
}
}
return int64(value), nil
return twosComplement(value, count), 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
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
return readBits(buffer, bitOffset, count), 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))
}
}

View File

@ -17,50 +17,26 @@ var data = []byte{
0xEF, // 11101111
}
func testRead(test *testing.T, reader *BitReader, count int, expected uint64) {
value, err := reader.ReadUint64(count)
if err != nil {
log.Printf("error: %s\n", err.Error())
test.Fail()
} else if value != expected {
log.Printf("value: %.16x, expected: %.16x\n", value, expected)
test.Fail()
func TestBitReader(t *testing.T) {
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())
t.Fail()
} else if value != expected {
log.Printf("value: %.16x, expected: %.16x\n", value, expected)
t.Fail()
}
}
}
func TestUnsigned(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)&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)
test(0, 0x00)
test(8, 0x01)
test(16, 0x4523)
test(3, 0x67&0x07)
test(13, 0x8967>>3)
test(13, 0xcdab&0x1fff)
test(2, (0xcdab>>13)&3)
test(9, 0xefcdab>>15)
}