go_goutils/iox/funcs.go

package iox

import (
	`context`
	`io`
)

/*
CopyBufN is a mix between [io.CopyN] and [io.CopyBuffer].

Despite what the docs may suggest, [io.CopyN] does NOT *read* n bytes from src AND write n bytes to dst.
Instead, it always reads 32 KiB from src, and writes n bytes to dst.

There are cases where this is dreadfully undesired.

One can, of course, use [io.CopyBuffer], but this is a bit annoying since you then have to provide a buffer yourself.

This convenience-wraps [io.CopyBuffer] to have a similar signature to [io.CopyN] but properly uses n for both reading and writing.
*/
func CopyBufN(dst io.Writer, src io.Reader, n int64) (written int64, err error) {

	var b []byte

	if n <= 0 {
		err = ErrBufTooSmall
		return
	}

	b = make([]byte, n)

	written, err = io.CopyBuffer(dst, src, b)

	return
}

// CopyCtxBufN copies from `src` to `dst`, `n` bytes at a time, interruptible by `ctx`.
func CopyCtxBufN(ctx context.Context, dst io.Writer, src io.Reader, n int64) (written int64, err error) {

	var nr int
	var nw int
	var end bool
	var buf []byte

	if ctx == nil {
		err = ErrNilCtx
		return
	}
	if n <= 0 {
		err = ErrBufTooSmall
		return
	}

endCopy:
	for {
		select {
		case <-ctx.Done():
			err = ctx.Err()
			return
		default:
			buf = make([]byte, n)
			nr, err = src.Read(buf)
			if err == io.EOF {
				err = nil
				end = true
			} else if err != nil {
				return
			}
			buf = buf[:nr]

			if nw, err = dst.Write(buf); err != nil {
				written += int64(nw)
				return
			}
			written += int64(nw)
			if len(buf) != nw {
				err = io.ErrShortWrite
				return
			}
			if end {
				break endCopy
			}
		}
	}

	return
}

/*
CopyBufWith allows for specifying a buffer allocator function, otherwise acts as [CopyBufN].

bufFunc *MUST NOT* return a nil or len == 0 buffer. [ErrBufTooSmall] will be returned if it does.

This uses a fixed buffer size from a single call to `bufFunc`.
If you need something with dynamic buffer sizing according to some state, use [CopyBufWithDynamic] instead.
(Note that CopyBufWithDynamic is generally a little slower, but it should only be noticeable on very large amounts of data.)
*/
func CopyBufWith(dst io.Writer, src io.Reader, bufFunc func() (b []byte)) (written int64, err error) {

	var buf []byte = bufFunc()

	if buf == nil || len(buf) == 0 {
		err = ErrBufTooSmall
		return
	}

	written, err = io.CopyBuffer(dst, src, buf)

	return
}

/*
CopyBufWithDynamic is like [CopyBufWith] except it will call bufFunc after each previous buffer is written.

That is to say (using a particularly contrived example):

	import time

	func dynBuf() (b []byte) {

		var t time.Time = time.Now()

		b = make([]byte, t.Seconds())

		return
	}

Then:

	CopyBufWithDynamic(w, r, dynBuf)

will use a buffer sized to the seconds of the time it reads in/writes out the next buffer, whereas with [CopyBufWith]:

	CopyBufWith(w, r, dynBuf)

would use a *fixed* buffer size of whatever the seconds was equal to at the time of the *first call* to dynBuf.

`src` MUST return an [io.EOF] when its end is reached, but (as per e.g. [io.CopyBuffer]) the io.EOF error will not
be returned from CopyBufWithDynamic. (Any/all other errors encountered will be returned, however, and copying will
immediately cease.)
*/
func CopyBufWithDynamic(dst io.Writer, src io.Reader, bufFunc func() (b []byte)) (written int64, err error) {

	var nr int
	var nw int
	var end bool
	var buf []byte

	for {
		buf = bufFunc()
		if buf == nil || len(buf) == 0 {
			err = ErrBufTooSmall
			return
		}
		nr, err = src.Read(buf)
		if err == io.EOF {
			err = nil
			end = true
		} else if err != nil {
			return
		}
		buf = buf[:nr]

		if nw, err = dst.Write(buf); err != nil {
			written += int64(nw)
			return
		}
		written += int64(nw)
		if len(buf) != nw {
			err = ErrShortWrite
			return
		}
		if end {
			break
		}
	}

	return
}

// NewChunker returns a [ChunkLocker] ready to use.
func NewChunker(chunkSize uint) (c *ChunkLocker, err error) {

	c = &ChunkLocker{}
	err = c.SetChunkLen(chunkSize)

	return
}

// NewCtxIO returns a [CtxIO].
func NewCtxIO(ctx context.Context, r io.Reader, w io.Writer, chunkSize uint) (c *CtxIO, err error) {

	if r == nil {
		err = ErrNilReader
		return
	}
	if w == nil {
		err = ErrNilWriter
		return
	}

	if chunkSize == 0 {
		err = ErrInvalidChunkSize
		return
	}

	if ctx == nil {
		err = ErrNilCtx
		return
	}

	c = &CtxIO{
		r: r,
		w: w,
		l: ChunkLocker{
			chunkLen: chunkSize,
		},
		ctx: ctx,
	}

	return
}

/*
	NewXIO returns a nil [XIO].

	A weird "feature" of Golang is that a nil XIO is perfectly fine to use;
	it's completely stateless and only has pointer receivers that only work with passed in
	values so `new(XIO)` is completely unnecessary (as is NewXCopier).
	In other words, this works fine:

		var xc *iox.XIO

		if n, err = xc.Copy(w, r); err != nil {
			return
		}

	This function is just to maintain cleaner-looking code if you should so need it,
	or want an XIO without declaring one:

		if n, err = iox.NewXCopier().Copy(w, r); err != nil {
			return
		}
*/
func NewXIO() (x *XIO) {
	// No-op lel
	return
}