reader.go

// Copyright 2013 Benoît Amiaux. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

/*
Package iobit provides primitives for reading & writing bits

The main purpose of this library is to remove the need to write
custom bit-masks when reading or writing bitstreams, and to ease
maintenance. This is true especially when you need to read/write
data which is not aligned on bytes.

Errors are sticky so you can check for errors after a chunk of
meaningful work rather than after every operation.

For example, with iobit you can read an MPEG-TS PCR like this:

	r := iobit.NewReader(buffer)
	base := r.Uint64(33)     // PCR base is 33-bits
	r.Skip(6)                // 6-bits are reserved
	extension := r.Uint64(9) // PCR extension is 9-bits

instead of:

	base  = uint64(buffer[0]) << 25
	base |= uint64(buffer[1]) << 17
	base |= uint64(buffer[2]) << 9
	base |= uint64(buffer[3]) << 1
	base |= buffer[4] >> 7
	extension := uint16(buffer[4] & 0x1) << 8
	extension |= buffer[5]

and write it like this:

	w := iobit.NewWriter(buffer)
	w.PutUint64(33, base)
	w.PutUint32(6, 0)
	w.PutUint32(9, extension)
*/
package iobit

// Reader wraps a raw byte array and provides multiple methods to read and
// skip data bit-by-bit.
// Its methods don't return the usual error as it is too expensive.
// Instead, read errors can be checked with the Check() method
type Reader struct {
	src  []byte
	buf  uint64 // internal bit buffer (MSB aligned)
	off  int    // next byte offset to load from src
	bit  int    // number of bits currently available in buf
	idx  uint   // total bits consumed (for position tracking)
	size uint   // source size in bytes
}

// NewReader returns a new reader reading from <src> byte array.
func NewReader(src []byte) Reader {
	return Reader{
		src:  src,
		size: uint(len(src)),
	}
}

func min(a, b uint) uint {
	if a > b {
		return b
	}
	return a
}

func bswap16(v uint16) uint16 {
	return v>>8 | v<<8
}

func bswap32(val uint32) uint32 {
	return uint32(bswap16(uint16(val>>16))) | uint32(bswap16(uint16(val)))<<16
}

// fill pulls bytes from src to keep the 64-bit buffer as full as possible.
func (r *Reader) fill() {
	for r.bit <= 56 && r.off < int(r.size) {
		r.buf |= uint64(r.src[r.off]) << uint(56-r.bit)
		r.off++
		r.bit += 8
	}
}

// readBits reads up to 64 bits from the buffer.
func (r *Reader) readBits(bits uint) uint64 {
	if bits == 0 {
		return 0
	}

	r.fill()

	// Fast path: buffer has enough bits
	if uint(r.bit) >= bits {
		val := r.buf >> (64 - bits)
		r.buf <<= bits
		r.bit -= int(bits)
		r.idx += bits
		return val
	}

	// Slow path: split read across two fills (needed for >56 bit reads
	// when not byte-aligned)
	have := uint(r.bit)
	if have > 0 {
		hi := r.buf >> (64 - have)
		r.buf = 0
		r.bit = 0
		remaining := bits - have
		r.fill()
		if uint(r.bit) >= remaining {
			lo := r.buf >> (64 - remaining)
			r.buf <<= remaining
			r.bit -= int(remaining)
			r.idx += bits
			return hi<<remaining | lo
		}
	}

	// Underflow: not enough source data - advance past end for Error()
	r.idx += bits
	r.buf = 0
	r.bit = 0
	return 0
}

func signExtend(val uint64, bits uint) int64 {
	if bits == 0 {
		return 0
	}
	shift := 64 - bits
	return int64(val<<shift) >> shift
}

// Bit reads the next bit as a boolean.
func (r *Reader) Bit() bool {
	return r.readBits(1) != 0
}

// Byte reads one byte.
func (r *Reader) Byte() uint8 {
	return uint8(r.readBits(8))
}

// Uint8 reads up to 8 unsigned bits in big-endian order.
func (r *Reader) Uint8(bits uint) uint8 {
	return uint8(r.readBits(bits))
}

