niki/vendor/github.com/mailru/easyjson/buffer/pool.go

// Package buffer implements a buffer for serialization, consisting of a chain of []byte-s to

// reduce copying and to allow reuse of individual chunks.

package buffer

import (
	"io"
	"net"
	"sync"
)

// PoolConfig contains configuration for the allocation and reuse strategy.

type PoolConfig struct {
	StartSize int // Minimum chunk size that is allocated.

	PooledSize int // Minimum chunk size that is reused, reusing chunks too small will result in overhead.

	MaxSize int // Maximum chunk size that will be allocated.

}

var config = PoolConfig{

	StartSize: 128,

	PooledSize: 512,

	MaxSize: 32768,
}

// Reuse pool: chunk size -> pool.

var buffers = map[int]*sync.Pool{}

func initBuffers() {

	for l := config.PooledSize; l <= config.MaxSize; l *= 2 {

		buffers[l] = new(sync.Pool)

	}

}

func init() {

	initBuffers()

}

// Init sets up a non-default pooling and allocation strategy. Should be run before serialization is done.

func Init(cfg PoolConfig) {

	config = cfg

	initBuffers()

}

// putBuf puts a chunk to reuse pool if it can be reused.

func putBuf(buf []byte) {

	size := cap(buf)

	if size < config.PooledSize {

		return

	}

	if c := buffers[size]; c != nil {

		c.Put(buf[:0])

	}

}

// getBuf gets a chunk from reuse pool or creates a new one if reuse failed.

func getBuf(size int) []byte {

	if size >= config.PooledSize {

		if c := buffers[size]; c != nil {

			v := c.Get()

			if v != nil {

				return v.([]byte)

			}

		}

	}

	return make([]byte, 0, size)

}

// Buffer is a buffer optimized for serialization without extra copying.

type Buffer struct {

	// Buf is the current chunk that can be used for serialization.

	Buf []byte

	toPool []byte

	bufs [][]byte
}

// EnsureSpace makes sure that the current chunk contains at least s free bytes,

// possibly creating a new chunk.

func (b *Buffer) EnsureSpace(s int) {

	if cap(b.Buf)-len(b.Buf) < s {

		b.ensureSpaceSlow(s)

	}

}

func (b *Buffer) ensureSpaceSlow(s int) {

	l := len(b.Buf)

	if l > 0 {

		if cap(b.toPool) != cap(b.Buf) {

			// Chunk was reallocated, toPool can be pooled.

			putBuf(b.toPool)

		}

		if cap(b.bufs) == 0 {

			b.bufs = make([][]byte, 0, 8)

		}

		b.bufs = append(b.bufs, b.Buf)

		l = cap(b.toPool) * 2

	} else {

		l = config.StartSize

	}

	if l > config.MaxSize {

		l = config.MaxSize

	}

	b.Buf = getBuf(l)

	b.toPool = b.Buf

}

// AppendByte appends a single byte to buffer.

func (b *Buffer) AppendByte(data byte) {

	b.EnsureSpace(1)

	b.Buf = append(b.Buf, data)

}

// AppendBytes appends a byte slice to buffer.

func (b *Buffer) AppendBytes(data []byte) {

	if len(data) <= cap(b.Buf)-len(b.Buf) {

		b.Buf = append(b.Buf, data...) // fast path

	} else {

		b.appendBytesSlow(data)

	}

}

func (b *Buffer) appendBytesSlow(data []byte) {

	for len(data) > 0 {

		b.EnsureSpace(1)

		sz := cap(b.Buf) - len(b.Buf)

		if sz > len(data) {

			sz = len(data)

		}

		b.Buf = append(b.Buf, data[:sz]...)

		data = data[sz:]

	}

}

// AppendString appends a string to buffer.

func (b *Buffer) AppendString(data string) {

	if len(data) <= cap(b.Buf)-len(b.Buf) {

		b.Buf = append(b.Buf, data...) // fast path

	} else {

		b.appendStringSlow(data)

	}

}

func (b *Buffer) appendStringSlow(data string) {

	for len(data) > 0 {

		b.EnsureSpace(1)

		sz := cap(b.Buf) - len(b.Buf)

		if sz > len(data) {

			sz = len(data)

		}

		b.Buf = append(b.Buf, data[:sz]...)

		data = data[sz:]

	}

}

// Size computes the size of a buffer by adding sizes of every chunk.

func (b *Buffer) Size() int {

	size := len(b.Buf)

	for _, buf := range b.bufs {

		size += len(buf)

	}

	return size

}

// DumpTo outputs the contents of a buffer to a writer and resets the buffer.

func (b *Buffer) DumpTo(w io.Writer) (written int, err error) {

	bufs := net.Buffers(b.bufs)

	if len(b.Buf) > 0 {

		bufs = append(bufs, b.Buf)

	}

	n, err := bufs.WriteTo(w)

	for _, buf := range b.bufs {

		putBuf(buf)

	}

	putBuf(b.toPool)

	b.bufs = nil

	b.Buf = nil

	b.toPool = nil

	return int(n), err

}

// BuildBytes creates a single byte slice with all the contents of the buffer. Data is

// copied if it does not fit in a single chunk. You can optionally provide one byte

// slice as argument that it will try to reuse.

func (b *Buffer) BuildBytes(reuse ...[]byte) []byte {

	if len(b.bufs) == 0 {

		ret := b.Buf

		b.toPool = nil

		b.Buf = nil

		return ret

	}

	var ret []byte

	size := b.Size()

	// If we got a buffer as argument and it is big enough, reuse it.

	if len(reuse) == 1 && cap(reuse[0]) >= size {

		ret = reuse[0][:0]

	} else {

		ret = make([]byte, 0, size)

	}

	for _, buf := range b.bufs {

		ret = append(ret, buf...)

		putBuf(buf)

	}

	ret = append(ret, b.Buf...)

	putBuf(b.toPool)

	b.bufs = nil

	b.toPool = nil

	b.Buf = nil

	return ret

}

type readCloser struct {
	offset int

	bufs [][]byte
}

func (r *readCloser) Read(p []byte) (n int, err error) {

	for _, buf := range r.bufs {

		// Copy as much as we can.

		x := copy(p[n:], buf[r.offset:])

		n += x // Increment how much we filled.

		// Did we empty the whole buffer?

		if r.offset+x == len(buf) {

			// On to the next buffer.

			r.offset = 0

			r.bufs = r.bufs[1:]

			// We can release this buffer.

			putBuf(buf)

		} else {

			r.offset += x

		}

		if n == len(p) {

			break

		}

	}

	// No buffers left or nothing read?

	if len(r.bufs) == 0 {

		err = io.EOF

	}

	return

}

func (r *readCloser) Close() error {

	// Release all remaining buffers.

	for _, buf := range r.bufs {

		putBuf(buf)

	}

	// In case Close gets called multiple times.

	r.bufs = nil

	return nil

}

// ReadCloser creates an io.ReadCloser with all the contents of the buffer.

func (b *Buffer) ReadCloser() io.ReadCloser {

	ret := &readCloser{0, append(b.bufs, b.Buf)}

	b.bufs = nil

	b.toPool = nil

	b.Buf = nil

	return ret

}