chunked_writer.go

// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package bytes2 provides alternate implementations of functionality similar
// to go's bytes package.
package bytes2

import (
	"fmt"
	"io"
	"unicode/utf8"

	"github.com/youtube/vitess/go/hack"
)

// ChunkedWriter has the same interface as bytes.Buffer's write functions.
// It additionally provides a Reserve function that returns a []byte that
// the caller can directly change.
type ChunkedWriter struct {
	bufs [][]byte
}

func NewChunkedWriter(chunkSize int) *ChunkedWriter {
	cw := &ChunkedWriter{make([][]byte, 1)}
	cw.bufs[0] = make([]byte, 0, chunkSize)
	return cw
}

// Bytes This function can get expensive for large buffers.
func (cw *ChunkedWriter) Bytes() (b []byte) {
	if len(cw.bufs) == 1 {
		return cw.bufs[0]
	}
	b = make([]byte, 0, cw.Len())
	for _, buf := range cw.bufs {
		b = append(b, buf...)
	}
	return b
}

func (cw *ChunkedWriter) Len() int {
	l := 0
	for _, buf := range cw.bufs {
		l += len(buf)
	}
	return l
}

func (cw *ChunkedWriter) Reset() {
	b := cw.bufs[0][:0]
	cw.bufs = make([][]byte, 1)
	cw.bufs[0] = b
}

func (cw *ChunkedWriter) Truncate(n int) {
	for i, buf := range cw.bufs {
		if n > len(buf) {
			n -= len(buf)
			continue
		}
		cw.bufs[i] = buf[:n]
		cw.bufs = cw.bufs[:i+1]
		return
	}
	panic("bytes.ChunkedBuffer: truncation out of range")
}

func (cw *ChunkedWriter) Write(p []byte) (n int, err error) {
	return cw.WriteString(hack.String(p))
}

func (cw *ChunkedWriter) WriteString(p string) (n int, err error) {
	n = len(p)
	lastbuf := cw.bufs[len(cw.bufs)-1]
	for {
		available := cap(lastbuf) - len(lastbuf)
		required := len(p)
		if available >= required {
			cw.bufs[len(cw.bufs)-1] = append(lastbuf, p...)
			return
		}
		cw.bufs[len(cw.bufs)-1] = append(lastbuf, p[:available]...)
		p = p[available:]
		lastbuf = make([]byte, 0, cap(cw.bufs[0]))
		cw.bufs = append(cw.bufs, lastbuf)
	}
}

func (cw *ChunkedWriter) Reserve(n int) (b []byte) {
	if n > cap(cw.bufs[0]) {
		panic(fmt.Sprintf("bytes.ChunkedBuffer: Reserve request too high: %d > %d", n, cap(cw.bufs[0])))
	}
	lastbuf := cw.bufs[len(cw.bufs)-1]
	if n > cap(lastbuf)-len(lastbuf) {
		b = make([]byte, n, cap(cw.bufs[0]))
		cw.bufs = append(cw.bufs, b)
		return b
	}
	l := len(lastbuf)
	b = lastbuf[l : n+l]
	cw.bufs[len(cw.bufs)-1] = lastbuf[:n+l]
	return b
}

func (cw *ChunkedWriter) WriteByte(c byte) error {
	cw.Reserve(1)[0] = c
	return nil
}

func (cw *ChunkedWriter) WriteRune(r rune) (n int, err error) {
	n = utf8.EncodeRune(cw.Reserve(utf8.RuneLen(r)), r)
	return n, nil
}

func (cw *ChunkedWriter) WriteTo(w io.Writer) (n int64, err error) {
	for _, buf := range cw.bufs {
		m, err := w.Write(buf)
		n += int64(m)
		if err != nil {
			return n, err
		}
		if m != len(buf) {
			return n, io.ErrShortWrite
		}
	}
	cw.Reset()
	return n, nil
}