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

diff --git a/src/pkg/compress/gzip/gzip.go b/src/pkg/compress/gzip/gzip.go
index 45558b74289409559751a927ecf47eaf8d417eba..fe32d6871ae950985b1775c354745fca73f16048 100644 (file)
--- a/src/pkg/compress/gzip/gzip.go
+++ b/src/pkg/compress/gzip/gzip.go
@@ -26,14 +26,15 @@ const (
 // to its wrapped io.Writer.
 type Writer struct {
        Header
-       w          io.Writer
-       level      int
-       compressor *flate.Writer
-       digest     hash.Hash32
-       size       uint32
-       closed     bool
-       buf        [10]byte
-       err        error
+       w           io.Writer
+       level       int
+       wroteHeader bool
+       compressor  *flate.Writer
+       digest      hash.Hash32
+       size        uint32
+       closed      bool
+       buf         [10]byte
+       err         error
 }
 
 // NewWriter creates a new Writer that satisfies writes by compressing data
@@ -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
+               w:          w,
+               level:      level,
+               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,7 +210,9 @@ func (z *Writer) Write(p []byte) (int, error) {
                                return n, z.err
                        }
                }
-               z.compressor, _ = flate.NewWriter(z.w, z.level)
+               if z.compressor == nil {
+                       z.compressor, _ = flate.NewWriter(z.w, z.level)
+               }
        }
        z.size += uint32(len(p))
        z.digest.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

diff --git a/src/pkg/compress/gzip/gzip_test.go b/src/pkg/compress/gzip/gzip_test.go
index 39afb331dafbb9afcbddceae2a71d6d325f6d766..119be2e135be34defe891aa2b7226f92305bfa83 100644 (file)
--- a/src/pkg/compress/gzip/gzip_test.go
+++ b/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())
+       }
+}