reader.go

package shared

import (
	"bytes"
	"errors"
	"io"
)

const (
	scratchByteArrayLen = 32
)

var (
	zeroByteSlice = []byte{}[:0:0]
)

var (
	_ SlickReader = &SlickReaderStream{}
	_ SlickReader = &SlickReaderSlice{}
)

func NewReader(r io.Reader) SlickReader {
	return &SlickReaderStream{br: &readerToScanner{r: r}}
}

func NewBytesReader(buf *bytes.Buffer) SlickReader {
	return &SlickReaderStream{br: buf}
}

func NewSliceReader(b []byte) SlickReader {
	return &SlickReaderSlice{b: b}
}

// SlickReader is a hybrid of reader and buffer interfaces with methods giving
// specific attention to the performance needs found in a decoder.
// Implementations cover io.Reader as well as []byte directly.
//
// In particular, it allows:
//
//   - returning byte-slices with zero-copying (you were warned!) when possible
//   - returning byte-slices for short reads which will be reused (you were warned!)
//   - putting a 'track' point in the buffer, and later yielding all those bytes at once
//   - counting the number of bytes read (for use in parser error messages, mainly)
//
// All of these shortcuts mean correct usage is essential to avoid unexpected behaviors,
// but in return allow avoiding many, many common sources of memory allocations in a parser.
type SlickReader interface {

	// Read n bytes into a byte slice which may be shared and must not be reused
	// After any additional calls to this reader.
	// Readnzc will use the implementation scratch buffer if possible,
	// i.e. n < len(scratchbuf), or may return a view of the []byte being decoded from.
	// Requesting a zero length read will return `zeroByteSlice`, a len-zero cap-zero slice.
	Readnzc(n int) ([]byte, error)

	// Read n bytes into a new byte slice.
	// If zero-copy views into existing buffers are acceptable (e.g. you know you
	// won't later mutate, reference or expose this memory again), prefer `Readnzc`.
	// If you already have an existing slice of sufficient size to reuse, prefer `Readb`.
	// Requesting a zero length read will return `zeroByteSlice`, a len-zero cap-zero slice.
	Readn(n int) ([]byte, error)

	// Read `len(b)` bytes into the given slice, starting at its beginning,
	// overwriting all values, and disregarding any extra capacity.
	Readb(b []byte) error

	Readn1() (uint8, error)
	Unreadn1()
	NumRead() int // number of bytes read
	Track()
	StopTrack() []byte
}

// SlickReaderStream is a SlickReader that reads off an io.Reader.
// Initialize it by wrapping an ioDecByteScanner around your io.Reader and dumping it in.
// While this implementation does use some internal buffers, it's still advisable
// to use a buffered reader to avoid small reads for any external IO like disk or network.
type SlickReaderStream struct {
	br         readerScanner
	scratch    [scratchByteArrayLen]byte // temp byte array re-used internally for efficiency during read.
	n          int                       // num read
	tracking   []byte                    // tracking bytes read
	isTracking bool
}

func (z *SlickReaderStream) NumRead() int {
	return z.n
}

func (z *SlickReaderStream) Readnzc(n int) (bs []byte, err error) {
	if n == 0 {
		return zeroByteSlice, nil
	}
	if n < len(z.scratch) {
		bs = z.scratch[:n]
	} else {
		bs = make([]byte, n)
	}
	err = z.Readb(bs)
	return
}

func (z *SlickReaderStream) Readn(n int) (bs []byte, err error) {
	if n == 0 {
		return zeroByteSlice, nil
	}
	bs = make([]byte, n)
	err = z.Readb(bs)
	return
}

func (z *SlickReaderStream) Readb(bs []byte) error {
	if len(bs) == 0 {
		return nil
	}
	n, err := io.ReadAtLeast(z.br, bs, len(bs))
	z.n += n
	if z.isTracking {
		z.tracking = append(z.tracking, bs...)
	}
	return err
}

func (z *SlickReaderStream) Readn1() (b uint8, err error) {
	b, err = z.br.ReadByte()
	if err != nil {
		return
	}
	z.n++
	if z.isTracking {
		z.tracking = append(z.tracking, b)
	}
	return
}