// Int8 reads up to 8 signed bits in big-endian order.
func (r *Reader) Int8(bits uint) int8 {
	return int8(signExtend(r.readBits(bits), bits))
}

// Be16 reads 16 unsigned bits in big-endian order.
func (r *Reader) Be16() uint16 {
	return uint16(r.readBits(16))
}

// Uint16 reads up to 16 unsigned bits in big-endian order.
func (r *Reader) Uint16(bits uint) uint16 {
	return uint16(r.readBits(bits))
}

// Int16 reads up to 16 signed bits in big-endian order.
func (r *Reader) Int16(bits uint) int16 {
	return int16(signExtend(r.readBits(bits), bits))
}

// Le16 reads 16 unsigned bits in litle-endian order.
func (r *Reader) Le16() uint16 {
	return bswap16(uint16(r.readBits(16)))
}

// Be32 reads 32 unsigned bits in big-endian order.
func (r *Reader) Be32() uint32 {
	return uint32(r.readBits(32))
}

// Uint32 reads up to 32 unsigned bits in big-endian order.
func (r *Reader) Uint32(bits uint) uint32 {
	return uint32(r.readBits(bits))
}

// Int32 reads up to 32 signed bits in big-endian order.
func (r *Reader) Int32(bits uint) int32 {
	return int32(signExtend(r.readBits(bits), bits))
}

// Le32 reads 32 unsigned bits in little-endian order.
func (r *Reader) Le32() uint32 {
	return bswap32(uint32(r.readBits(32)))
}

// Be64 reads 64 unsigned bits in big-endian order.
func (r *Reader) Be64() uint64 {
	return r.readBits(64)
}

// Le64 reads 64 unsigned bits in little-endian order.
func (r *Reader) Le64() uint64 {
	v := r.readBits(64)
	return uint64(bswap32(uint32(v)))<<32 | uint64(bswap32(uint32(v>>32)))
}

// Uint64 reads up to 64 unsigned bits in big-endian order.
func (r *Reader) Uint64(bits uint) uint64 {
	return r.readBits(bits)
}

// Int64 reads up to 64 signed bits in big-endian order.
func (r *Reader) Int64(bits uint) int64 {
	return signExtend(r.readBits(bits), bits)
}

// Bytes returns a byte-array of input size.
func (r *Reader) Bytes(size int) []byte {
	d := r.LeftBytes()
	if size > len(d) {
		size = len(d)
	}
	r.Skip(uint(size) * 8)
	return d[:size]
}

// String returns a string of input size.
func (r *Reader) String(size int) string {
	return string(r.Bytes(size))
}

// Peek returns a reader copy.
// Useful to read data without advancing the original reader.
func (r *Reader) Peek() *Reader {
	p := *r
	return &p
}

// Skip skips n bits.
func (r *Reader) Skip(bits uint) {
	// Fast path: skip within buffer
	if uint(r.bit) >= bits {
		r.buf <<= bits
		r.bit -= int(bits)
		r.idx += bits
		return
	}
	// Slow path: reposition
	r.idx += bits
	r.buf = 0
	r.bit = 0
	r.off = int(r.idx >> 3)
	subBits := r.idx & 7
	if subBits > 0 && r.off < int(r.size) {
		r.fill()
		if uint(r.bit) >= subBits {
			r.buf <<= subBits
			r.bit -= int(subBits)
		}
	}
}

// At returns the current reader position in bits.
func (r *Reader) At() uint {
	return r.idx
}

// LeftBits returns the number of bits left to read.
func (r *Reader) LeftBits() uint {
	return r.size<<3 - min(r.idx, r.size<<3)
}

// LeftBytes returns a slice of the contents of the unread reader portion.
// Note that this slice is byte aligned even if the reader is not.
func (r *Reader) LeftBytes() []byte {
	skip := r.idx >> 3
	if skip >= r.size {
		return r.src[:0]
	}
	return r.src[skip:r.size]
}

// Reset resets the reader to its initial position.
func (r *Reader) Reset() {
	r.idx = 0
	r.off = 0
	r.bit = 0
	r.buf = 0
}

// Error returns whether the reader encountered an error.
func (r *Reader) Error() error {
	if r.idx > r.size<<3 {
		return ErrOverflow
	}
	return nil
}