diff --git a/streaming/bitreader.go b/streaming/bitreader.go index dbf7f71..945b6f1 100644 --- a/streaming/bitreader.go +++ b/streaming/bitreader.go @@ -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)) + } +} diff --git a/streaming/streaming_test.go b/streaming/streaming_test.go index 5985743..df8f35e 100644 --- a/streaming/streaming_test.go +++ b/streaming/streaming_test.go @@ -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) }