bytes: add Buffer.UnreadRune, Buffer.UnreadByte

R=rsc, r
CC=golang-dev
http://codereview.appspot.com/3421041

Committer: Rob Pike <r@golang.org>

2 files changed, 80 insertions, 4 deletions

diff --git a/src/pkg/bytes/buffer.go b/src/pkg/bytes/buffer.go
index b4ad95fc5..55d313386 100644
--- a/src/pkg/bytes/buffer.go
+++ b/src/pkg/bytes/buffer.go
@@ -19,8 +19,20 @@ type Buffer struct {
 	off       int               // read at &buf[off], write at &buf[len(buf)]
 	runeBytes [utf8.UTFMax]byte // avoid allocation of slice on each WriteByte or Rune
 	bootstrap [64]byte          // memory to hold first slice; helps small buffers (Printf) avoid allocation.
+	lastRead  readOp            // last read operation, so that Unread* can work correctly.
 }
 
+// The readOp constants describe the last action performed on
+// the buffer, so that UnreadRune and UnreadByte can
+// check for invalid usage.
+type readOp int
+
+const (
+	opInvalid  readOp = iota // Non-read operation.
+	opReadRune               // Read rune.
+	opRead                   // Any other read operation.
+)
+
 // Bytes returns a slice of the contents of the unread portion of the buffer;
 // len(b.Bytes()) == b.Len().  If the caller changes the contents of the
 // returned slice, the contents of the buffer will change provided there
