buffer.go 9.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293
  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. package util
  5. // This a copy of Go std bytes.Buffer with some modification
  6. // and some features stripped.
  7. import (
  8. "bytes"
  9. "io"
  10. )
  11. // A Buffer is a variable-sized buffer of bytes with Read and Write methods.
  12. // The zero value for Buffer is an empty buffer ready to use.
  13. type Buffer struct {
  14. buf []byte // contents are the bytes buf[off : len(buf)]
  15. off int // read at &buf[off], write at &buf[len(buf)]
  16. bootstrap [64]byte // memory to hold first slice; helps small buffers (Printf) avoid allocation.
  17. }
  18. // Bytes returns a slice of the contents of the unread portion of the buffer;
  19. // len(b.Bytes()) == b.Len(). If the caller changes the contents of the
  20. // returned slice, the contents of the buffer will change provided there
  21. // are no intervening method calls on the Buffer.
  22. func (b *Buffer) Bytes() []byte { return b.buf[b.off:] }
  23. // String returns the contents of the unread portion of the buffer
  24. // as a string. If the Buffer is a nil pointer, it returns "<nil>".
  25. func (b *Buffer) String() string {
  26. if b == nil {
  27. // Special case, useful in debugging.
  28. return "<nil>"
  29. }
  30. return string(b.buf[b.off:])
  31. }
  32. // Len returns the number of bytes of the unread portion of the buffer;
  33. // b.Len() == len(b.Bytes()).
  34. func (b *Buffer) Len() int { return len(b.buf) - b.off }
  35. // Truncate discards all but the first n unread bytes from the buffer.
  36. // It panics if n is negative or greater than the length of the buffer.
  37. func (b *Buffer) Truncate(n int) {
  38. switch {
  39. case n < 0 || n > b.Len():
  40. panic("leveldb/util.Buffer: truncation out of range")
  41. case n == 0:
  42. // Reuse buffer space.
  43. b.off = 0
  44. }
  45. b.buf = b.buf[0 : b.off+n]
  46. }
  47. // Reset resets the buffer so it has no content.
  48. // b.Reset() is the same as b.Truncate(0).
  49. func (b *Buffer) Reset() { b.Truncate(0) }
  50. // grow grows the buffer to guarantee space for n more bytes.
  51. // It returns the index where bytes should be written.
  52. // If the buffer can't grow it will panic with bytes.ErrTooLarge.
  53. func (b *Buffer) grow(n int) int {
  54. m := b.Len()
  55. // If buffer is empty, reset to recover space.
  56. if m == 0 && b.off != 0 {
  57. b.Truncate(0)
  58. }
  59. if len(b.buf)+n > cap(b.buf) {
  60. var buf []byte
  61. if b.buf == nil && n <= len(b.bootstrap) {
  62. buf = b.bootstrap[0:]
  63. } else if m+n <= cap(b.buf)/2 {
  64. // We can slide things down instead of allocating a new
  65. // slice. We only need m+n <= cap(b.buf) to slide, but
  66. // we instead let capacity get twice as large so we
  67. // don't spend all our time copying.
  68. copy(b.buf[:], b.buf[b.off:])
  69. buf = b.buf[:m]
  70. } else {
  71. // not enough space anywhere
  72. buf = makeSlice(2*cap(b.buf) + n)
  73. copy(buf, b.buf[b.off:])
  74. }
  75. b.buf = buf
  76. b.off = 0
  77. }
  78. b.buf = b.buf[0 : b.off+m+n]
  79. return b.off + m
  80. }
  81. // Alloc allocs n bytes of slice from the buffer, growing the buffer as
  82. // needed. If n is negative, Alloc will panic.
  83. // If the buffer can't grow it will panic with bytes.ErrTooLarge.
  84. func (b *Buffer) Alloc(n int) []byte {
  85. if n < 0 {
  86. panic("leveldb/util.Buffer.Alloc: negative count")
  87. }
  88. m := b.grow(n)
  89. return b.buf[m:]
  90. }
  91. // Grow grows the buffer's capacity, if necessary, to guarantee space for
  92. // another n bytes. After Grow(n), at least n bytes can be written to the
  93. // buffer without another allocation.
  94. // If n is negative, Grow will panic.
  95. // If the buffer can't grow it will panic with bytes.ErrTooLarge.
  96. func (b *Buffer) Grow(n int) {
  97. if n < 0 {
  98. panic("leveldb/util.Buffer.Grow: negative count")
  99. }
  100. m := b.grow(n)
  101. b.buf = b.buf[0:m]
  102. }
  103. // Write appends the contents of p to the buffer, growing the buffer as
  104. // needed. The return value n is the length of p; err is always nil. If the
  105. // buffer becomes too large, Write will panic with bytes.ErrTooLarge.
  106. func (b *Buffer) Write(p []byte) (n int, err error) {
  107. m := b.grow(len(p))
  108. return copy(b.buf[m:], p), nil
  109. }
  110. // MinRead is the minimum slice size passed to a Read call by
  111. // Buffer.ReadFrom. As long as the Buffer has at least MinRead bytes beyond
  112. // what is required to hold the contents of r, ReadFrom will not grow the
  113. // underlying buffer.
  114. const MinRead = 512
  115. // ReadFrom reads data from r until EOF and appends it to the buffer, growing
  116. // the buffer as needed. The return value n is the number of bytes read. Any
  117. // error except io.EOF encountered during the read is also returned. If the
  118. // buffer becomes too large, ReadFrom will panic with bytes.ErrTooLarge.
  119. func (b *Buffer) ReadFrom(r io.Reader) (n int64, err error) {
  120. // If buffer is empty, reset to recover space.
  121. if b.off >= len(b.buf) {
  122. b.Truncate(0)
  123. }
  124. for {
  125. if free := cap(b.buf) - len(b.buf); free < MinRead {
  126. // not enough space at end
  127. newBuf := b.buf
  128. if b.off+free < MinRead {
  129. // not enough space using beginning of buffer;
  130. // double buffer capacity
  131. newBuf = makeSlice(2*cap(b.buf) + MinRead)
  132. }
  133. copy(newBuf, b.buf[b.off:])
  134. b.buf = newBuf[:len(b.buf)-b.off]
  135. b.off = 0
  136. }
  137. m, e := r.Read(b.buf[len(b.buf):cap(b.buf)])
  138. b.buf = b.buf[0 : len(b.buf)+m]
  139. n += int64(m)
  140. if e == io.EOF {
  141. break
  142. }
  143. if e != nil {
  144. return n, e
  145. }
  146. }
  147. return n, nil // err is EOF, so return nil explicitly
  148. }
  149. // makeSlice allocates a slice of size n. If the allocation fails, it panics
  150. // with bytes.ErrTooLarge.
  151. func makeSlice(n int) []byte {
  152. // If the make fails, give a known error.
  153. defer func() {
  154. if recover() != nil {
  155. panic(bytes.ErrTooLarge)
  156. }
  157. }()
  158. return make([]byte, n)
  159. }
  160. // WriteTo writes data to w until the buffer is drained or an error occurs.
  161. // The return value n is the number of bytes written; it always fits into an
  162. // int, but it is int64 to match the io.WriterTo interface. Any error
  163. // encountered during the write is also returned.
  164. func (b *Buffer) WriteTo(w io.Writer) (n int64, err error) {
  165. if b.off < len(b.buf) {
  166. nBytes := b.Len()
  167. m, e := w.Write(b.buf[b.off:])
  168. if m > nBytes {
  169. panic("leveldb/util.Buffer.WriteTo: invalid Write count")
  170. }
  171. b.off += m
  172. n = int64(m)
  173. if e != nil {
  174. return n, e
  175. }
  176. // all bytes should have been written, by definition of
  177. // Write method in io.Writer
  178. if m != nBytes {
  179. return n, io.ErrShortWrite
  180. }
  181. }
  182. // Buffer is now empty; reset.
  183. b.Truncate(0)
  184. return
  185. }
  186. // WriteByte appends the byte c to the buffer, growing the buffer as needed.
  187. // The returned error is always nil, but is included to match bufio.Writer's
  188. // WriteByte. If the buffer becomes too large, WriteByte will panic with
  189. // bytes.ErrTooLarge.
  190. func (b *Buffer) WriteByte(c byte) error {
  191. m := b.grow(1)
  192. b.buf[m] = c
  193. return nil
  194. }
  195. // Read reads the next len(p) bytes from the buffer or until the buffer
  196. // is drained. The return value n is the number of bytes read. If the
  197. // buffer has no data to return, err is io.EOF (unless len(p) is zero);
  198. // otherwise it is nil.
  199. func (b *Buffer) Read(p []byte) (n int, err error) {
  200. if b.off >= len(b.buf) {
  201. // Buffer is empty, reset to recover space.
  202. b.Truncate(0)
  203. if len(p) == 0 {
  204. return
  205. }
  206. return 0, io.EOF
  207. }
  208. n = copy(p, b.buf[b.off:])
  209. b.off += n
  210. return
  211. }
  212. // Next returns a slice containing the next n bytes from the buffer,
  213. // advancing the buffer as if the bytes had been returned by Read.
  214. // If there are fewer than n bytes in the buffer, Next returns the entire buffer.
  215. // The slice is only valid until the next call to a read or write method.
  216. func (b *Buffer) Next(n int) []byte {
  217. m := b.Len()
  218. if n > m {
  219. n = m
  220. }
  221. data := b.buf[b.off : b.off+n]
  222. b.off += n
  223. return data
  224. }
  225. // ReadByte reads and returns the next byte from the buffer.
  226. // If no byte is available, it returns error io.EOF.
  227. func (b *Buffer) ReadByte() (c byte, err error) {
  228. if b.off >= len(b.buf) {
  229. // Buffer is empty, reset to recover space.
  230. b.Truncate(0)
  231. return 0, io.EOF
  232. }
  233. c = b.buf[b.off]
  234. b.off++
  235. return c, nil
  236. }
  237. // ReadBytes reads until the first occurrence of delim in the input,
  238. // returning a slice containing the data up to and including the delimiter.
  239. // If ReadBytes encounters an error before finding a delimiter,
  240. // it returns the data read before the error and the error itself (often io.EOF).
  241. // ReadBytes returns err != nil if and only if the returned data does not end in
  242. // delim.
  243. func (b *Buffer) ReadBytes(delim byte) (line []byte, err error) {
  244. slice, err := b.readSlice(delim)
  245. // return a copy of slice. The buffer's backing array may
  246. // be overwritten by later calls.
  247. line = append(line, slice...)
  248. return
  249. }
  250. // readSlice is like ReadBytes but returns a reference to internal buffer data.
  251. func (b *Buffer) readSlice(delim byte) (line []byte, err error) {
  252. i := bytes.IndexByte(b.buf[b.off:], delim)
  253. end := b.off + i + 1
  254. if i < 0 {
  255. end = len(b.buf)
  256. err = io.EOF
  257. }
  258. line = b.buf[b.off:end]
  259. b.off = end
  260. return line, err
  261. }
  262. // NewBuffer creates and initializes a new Buffer using buf as its initial
  263. // contents. It is intended to prepare a Buffer to read existing data. It
  264. // can also be used to size the internal buffer for writing. To do that,
  265. // buf should have the desired capacity but a length of zero.
  266. //
  267. // In most cases, new(Buffer) (or just declaring a Buffer variable) is
  268. // sufficient to initialize a Buffer.
  269. func NewBuffer(buf []byte) *Buffer { return &Buffer{buf: buf} }