buffer/reader.go

package buffer

import (
	"errors"
	"io"
	"unicode/utf8"
)

type reader struct {
	options      Options
	in           io.Reader
	segments     [][]byte
	offset       int
	len          int
	lastSegStart int
	err          error
}

func (r *reader) fillSegment() int {
	var n int
	seg := r.segments[len(r.segments)-1]
	start := r.offset + r.len - r.lastSegStart
	n, r.err = r.in.Read(seg[start:len(seg)])
	if n == 0 && r.err == nil {
		n, r.err = r.in.Read(seg[start:len(seg)])
	}

	r.len += n
	return n
}

func (r reader) segmentFull() bool {
	if len(r.segments) == 0 {
		return true
	}

	segmentPosition := r.offset + r.len - r.lastSegStart
	lastSeg := r.segments[len(r.segments)-1]
	return segmentPosition == len(lastSeg)
}

func (r *reader) allocate() {
	segment, err := r.options.BufferPool.Get()
	if err != nil {
		r.err = err
		return
	}

	if len(segment) == 0 {
		r.err = ErrZeroAllocation
		return
	}

	if len(r.segments) > 0 {
		lastSeg := r.segments[len(r.segments)-1]
		r.lastSegStart += len(lastSeg)
	}

	r.segments = append(r.segments, segment)
}

func (r *reader) fill(to int) int {
	var n int
	for {
		if r.err != nil {
			return n
		}

		if r.len >= to {
			return n
		}

		if r.segmentFull() {
			r.allocate()
			continue
		}

		fn := r.fillSegment()
		n += fn
		if fn == 0 || !r.segmentFull() {
			return n
		}
	}
}

func (r reader) copy(p []byte) (n int) {
	var seg, segStart int
	offset := r.offset
	for {
		if len(p) == 0 {
			return
		}

		if seg == len(r.segments) {
			return
		}

		rl := r.offset + r.len - segStart
		if rl > len(r.segments[seg]) {
			rl = len(r.segments[seg])
		}

		ni := copy(p, r.segments[seg][offset:rl])
		n += ni
		p = p[ni:]
		segStart += len(r.segments[seg])
		seg++
		offset = 0
	}
}

func (r *reader) consume(n int) {
	r.offset += n
	r.len -= n
	seg := 0
	if len(r.options.InitialSegment) > 0 {
		seg = 1
	}

	for {
		if len(r.segments) <= 1 {
			break
		}

		segEnd := len(r.segments[seg])
		if seg == 1 {
			segEnd += len(r.segments[0])
		}

		if r.offset < segEnd {
			break
		}

		r.offset -= len(r.segments[seg])
		r.lastSegStart -= len(r.segments[seg])
		r.options.BufferPool.Put(r.segments[seg])
		r.segments = r.segments[seg+1:]
	}

	if len(r.segments) == 1 {
		copy(r.segments[0], r.segments[0][r.offset:r.offset+r.len])
		r.offset = 0
	}
}

func (r *reader) free() {
	if len(r.segments) > 0 && len(r.options.InitialSegment) > 0 {
		r.segments = r.segments[1:]
	}

	for {
		if len(r.segments) == 0 {
			break
		}

		r.options.BufferPool.Put(r.segments[0])
		r.segments[0] = nil
		r.segments = r.segments[1:]
	}

	r.segments = nil
	r.len = 0
	r.options.BufferPool = nil
	if c, ok := r.in.(interface{ close() error }); ok {
		if err := c.close(); err != nil {
			r.err = errors.Join(r.err, err)
		}
	}

	r.in = nil
	if r.err == nil {
		r.err = ErrAbort
	}
}

func (r reader) findSegmentUp(i int) (int, int) {
	var seg, segStart int
	for {
		if r.offset+i >= segStart && r.offset+i < segStart+len(r.segments[seg]) {
			return seg, segStart
		}

		segStart += len(r.segments[seg])
		seg++
	}
}

func (r reader) findSegmentDown(i int) (int, int) {
	seg := len(r.segments) - 1
	segStart := r.lastSegStart
	for {
		if r.offset+i >= segStart && r.offset+i < segStart+len(r.segments[seg]) {
			return seg, segStart
		}

		seg--
		segStart -= len(r.segments[seg])
	}
}

func match(p0, p1 []byte) (bool, int) {
	var i int
	for {
		if len(p0) == i || len(p1) == i {
			return true, len(p0) - len(p1)
		}

		if p0[i] != p1[i] {
			return false, len(p0) - len(p1)
		}

		i++
	}
}

func (r *reader) fillToDelimiter(delimiter []byte, max int) (int, bool) {
	var i, seg, segStart int
	d := delimiter
	for {
		if i+len(d) > max {
			return 0, false
		}

		if r.len < i+len(d) {
			fn := r.fill(i + len(d))
			if fn == 0 {
				return 0, false
			}
		}

		if r.offset+i >= segStart+len(r.segments[seg]) {
			segStart += len(r.segments[seg])
			seg++
		}

		first := r.offset + i - segStart
		last := first + len(d)
		if last > len(r.segments[seg]) {
			last = len(r.segments[seg])
		}

		m, c := match(d, r.segments[seg][first:last])
		if !m {
			i += 1 - len(delimiter) + len(d)
			d = delimiter
			if i < r.len {
				seg, segStart = r.findSegmentDown(i)
			}

			continue
		}

		if c > 0 {
			c = len(d) - c
			i += c
			d = d[c:]
			continue
		}

		return i + len(d), true
	}
}

