compress/gzip: add Writer.Reset

compress/flate is https://golang.org/cl/12953048
compress/zlib is https://golang.org/cl/13171046

Update #6138

R=golang-dev, iant
CC=golang-dev
https://golang.org/cl/13435043

src/pkg/compress/gzip/gzip.go

@@ -28,6 +28,7 @@ type Writer struct {
 	Header
 	w           io.Writer
 	level       int
+	wroteHeader bool
 	compressor  *flate.Writer
 	digest      hash.Hash32
 	size        uint32
@@ -62,14 +63,39 @@ func NewWriterLevel(w io.Writer, level int) (*Writer, error) {
 	if level < DefaultCompression || level > BestCompression {
 		return nil, fmt.Errorf("gzip: invalid compression level: %d", level)
 	}
-	return &Writer{
+	z := new(Writer)
+	z.init(w, level)
+	return z, nil
+}
+
+func (z *Writer) init(w io.Writer, level int) {
+	digest := z.digest
+	if digest != nil {
+		digest.Reset()
+	} else {
+		digest = crc32.NewIEEE()
+	}
+	compressor := z.compressor
+	if compressor != nil {
+		compressor.Reset(w)
+	}
+	*z = Writer{
 		Header: Header{
 			OS: 255, // unknown
 		},
 		w:          w,
 		level:      level,
-		digest: crc32.NewIEEE(),
-	}, nil
+		digest:     digest,
+		compressor: compressor,
+	}
+}
+
+// Reset discards the Writer z's state and makes it equivalent to the
+// result of its original state from NewWriter or NewWriterLevel, but
+// writing to w instead. This permits reusing a Writer rather than
+// allocating a new one.
+func (z *Writer) Reset(w io.Writer) {
+	z.init(w, z.level)
 }
 
 // GZIP (RFC 1952) is little-endian, unlike ZLIB (RFC 1950).
@@ -138,7 +164,8 @@ func (z *Writer) Write(p []byte) (int, error) {
 	}
 	var n int
 	// Write the GZIP header lazily.
-	if z.compressor == nil {
+	if !z.wroteHeader {
+		z.wroteHeader = true
 		z.buf[0] = gzipID1
 		z.buf[1] = gzipID2
 		z.buf[2] = gzipDeflate
@@ -183,8 +210,10 @@ func (z *Writer) Write(p []byte) (int, error) {
 				return n, z.err
 			}
 		}
+		if z.compressor == nil {
 			z.compressor, _ = flate.NewWriter(z.w, z.level)
 		}
+	}
 	z.size += uint32(len(p))
 	z.digest.Write(p)
 	n, z.err = z.compressor.Write(p)
@@ -206,8 +235,11 @@ func (z *Writer) Flush() error {
 	if z.closed {
 		return nil
 	}
-	if z.compressor == nil {
+	if !z.wroteHeader {
 		z.Write(nil)
+		if z.err != nil {
+			return z.err
+		}
 	}
 	z.err = z.compressor.Flush()
 	return z.err
@@ -222,7 +254,7 @@ func (z *Writer) Close() error {
 		return nil
 	}
 	z.closed = true
-	if z.compressor == nil {
+	if !z.wroteHeader {
 		z.Write(nil)
 		if z.err != nil {
 			return z.err

src/pkg/compress/gzip/gzip_test.go

@@ -214,3 +214,18 @@ func TestConcat(t *testing.T) {
 		t.Fatalf("ReadAll = %q, %v, want %q, nil", data, err, "hello world")
 	}
 }
+
+func TestWriterReset(t *testing.T) {
+	buf := new(bytes.Buffer)
+	buf2 := new(bytes.Buffer)
+	z := NewWriter(buf)
+	msg := []byte("hello world")
+	z.Write(msg)
+	z.Close()
+	z.Reset(buf2)
+	z.Write(msg)
+	z.Close()
+	if buf.String() != buf2.String() {
+		t.Errorf("buf2 %q != original buf of %q", buf2.String(), buf.String())
+	}
+}