131 lines
2.9 KiB
Go
131 lines
2.9 KiB
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 gives you alternate implementations of functionality
|
|
// similar to go's bytes package
|
|
|
|
package bytes2
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"unicode/utf8"
|
|
|
|
"github.com/siddontang/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
|
|
}
|