1 files changed, 54 insertions, 26 deletions

diff --git a/src/compress/zlib/reader.go b/src/compress/zlib/reader.go
index e1191816d..816f1bf6b 100644
--- a/src/compress/zlib/reader.go
+++ b/src/compress/zlib/reader.go
@@ -51,10 +51,21 @@ type reader struct {
 	scratch      [4]byte
 }
 
-// NewReader creates a new io.ReadCloser.
-// Reads from the returned io.ReadCloser read and decompress data from r.
+// Resetter resets a ReadCloser returned by NewReader or NewReaderDict to
+// to switch to a new underlying Reader. This permits reusing a ReadCloser
+// instead of allocating a new one.
+type Resetter interface {
+	// Reset discards any buffered data and resets the Resetter as if it was
+	// newly initialized with the given reader.
+	Reset(r io.Reader, dict []byte) error
+}
+
+// NewReader creates a new ReadCloser.
+// Reads from the returned ReadCloser read and decompress data from r.
 // The implementation buffers input and may read more data than necessary from r.
 // It is the caller's responsibility to call Close on the ReadCloser when done.
+//
+// The ReadCloser returned by NewReader also implements Resetter.
 func NewReader(r io.Reader) (io.ReadCloser, error) {
 	return NewReaderDict(r, nil)
 }
@@ -62,35 +73,14 @@ func NewReader(r io.Reader) (io.ReadCloser, error) {
 // NewReaderDict is like NewReader but uses a preset dictionary.
 // NewReaderDict ignores the dictionary if the compressed data does not refer to it.
 // If the compressed data refers to a different dictionary, NewReaderDict returns ErrDictionary.
+//
+// The ReadCloser returned by NewReaderDict also implements Resetter.
 func NewReaderDict(r io.Reader, dict []byte) (io.ReadCloser, error) {
 	z := new(reader)
-	if fr, ok := r.(flate.Reader); ok {
-		z.r = fr
-	} else {
-		z.r = bufio.NewReader(r)
-	}
-	_, err := io.ReadFull(z.r, z.scratch[0:2])
+	err := z.Reset(r, dict)
 	if err != nil {
 		return nil, err
 	}
-	h := uint(z.scratch[0])<<8 | uint(z.scratch[1])
-	if (z.scratch[0]&0x0f != zlibDeflate) || (h%31 != 0) {
-		return nil, ErrHeader
-	}
-	if z.scratch[1]&0x20 != 0 {
-		_, err = io.ReadFull(z.r, z.scratch[0:4])
-		if err != nil {
-			return nil, err
-		}
-		checksum := uint32(z.scratch[0])<<24 | uint32(z.scratch[1])<<16 | uint32(z.scratch[2])<<8 | uint32(z.scratch[3])
-		if checksum != adler32.Checksum(dict) {
-			return nil, ErrDictionary
-		}
-		z.decompressor = flate.NewReaderDict(z.r, dict)
-	} else {
-		z.decompressor = flate.NewReader(z.r)
-	}
-	z.digest = adler32.New()
 	return z, nil
 }
 
@@ -131,3 +121,41 @@ func (z *reader) Close() error {
 	z.err = z.decompressor.Close()
 	return z.err
 }
+
+func (z *reader) Reset(r io.Reader, dict []byte) error {
+	if fr, ok := r.(flate.Reader); ok {
+		z.r = fr
+	} else {
+		z.r = bufio.NewReader(r)
+	}
+	_, err := io.ReadFull(z.r, z.scratch[0:2])
+	if err != nil {
+		return err
+	}
+	h := uint(z.scratch[0])<<8 | uint(z.scratch[1])
+	if (z.scratch[0]&0x0f != zlibDeflate) || (h%31 != 0) {
+		return ErrHeader
+	}
+	haveDict := z.scratch[1]&0x20 != 0
+	if haveDict {
+		_, err = io.ReadFull(z.r, z.scratch[0:4])
+		if err != nil {
+			return err
+		}
+		checksum := uint32(z.scratch[0])<<24 | uint32(z.scratch[1])<<16 | uint32(z.scratch[2])<<8 | uint32(z.scratch[3])
+		if checksum != adler32.Checksum(dict) {
+			return ErrDictionary
+		}
+	}
+	if z.decompressor == nil {
+		if haveDict {
+			z.decompressor = flate.NewReaderDict(z.r, dict)
+		} else {
+			z.decompressor = flate.NewReader(z.r)
+		}
+	} else {
+		z.decompressor.(flate.Resetter).Reset(z.r, dict)
+	}
+	z.digest = adler32.New()
+	return nil
+}