func (z *SlickReaderStream) Unreadn1() {
	err := z.br.UnreadByte()
	if err != nil {
		panic(err)
	}
	z.n--
	if z.isTracking {
		if l := len(z.tracking) - 1; l >= 0 {
			z.tracking = z.tracking[:l]
		}
	}
}

func (z *SlickReaderStream) Track() {
	if z.tracking != nil {
		z.tracking = z.tracking[:0]
	}
	z.isTracking = true
}

func (z *SlickReaderStream) StopTrack() (bs []byte) {
	z.isTracking = false
	return z.tracking
}

// SlickReaderSlice implements SlickReader by reading a byte slice directly.
// Often this means the zero-copy methods can simply return subslices.
type SlickReaderSlice struct {
	b []byte // data
	c int    // cursor
	a int    // available
	t int    // track start
}

func (z *SlickReaderSlice) reset(in []byte) {
	z.b = in
	z.a = len(in)
	z.c = 0
	z.t = 0
}

func (z *SlickReaderSlice) NumRead() int {
	return z.c
}

func (z *SlickReaderSlice) Unreadn1() {
	if z.c == 0 || len(z.b) == 0 {
		panic(errors.New("cannot unread last byte read"))
	}
	z.c--
	z.a++
	return
}

func (z *SlickReaderSlice) Readnzc(n int) (bs []byte, err error) {
	if n == 0 {
		return zeroByteSlice, nil
	} else if z.a == 0 {
		return zeroByteSlice, io.EOF
	} else if n > z.a {
		return zeroByteSlice, io.ErrUnexpectedEOF
	} else {
		c0 := z.c
		z.c = c0 + n
		z.a = z.a - n
		bs = z.b[c0:z.c]
	}
	return
}

func (z *SlickReaderSlice) Readn(n int) (bs []byte, err error) {
	if n == 0 {
		return zeroByteSlice, nil
	}
	bs = make([]byte, n)
	err = z.Readb(bs)
	return
}

func (z *SlickReaderSlice) Readn1() (v uint8, err error) {
	if z.a == 0 {
		panic(io.EOF)
	}
	v = z.b[z.c]
	z.c++
	z.a--
	return
}

func (z *SlickReaderSlice) Readb(bs []byte) error {
	bs2, err := z.Readnzc(len(bs))
	copy(bs, bs2)
	return err
}

func (z *SlickReaderSlice) Track() {
	z.t = z.c
}

func (z *SlickReaderSlice) StopTrack() (bs []byte) {
	return z.b[z.t:z.c]
}

// conjoin the io.Reader and io.ByteScanner interfaces.
type readerScanner interface {
	io.Reader
	io.ByteScanner
}

// readerToScanner decorates an `io.Reader` with all the methods to also
// fulfill the `io.ByteScanner` interface.
type readerToScanner struct {
	r  io.Reader
	l  byte    // last byte
	ls byte    // last byte status. 0: init-canDoNothing, 1: canRead, 2: canUnread
	b  [1]byte // tiny buffer for reading single bytes
}

func (z *readerToScanner) Read(p []byte) (n int, err error) {
	var firstByte bool
	if z.ls == 1 {
		z.ls = 2
		p[0] = z.l
		if len(p) == 1 {
			n = 1
			return
		}
		firstByte = true
		p = p[1:]
	}
	n, err = z.r.Read(p)
	if n > 0 {
		if err == io.EOF && n == len(p) {
			err = nil // read was successful, so postpone EOF (till next time)
		}
		z.l = p[n-1]
		z.ls = 2
	}
	if firstByte {
		n++
	}
	return
}

func (z *readerToScanner) ReadByte() (c byte, err error) {
	n, err := z.Read(z.b[:])
	if n == 1 {
		c = z.b[0]
		if err == io.EOF {
			err = nil // read was successful, so postpone EOF (till next time)
		}
	}
	return
}

func (z *readerToScanner) UnreadByte() (err error) {
	x := z.ls
	if x == 0 {
		err = errors.New("cannot unread - nothing has been read")
	} else if x == 1 {
		err = errors.New("cannot unread - last byte has not been read")
	} else if x == 2 {
		z.ls = 1
	}
	return
}