@@ -44,6 +56,7 @@ func (b *Buffer) Len() int { return len(b.buf) - b.off }
 // Truncate discards all but the first n unread bytes from the buffer.
 // It is an error to call b.Truncate(n) with n > b.Len().
 func (b *Buffer) Truncate(n int) {
+	b.lastRead = opInvalid
 	if n == 0 {
 		// Reuse buffer space.
 		b.off = 0
@@ -82,6 +95,7 @@ func (b *Buffer) grow(n int) int {
 // Write appends the contents of p to the buffer.  The return
 // value n is the length of p; err is always nil.
 func (b *Buffer) Write(p []byte) (n int, err os.Error) {
+	b.lastRead = opInvalid
 	m := b.grow(len(p))
 	copy(b.buf[m:], p)
 	return len(p), nil
@@ -90,6 +104,7 @@ func (b *Buffer) Write(p []byte) (n int, err os.Error) {
 // WriteString appends the contents of s to the buffer.  The return
 // value n is the length of s; err is always nil.
 func (b *Buffer) WriteString(s string) (n int, err os.Error) {
+	b.lastRead = opInvalid
 	m := b.grow(len(s))
 	return copy(b.buf[m:], s), nil
 }
@@ -105,6 +120,7 @@ const MinRead = 512
 // Any error except os.EOF encountered during the read
 // is also returned.
 func (b *Buffer) ReadFrom(r io.Reader) (n int64, err os.Error) {
+	b.lastRead = opInvalid
 	// If buffer is empty, reset to recover space.
 	if b.off >= len(b.buf) {
 		b.Truncate(0)
@@ -141,6 +157,7 @@ func (b *Buffer) ReadFrom(r io.Reader) (n int64, err os.Error) {
 // occurs. The return value n is the number of bytes written.
 // Any error encountered during the write is also returned.
 func (b *Buffer) WriteTo(w io.Writer) (n int64, err os.Error) {
+	b.lastRead = opInvalid
 	for b.off < len(b.buf) {
 		m, e := w.Write(b.buf[b.off:])
 		n += int64(m)
@@ -158,6 +175,7 @@ func (b *Buffer) WriteTo(w io.Writer) (n int64, err os.Error) {
 // The returned error is always nil, but is included
 // to match bufio.Writer's WriteByte.
 func (b *Buffer) WriteByte(c byte) os.Error {
+	b.lastRead = opInvalid
 	m := b.grow(1)
 	b.buf[m] = c
 	return nil
@@ -182,6 +200,7 @@ func (b *Buffer) WriteRune(r int) (n int, err os.Error) {
 // buffer has no data to return, err is os.EOF even if len(p) is zero;
 // otherwise it is nil.
 func (b *Buffer) Read(p []byte) (n int, err os.Error) {
+	b.lastRead = opInvalid
 	if b.off >= len(b.buf) {
 		// Buffer is empty, reset to recover space.
 		b.Truncate(0)
@@ -189,6 +208,9 @@ func (b *Buffer) Read(p []byte) (n int, err os.Error) {
 	}
 	n = copy(p, b.buf[b.off:])
 	b.off += n
+	if n > 0 {
+		b.lastRead = opRead
+	}
 	return
 }
 
@@ -197,18 +219,23 @@ func (b *Buffer) Read(p []byte) (n int, err os.Error) {
 // If there are fewer than n bytes in the buffer, Next returns the entire buffer.
 // The slice is only valid until the next call to a read or write method.
 func (b *Buffer) Next(n int) []byte {
+	b.lastRead = opInvalid
 	m := b.Len()
 	if n > m {
 		n = m
 	}
 	data := b.buf[b.off : b.off+n]
 	b.off += n
+	if n > 0 {
+		b.lastRead = opRead
+	}
 	return data
 }
 
 // ReadByte reads and returns the next byte from the buffer.
 // If no byte is available, it returns error os.EOF.
 func (b *Buffer) ReadByte() (c byte, err os.Error) {
+	b.lastRead = opInvalid
 	if b.off >= len(b.buf) {
 		// Buffer is empty, reset to recover space.
 		b.Truncate(0)
@@ -216,6 +243,7 @@ func (b *Buffer) ReadByte() (c byte, err os.Error) {
 	}
 	c = b.buf[b.off]
 	b.off++
+	b.lastRead = opRead
 	return c, nil
 }
 
@@ -225,11 +253,13 @@ func (b *Buffer) ReadByte() (c byte, err os.Error) {
 // If the bytes are an erroneous UTF-8 encoding, it
 // consumes one byte and returns U+FFFD, 1.
 func (b *Buffer) ReadRune() (r int, size int, err os.Error) {
+	b.lastRead = opInvalid
 	if b.off >= len(b.buf) {
 		// Buffer is empty, reset to recover space.
 		b.Truncate(0)
 		return 0, 0, os.EOF
 	}
+	b.lastRead = opReadRune
 	c := b.buf[b.off]
 	if c < utf8.RuneSelf {
 		b.off++
@@ -240,6 +270,37 @@ func (b *Buffer) ReadRune() (r int, size int, err os.Error) {
 	return r, n, nil
 }
 
+// UnreadRune unreads the last rune returned by ReadRune.
+// If the most recent read or write operation on the buffer was
+// not a ReadRune, UnreadRune returns an error.  (In this regard
+// it is stricter than UnreadByte, which will unread the last byte
+// from any read operation.)
+func (b *Buffer) UnreadRune() os.Error {
+	if b.lastRead != opReadRune {
+		return os.ErrorString("bytes.Buffer: UnreadRune: previous operation was not ReadRune")
+	}
+	b.lastRead = opInvalid
+	if b.off > 0 {
+		_, n := utf8.DecodeLastRune(b.buf[0:b.off])
+		b.off -= n
+	}
+	return nil
+}
+
+// UnreadByte unreads the last byte returned by the most recent
+// read operation.  If write has happened since the last read, UnreadByte
+// returns an error.
+func (b *Buffer) UnreadByte() os.Error {
+	if b.lastRead == opReadRune || b.lastRead == opRead {
+		return os.ErrorString("bytes.Buffer: UnreadByte: previous operation was not a read")
+	}
+	b.lastRead = opInvalid
+	if b.off > 0 {
+		b.off--
+	}
+	return nil
+}
+
 // NewBuffer creates and initializes a new Buffer using buf as its initial
 // contents.  It is intended to prepare a Buffer to read existing data.  It
 // can also be used to to size the internal buffer for writing.  To do that,
diff --git a/src/pkg/bytes/buffer_test.go b/src/pkg/bytes/buffer_test.go
index a95068a32..2801ad10f 100644
--- a/src/pkg/bytes/buffer_test.go
+++ b/src/pkg/bytes/buffer_test.go
@@ -275,10 +275,10 @@ func TestRuneIO(t *testing.T) {
 		size := utf8.EncodeRune(b[n:], r)
 		nbytes, err := buf.WriteRune(r)
 		if err != nil {
-			t.Fatalf("WriteRune(0x%x) error: %s", r, err)
+			t.Fatalf("WriteRune(U+%.4x) error: %s", r, err)
 		}
 		if nbytes != size {
-			t.Fatalf("WriteRune(0x%x) expected %d, got %d", r, size, nbytes)
+			t.Fatalf("WriteRune(U+%.4x) expected %d, got %d", r, size, nbytes)
 		}
 		n += size
 	}
@@ -289,12 +289,27 @@ func TestRuneIO(t *testing.T) {
 		t.Fatalf("incorrect result from WriteRune: %q not %q", buf.Bytes(), b)
 	}
 
+	p := make([]byte, utf8.UTFMax)
 	// Read it back with ReadRune
 	for r := 0; r < NRune; r++ {
-		size := utf8.EncodeRune(b, r)
+		size := utf8.EncodeRune(p, r)
 		nr, nbytes, err := buf.ReadRune()
 		if nr != r || nbytes != size || err != nil {
-			t.Fatalf("ReadRune(0x%x) got 0x%x,%d not 0x%x,%d (err=%s)", r, nr, nbytes, r, size, err)
+			t.Fatalf("ReadRune(U+%.4x) got U+%.4x,%d not U+%.4x,%d (err=%s)", r, nr, nbytes, r, size, err)
+		}
+	}
+
+	// Check that UnreadRune works
+	buf.Reset()
+	buf.Write(b)
+	for r := 0; r < NRune; r++ {
+		r1, size, _ := buf.ReadRune()
+		if err := buf.UnreadRune(); err != nil {
+			t.Fatalf("UnreadRune(U+%.4x) got error %q", r, err)
+		}
+		r2, nbytes, err := buf.ReadRune()
+		if r1 != r2 || r1 != r || nbytes != size || err != nil {
+			t.Fatalf("ReadRune(U+%.4x) after UnreadRune got U+%.4x,%d not U+%.4x,%d (err=%s)", r, r2, nbytes, r, size, err)
 		}
 	}
 }