github.com/yanyiwu/go@v0.0.0-20150106053140-03d6637dbb7f/src/bytes/buffer.go (about)

     1  // Copyright 2009 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package bytes
     6  
     7  // Simple byte buffer for marshaling data.
     8  
     9  import (
    10  	"errors"
    11  	"io"
    12  	"unicode/utf8"
    13  )
    14  
    15  // A Buffer is a variable-sized buffer of bytes with Read and Write methods.
    16  // The zero value for Buffer is an empty buffer ready to use.
    17  type Buffer struct {
    18  	buf       []byte            // contents are the bytes buf[off : len(buf)]
    19  	off       int               // read at &buf[off], write at &buf[len(buf)]
    20  	runeBytes [utf8.UTFMax]byte // avoid allocation of slice on each WriteByte or Rune
    21  	bootstrap [64]byte          // memory to hold first slice; helps small buffers (Printf) avoid allocation.
    22  	lastRead  readOp            // last read operation, so that Unread* can work correctly.
    23  }
    24  
    25  // The readOp constants describe the last action performed on
    26  // the buffer, so that UnreadRune and UnreadByte can
    27  // check for invalid usage.
    28  type readOp int
    29  
    30  const (
    31  	opInvalid  readOp = iota // Non-read operation.
    32  	opReadRune               // Read rune.
    33  	opRead                   // Any other read operation.
    34  )
    35  
    36  // ErrTooLarge is passed to panic if memory cannot be allocated to store data in a buffer.
    37  var ErrTooLarge = errors.New("bytes.Buffer: too large")
    38  
    39  // Bytes returns a slice of the contents of the unread portion of the buffer;
    40  // len(b.Bytes()) == b.Len().  If the caller changes the contents of the
    41  // returned slice, the contents of the buffer will change provided there
    42  // are no intervening method calls on the Buffer.
    43  func (b *Buffer) Bytes() []byte { return b.buf[b.off:] }
    44  
    45  // String returns the contents of the unread portion of the buffer
    46  // as a string.  If the Buffer is a nil pointer, it returns "<nil>".
    47  func (b *Buffer) String() string {
    48  	if b == nil {
    49  		// Special case, useful in debugging.
    50  		return "<nil>"
    51  	}
    52  	return string(b.buf[b.off:])
    53  }
    54  
    55  // Len returns the number of bytes of the unread portion of the buffer;
    56  // b.Len() == len(b.Bytes()).
    57  func (b *Buffer) Len() int { return len(b.buf) - b.off }
    58  
    59  // Truncate discards all but the first n unread bytes from the buffer.
    60  // It panics if n is negative or greater than the length of the buffer.
    61  func (b *Buffer) Truncate(n int) {
    62  	b.lastRead = opInvalid
    63  	switch {
    64  	case n < 0 || n > b.Len():
    65  		panic("bytes.Buffer: truncation out of range")
    66  	case n == 0:
    67  		// Reuse buffer space.
    68  		b.off = 0
    69  	}
    70  	b.buf = b.buf[0 : b.off+n]
    71  }
    72  
    73  // Reset resets the buffer so it has no content.
    74  // b.Reset() is the same as b.Truncate(0).
    75  func (b *Buffer) Reset() { b.Truncate(0) }
    76  
    77  // grow grows the buffer to guarantee space for n more bytes.
    78  // It returns the index where bytes should be written.
    79  // If the buffer can't grow it will panic with ErrTooLarge.
    80  func (b *Buffer) grow(n int) int {
    81  	m := b.Len()
    82  	// If buffer is empty, reset to recover space.
    83  	if m == 0 && b.off != 0 {
    84  		b.Truncate(0)
    85  	}
    86  	if len(b.buf)+n > cap(b.buf) {
    87  		var buf []byte
    88  		if b.buf == nil && n <= len(b.bootstrap) {
    89  			buf = b.bootstrap[0:]
    90  		} else if m+n <= cap(b.buf)/2 {
    91  			// We can slide things down instead of allocating a new
    92  			// slice. We only need m+n <= cap(b.buf) to slide, but
    93  			// we instead let capacity get twice as large so we
    94  			// don't spend all our time copying.
    95  			copy(b.buf[:], b.buf[b.off:])
    96  			buf = b.buf[:m]
    97  		} else {
    98  			// not enough space anywhere
    99  			buf = makeSlice(2*cap(b.buf) + n)
   100  			copy(buf, b.buf[b.off:])
   101  		}
   102  		b.buf = buf
   103  		b.off = 0
   104  	}
   105  	b.buf = b.buf[0 : b.off+m+n]
   106  	return b.off + m
   107  }
   108  
   109  // Grow grows the buffer's capacity, if necessary, to guarantee space for
   110  // another n bytes. After Grow(n), at least n bytes can be written to the
   111  // buffer without another allocation.
   112  // If n is negative, Grow will panic.
   113  // If the buffer can't grow it will panic with ErrTooLarge.
   114  func (b *Buffer) Grow(n int) {
   115  	if n < 0 {
   116  		panic("bytes.Buffer.Grow: negative count")
   117  	}
   118  	m := b.grow(n)
   119  	b.buf = b.buf[0:m]
   120  }
   121  
   122  // Write appends the contents of p to the buffer, growing the buffer as
   123  // needed. The return value n is the length of p; err is always nil. If the
   124  // buffer becomes too large, Write will panic with ErrTooLarge.
   125  func (b *Buffer) Write(p []byte) (n int, err error) {
   126  	b.lastRead = opInvalid
   127  	m := b.grow(len(p))
   128  	return copy(b.buf[m:], p), nil
   129  }
   130  
   131  // WriteString appends the contents of s to the buffer, growing the buffer as
   132  // needed. The return value n is the length of s; err is always nil. If the
   133  // buffer becomes too large, WriteString will panic with ErrTooLarge.
   134  func (b *Buffer) WriteString(s string) (n int, err error) {
   135  	b.lastRead = opInvalid
   136  	m := b.grow(len(s))
   137  	return copy(b.buf[m:], s), nil
   138  }
   139  
   140  // MinRead is the minimum slice size passed to a Read call by
   141  // Buffer.ReadFrom.  As long as the Buffer has at least MinRead bytes beyond
   142  // what is required to hold the contents of r, ReadFrom will not grow the
   143  // underlying buffer.
   144  const MinRead = 512
   145  
   146  // ReadFrom reads data from r until EOF and appends it to the buffer, growing
   147  // the buffer as needed. The return value n is the number of bytes read. Any
   148  // error except io.EOF encountered during the read is also returned. If the
   149  // buffer becomes too large, ReadFrom will panic with ErrTooLarge.
   150  func (b *Buffer) ReadFrom(r io.Reader) (n int64, err error) {
   151  	b.lastRead = opInvalid
   152  	// If buffer is empty, reset to recover space.
   153  	if b.off >= len(b.buf) {
   154  		b.Truncate(0)
   155  	}
   156  	for {
   157  		if free := cap(b.buf) - len(b.buf); free < MinRead {
   158  			// not enough space at end
   159  			newBuf := b.buf
   160  			if b.off+free < MinRead {
   161  				// not enough space using beginning of buffer;
   162  				// double buffer capacity
   163  				newBuf = makeSlice(2*cap(b.buf) + MinRead)
   164  			}
   165  			copy(newBuf, b.buf[b.off:])
   166  			b.buf = newBuf[:len(b.buf)-b.off]
   167  			b.off = 0
   168  		}
   169  		m, e := r.Read(b.buf[len(b.buf):cap(b.buf)])
   170  		b.buf = b.buf[0 : len(b.buf)+m]
   171  		n += int64(m)
   172  		if e == io.EOF {
   173  			break
   174  		}
   175  		if e != nil {
   176  			return n, e
   177  		}
   178  	}
   179  	return n, nil // err is EOF, so return nil explicitly
   180  }
   181  
   182  // makeSlice allocates a slice of size n. If the allocation fails, it panics
   183  // with ErrTooLarge.
   184  func makeSlice(n int) []byte {
   185  	// If the make fails, give a known error.
   186  	defer func() {
   187  		if recover() != nil {
   188  			panic(ErrTooLarge)
   189  		}
   190  	}()
   191  	return make([]byte, n)
   192  }
   193  
   194  // WriteTo writes data to w until the buffer is drained or an error occurs.
   195  // The return value n is the number of bytes written; it always fits into an
   196  // int, but it is int64 to match the io.WriterTo interface. Any error
   197  // encountered during the write is also returned.
   198  func (b *Buffer) WriteTo(w io.Writer) (n int64, err error) {
   199  	b.lastRead = opInvalid
   200  	if b.off < len(b.buf) {
   201  		nBytes := b.Len()
   202  		m, e := w.Write(b.buf[b.off:])
   203  		if m > nBytes {
   204  			panic("bytes.Buffer.WriteTo: invalid Write count")
   205  		}
   206  		b.off += m
   207  		n = int64(m)
   208  		if e != nil {
   209  			return n, e
   210  		}
   211  		// all bytes should have been written, by definition of
   212  		// Write method in io.Writer
   213  		if m != nBytes {
   214  			return n, io.ErrShortWrite
   215  		}
   216  	}
   217  	// Buffer is now empty; reset.
   218  	b.Truncate(0)
   219  	return
   220  }
   221  
   222  // WriteByte appends the byte c to the buffer, growing the buffer as needed.
   223  // The returned error is always nil, but is included to match bufio.Writer's
   224  // WriteByte. If the buffer becomes too large, WriteByte will panic with
   225  // ErrTooLarge.
   226  func (b *Buffer) WriteByte(c byte) error {
   227  	b.lastRead = opInvalid
   228  	m := b.grow(1)
   229  	b.buf[m] = c
   230  	return nil
   231  }
   232  
   233  // WriteRune appends the UTF-8 encoding of Unicode code point r to the
   234  // buffer, returning its length and an error, which is always nil but is
   235  // included to match bufio.Writer's WriteRune. The buffer is grown as needed;
   236  // if it becomes too large, WriteRune will panic with ErrTooLarge.
   237  func (b *Buffer) WriteRune(r rune) (n int, err error) {
   238  	if r < utf8.RuneSelf {
   239  		b.WriteByte(byte(r))
   240  		return 1, nil
   241  	}
   242  	n = utf8.EncodeRune(b.runeBytes[0:], r)
   243  	b.Write(b.runeBytes[0:n])
   244  	return n, nil
   245  }
   246  
   247  // Read reads the next len(p) bytes from the buffer or until the buffer
   248  // is drained.  The return value n is the number of bytes read.  If the
   249  // buffer has no data to return, err is io.EOF (unless len(p) is zero);
   250  // otherwise it is nil.
   251  func (b *Buffer) Read(p []byte) (n int, err error) {
   252  	b.lastRead = opInvalid
   253  	if b.off >= len(b.buf) {
   254  		// Buffer is empty, reset to recover space.
   255  		b.Truncate(0)
   256  		if len(p) == 0 {
   257  			return
   258  		}
   259  		return 0, io.EOF
   260  	}
   261  	n = copy(p, b.buf[b.off:])
   262  	b.off += n
   263  	if n > 0 {
   264  		b.lastRead = opRead
   265  	}
   266  	return
   267  }
   268  
   269  // Next returns a slice containing the next n bytes from the buffer,
   270  // advancing the buffer as if the bytes had been returned by Read.
   271  // If there are fewer than n bytes in the buffer, Next returns the entire buffer.
   272  // The slice is only valid until the next call to a read or write method.
   273  func (b *Buffer) Next(n int) []byte {
   274  	b.lastRead = opInvalid
   275  	m := b.Len()
   276  	if n > m {
   277  		n = m
   278  	}
   279  	data := b.buf[b.off : b.off+n]
   280  	b.off += n
   281  	if n > 0 {
   282  		b.lastRead = opRead
   283  	}
   284  	return data
   285  }
   286  
   287  // ReadByte reads and returns the next byte from the buffer.
   288  // If no byte is available, it returns error io.EOF.
   289  func (b *Buffer) ReadByte() (c byte, err error) {
   290  	b.lastRead = opInvalid
   291  	if b.off >= len(b.buf) {
   292  		// Buffer is empty, reset to recover space.
   293  		b.Truncate(0)
   294  		return 0, io.EOF
   295  	}
   296  	c = b.buf[b.off]
   297  	b.off++
   298  	b.lastRead = opRead
   299  	return c, nil
   300  }
   301  
   302  // ReadRune reads and returns the next UTF-8-encoded
   303  // Unicode code point from the buffer.
   304  // If no bytes are available, the error returned is io.EOF.
   305  // If the bytes are an erroneous UTF-8 encoding, it
   306  // consumes one byte and returns U+FFFD, 1.
   307  func (b *Buffer) ReadRune() (r rune, size int, err error) {
   308  	b.lastRead = opInvalid
   309  	if b.off >= len(b.buf) {
   310  		// Buffer is empty, reset to recover space.
   311  		b.Truncate(0)
   312  		return 0, 0, io.EOF
   313  	}
   314  	b.lastRead = opReadRune
   315  	c := b.buf[b.off]
   316  	if c < utf8.RuneSelf {
   317  		b.off++
   318  		return rune(c), 1, nil
   319  	}
   320  	r, n := utf8.DecodeRune(b.buf[b.off:])
   321  	b.off += n
   322  	return r, n, nil
   323  }
   324  
   325  // UnreadRune unreads the last rune returned by ReadRune.
   326  // If the most recent read or write operation on the buffer was
   327  // not a ReadRune, UnreadRune returns an error.  (In this regard
   328  // it is stricter than UnreadByte, which will unread the last byte
   329  // from any read operation.)
   330  func (b *Buffer) UnreadRune() error {
   331  	if b.lastRead != opReadRune {
   332  		return errors.New("bytes.Buffer: UnreadRune: previous operation was not ReadRune")
   333  	}
   334  	b.lastRead = opInvalid
   335  	if b.off > 0 {
   336  		_, n := utf8.DecodeLastRune(b.buf[0:b.off])
   337  		b.off -= n
   338  	}
   339  	return nil
   340  }
   341  
   342  // UnreadByte unreads the last byte returned by the most recent
   343  // read operation.  If write has happened since the last read, UnreadByte
   344  // returns an error.
   345  func (b *Buffer) UnreadByte() error {
   346  	if b.lastRead != opReadRune && b.lastRead != opRead {
   347  		return errors.New("bytes.Buffer: UnreadByte: previous operation was not a read")
   348  	}
   349  	b.lastRead = opInvalid
   350  	if b.off > 0 {
   351  		b.off--
   352  	}
   353  	return nil
   354  }
   355  
   356  // ReadBytes reads until the first occurrence of delim in the input,
   357  // returning a slice containing the data up to and including the delimiter.
   358  // If ReadBytes encounters an error before finding a delimiter,
   359  // it returns the data read before the error and the error itself (often io.EOF).
   360  // ReadBytes returns err != nil if and only if the returned data does not end in
   361  // delim.
   362  func (b *Buffer) ReadBytes(delim byte) (line []byte, err error) {
   363  	slice, err := b.readSlice(delim)
   364  	// return a copy of slice. The buffer's backing array may
   365  	// be overwritten by later calls.
   366  	line = append(line, slice...)
   367  	return
   368  }
   369  
   370  // readSlice is like ReadBytes but returns a reference to internal buffer data.
   371  func (b *Buffer) readSlice(delim byte) (line []byte, err error) {
   372  	i := IndexByte(b.buf[b.off:], delim)
   373  	end := b.off + i + 1
   374  	if i < 0 {
   375  		end = len(b.buf)
   376  		err = io.EOF
   377  	}
   378  	line = b.buf[b.off:end]
   379  	b.off = end
   380  	b.lastRead = opRead
   381  	return line, err
   382  }
   383  
   384  // ReadString reads until the first occurrence of delim in the input,
   385  // returning a string containing the data up to and including the delimiter.
   386  // If ReadString encounters an error before finding a delimiter,
   387  // it returns the data read before the error and the error itself (often io.EOF).
   388  // ReadString returns err != nil if and only if the returned data does not end
   389  // in delim.
   390  func (b *Buffer) ReadString(delim byte) (line string, err error) {
   391  	slice, err := b.readSlice(delim)
   392  	return string(slice), err
   393  }
   394  
   395  // NewBuffer creates and initializes a new Buffer using buf as its initial
   396  // contents.  It is intended to prepare a Buffer to read existing data.  It
   397  // can also be used to size the internal buffer for writing. To do that,
   398  // buf should have the desired capacity but a length of zero.
   399  //
   400  // In most cases, new(Buffer) (or just declaring a Buffer variable) is
   401  // sufficient to initialize a Buffer.
   402  func NewBuffer(buf []byte) *Buffer { return &Buffer{buf: buf} }
   403  
   404  // NewBufferString creates and initializes a new Buffer using string s as its
   405  // initial contents. It is intended to prepare a buffer to read an existing
   406  // string.
   407  //
   408  // In most cases, new(Buffer) (or just declaring a Buffer variable) is
   409  // sufficient to initialize a Buffer.
   410  func NewBufferString(s string) *Buffer {
   411  	return &Buffer{buf: []byte(s)}
   412  }