func (r reader) view(i, max int) []byte {
	if len(r.segments) == 0 {
		return nil
	}

	if max == 0 {
		return nil
	}

	if i >= r.len {
		return nil
	}

	if i+max > r.len {
		max = r.len - i
	}

	var b []byte
	seg, segStart := r.findSegmentUp(i)
	for {
		if len(b) == max {
			return b
		}

		start := r.offset - segStart + i
		l := start + max - len(b)
		if l > len(r.segments[seg]) {
			l = len(r.segments[seg])
		}

		if len(b) == 0 {
			b = r.segments[seg][start:l]
		} else {
			b = append(b, r.segments[seg][start:l]...)
		}

		i += l - start
		if r.offset+i == segStart+len(r.segments[seg]) {
			segStart += len(r.segments[seg])
			seg++
		}
	}
}

func (r *reader) read(p []byte) (int, error) {
	if r.err != nil && r.len == 0 {
		return 0, r.err
	}

	if len(p) == 0 {
		return 0, nil
	}

	if len(r.segments) == 0 {
		r.allocate()
		if r.err != nil {
			return 0, r.err
		}
	}

	if r.len == 0 {
		r.fillSegment()
	}

	n := r.copy(p)
	r.consume(n)
	if r.err != nil && r.len == 0 {
		r.free()
		if n == 0 {
			return 0, r.err
		}
	}

	return n, nil
}

func (r *reader) readBytes(delimiter []byte, max int) ([]byte, error) {
	if r.err != nil && r.len == 0 {
		return nil, r.err
	}

	if len(delimiter) == 0 {
		return nil, nil
	}

	var (
		p []byte
		n int
	)

	l, ok := r.fillToDelimiter(delimiter, max)
	if ok {
		p = make([]byte, l)
		n = r.copy(p)
		r.consume(n)
	}

	if r.err != nil && r.len == 0 {
		r.free()
		if n == 0 {
			return nil, r.err
		}
	}

	return p, nil
}

func (r *reader) readUTF8(max int) ([]rune, int, error) {
	if r.err != nil && r.len == 0 {
		return nil, 0, r.err
	}

	var (
		runes []rune
		n     int
	)

	for {
		if len(runes) == max {
			break
		}

		var zeroRead bool
		b := r.view(n, 4)
		for len(b) < 4 && !utf8.FullRune(b) {
			fn := r.fill(n + 4)
			zeroRead = fn == 0
			if zeroRead {
				break
			}

			b = r.view(n, 4)
		}

		if len(b) == 0 {
			break
		}

		if zeroRead && r.err == nil {
			break
		}

		rn, s := utf8.DecodeRune(b)
		if rn == utf8.RuneError && s == 1 && len(runes) == 0 {
			n = 1
			break
		}

		if rn == utf8.RuneError && s == 1 {
			break
		}

		runes = append(runes, rn)
		n += s
	}

	r.consume(n)
	if r.err != nil && r.len == 0 {
		r.free()
		if n == 0 {
			return nil, 0, r.err
		}
	}

	return runes, n, nil
}

func (r *reader) peek(max int) ([]byte, error) {
	if r.err != nil && r.len == 0 {
		return nil, r.err
	}

	for r.len < max {
		if fn := r.fill(max); fn == 0 {
			break
		}
	}

	v := r.view(0, max)
	if r.err != nil && r.len == 0 {
		r.free()
		if len(v) == 0 {
			return nil, r.err
		}
	}

	return v, nil
}

func (r reader) buffered() []byte {
	return r.view(0, r.len)
}

func (r *reader) writeTo(w io.Writer) (int64, error) {
	if errors.Is(r.err, io.EOF) && r.len == 0 {
		return 0, nil
	}

	if r.err != nil && r.len == 0 {
		return 0, r.err
	}

	var (
		n    int64
		werr error
	)

	for {
		if r.len == 0 {
			if r.err != nil {
				break
			}

			if len(r.segments) == 0 {
				r.allocate()
				if r.err != nil {
					break
				}
			}

			fn := r.fillSegment()
			if fn == 0 && r.err == nil {
				return n, io.ErrNoProgress
			}
		}

		wl := len(r.segments[0]) - r.offset
		if wl > r.len {
			wl = r.len
		}

		ni, err := w.Write(r.segments[0][r.offset : r.offset+wl])
		r.consume(ni)
		n += int64(ni)
		if err != nil {
			werr = err
			break
		}

		if ni < wl {
			werr = io.ErrShortWrite
			break
		}
	}

	if r.err != nil && r.len == 0 {
		r.free()
	}

	if werr != nil {
		return n, werr
	}

	if r.err != nil && !errors.Is(r.err, io.EOF) && r.len == 0 {
		return n, r.err
	}

	return n, nil
}