buffer/writer.go

package buffer

import (
	"errors"
	"io"
)

type writer struct {
	out         io.Writer
	options     Options
	buffer      []byte
	offset, len int
	err         error
}

var errClosed = errors.New("buffer closed")

func (w *writer) write(p []byte) (int, error) {
	var n int
	for {
		if w.err != nil {
			return n, w.err
		}

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

		if len(w.buffer) == 0 {
			w.buffer, w.err = w.options.Pool.Get()
			if len(w.buffer) == 0 && w.err == nil {
				w.err = ErrZeroAllocation
			}

			continue
		}

		if w.offset+w.len == len(w.buffer) {
			w.flush()
			continue
		}

		ni := copy(w.buffer[w.offset+w.len:], p)
		w.len += ni
		p = p[ni:]
		n += ni
	}
}

func (w *writer) readFrom(r io.Reader) (int64, error) {
	var n int64
	for {
		if w.err != nil {
			return n, w.err
		}

		if len(w.buffer) == 0 {
			w.buffer, w.err = w.options.Pool.Get()
			if len(w.buffer) == 0 && w.err == nil {
				w.err = ErrZeroAllocation
			}

			continue
		}

		if w.offset+w.len == len(w.buffer) {
			w.flush()
			continue
		}

		ni, rerr := r.Read(w.buffer[w.offset+w.len:])
		if ni == 0 && rerr == nil {
			ni, rerr = r.Read(w.buffer[w.offset+w.len:])
		}

		if ni == 0 && rerr == nil {
			rerr = io.ErrNoProgress
		}

		w.len += ni
		n += int64(ni)
		if errors.Is(rerr, io.EOF) {
			w.flush()
			return n, w.err
		}

		if rerr != nil {
			return n, rerr
		}
	}
}

func (w *writer) flush() error {
	var zeroWrite bool
	for {
		if w.err != nil {
			return w.err
		}

		if w.len == 0 {
			w.offset = 0
			return nil
		}

		var n int
		n, w.err = w.out.Write(w.buffer[w.offset : w.offset+w.len])
		if n == 0 && w.err == nil && zeroWrite {
			w.err = io.ErrShortWrite
		}

		zeroWrite = n == 0 && w.err == nil
		w.offset += n
		w.len -= n
	}
}

func (w *writer) close() error {
	newErr := w.err == nil
	w.flush()
	if len(w.buffer) > 0 {
		w.options.Pool.Put(w.buffer)
		w.buffer = nil
	}

	if newErr && w.err != nil {
		return w.err
	}

	if w.err == nil {
		w.err = errClosed
	}

	return nil
}