Source file src/pkg/compress/zlib/writer.go
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 5 package zlib 6 7 import ( 8 "compress/flate" 9 "fmt" 10 "hash" 11 "hash/adler32" 12 "io" 13 ) 14 15 // These constants are copied from the flate package, so that code that imports 16 // "compress/zlib" does not also have to import "compress/flate". 17 const ( 18 NoCompression = flate.NoCompression 19 BestSpeed = flate.BestSpeed 20 BestCompression = flate.BestCompression 21 DefaultCompression = flate.DefaultCompression 22 ) 23 24 // A Writer takes data written to it and writes the compressed 25 // form of that data to an underlying writer (see NewWriter). 26 type Writer struct { 27 w io.Writer 28 level int 29 dict []byte 30 compressor *flate.Writer 31 digest hash.Hash32 32 err error 33 scratch [4]byte 34 wroteHeader bool 35 } 36 37 // NewWriter creates a new Writer that satisfies writes by compressing data 38 // written to w. 39 // 40 // It is the caller's responsibility to call Close on the WriteCloser when done. 41 // Writes may be buffered and not flushed until Close. 42 func NewWriter(w io.Writer) *Writer { 43 z, _ := NewWriterLevelDict(w, DefaultCompression, nil) 44 return z 45 } 46 47 // NewWriterLevel is like NewWriter but specifies the compression level instead 48 // of assuming DefaultCompression. 49 // 50 // The compression level can be DefaultCompression, NoCompression, or any 51 // integer value between BestSpeed and BestCompression inclusive. The error 52 // returned will be nil if the level is valid. 53 func NewWriterLevel(w io.Writer, level int) (*Writer, error) { 54 return NewWriterLevelDict(w, level, nil) 55 } 56 57 // NewWriterLevelDict is like NewWriterLevel but specifies a dictionary to 58 // compress with. 59 // 60 // The dictionary may be nil. If not, its contents should not be modified until 61 // the Writer is closed. 62 func NewWriterLevelDict(w io.Writer, level int, dict []byte) (*Writer, error) { 63 if level < DefaultCompression || level > BestCompression { 64 return nil, fmt.Errorf("zlib: invalid compression level: %d", level) 65 } 66 return &Writer{ 67 w: w, 68 level: level, 69 dict: dict, 70 }, nil 71 } 72 73 // writeHeader writes the ZLIB header. 74 func (z *Writer) writeHeader() (err error) { 75 z.wroteHeader = true 76 // ZLIB has a two-byte header (as documented in RFC 1950). 77 // The first four bits is the CINFO (compression info), which is 7 for the default deflate window size. 78 // The next four bits is the CM (compression method), which is 8 for deflate. 79 z.scratch[0] = 0x78 80 // The next two bits is the FLEVEL (compression level). The four values are: 81 // 0=fastest, 1=fast, 2=default, 3=best. 82 // The next bit, FDICT, is set if a dictionary is given. 83 // The final five FCHECK bits form a mod-31 checksum. 84 switch z.level { 85 case 0, 1: 86 z.scratch[1] = 0 << 6 87 case 2, 3, 4, 5: 88 z.scratch[1] = 1 << 6 89 case 6, -1: 90 z.scratch[1] = 2 << 6 91 case 7, 8, 9: 92 z.scratch[1] = 3 << 6 93 default: 94 panic("unreachable") 95 } 96 if z.dict != nil { 97 z.scratch[1] |= 1 << 5 98 } 99 z.scratch[1] += uint8(31 - (uint16(z.scratch[0])<<8+uint16(z.scratch[1]))%31) 100 if _, err = z.w.Write(z.scratch[0:2]); err != nil { 101 return err 102 } 103 if z.dict != nil { 104 // The next four bytes are the Adler-32 checksum of the dictionary. 105 checksum := adler32.Checksum(z.dict) 106 z.scratch[0] = uint8(checksum >> 24) 107 z.scratch[1] = uint8(checksum >> 16) 108 z.scratch[2] = uint8(checksum >> 8) 109 z.scratch[3] = uint8(checksum >> 0) 110 if _, err = z.w.Write(z.scratch[0:4]); err != nil { 111 return err 112 } 113 } 114 z.compressor, err = flate.NewWriterDict(z.w, z.level, z.dict) 115 if err != nil { 116 return err 117 } 118 z.digest = adler32.New() 119 return nil 120 } 121 122 // Write writes a compressed form of p to the underlying io.Writer. The 123 // compressed bytes are not necessarily flushed until the Writer is closed or 124 // explicitly flushed. 125 func (z *Writer) Write(p []byte) (n int, err error) { 126 if !z.wroteHeader { 127 z.err = z.writeHeader() 128 } 129 if z.err != nil { 130 return 0, z.err 131 } 132 if len(p) == 0 { 133 return 0, nil 134 } 135 n, err = z.compressor.Write(p) 136 if err != nil { 137 z.err = err 138 return 139 } 140 z.digest.Write(p) 141 return 142 } 143 144 // Flush flushes the Writer to its underlying io.Writer. 145 func (z *Writer) Flush() error { 146 if !z.wroteHeader { 147 z.err = z.writeHeader() 148 } 149 if z.err != nil { 150 return z.err 151 } 152 z.err = z.compressor.Flush() 153 return z.err 154 } 155 156 // Calling Close does not close the wrapped io.Writer originally passed to NewWriter. 157 func (z *Writer) Close() error { 158 if !z.wroteHeader { 159 z.err = z.writeHeader() 160 } 161 if z.err != nil { 162 return z.err 163 } 164 z.err = z.compressor.Close() 165 if z.err != nil { 166 return z.err 167 } 168 checksum := z.digest.Sum32() 169 // ZLIB (RFC 1950) is big-endian, unlike GZIP (RFC 1952). 170 z.scratch[0] = uint8(checksum >> 24) 171 z.scratch[1] = uint8(checksum >> 16) 172 z.scratch[2] = uint8(checksum >> 8) 173 z.scratch[3] = uint8(checksum >> 0) 174 _, z.err = z.w.Write(z.scratch[0:4]) 175 return z.err 176 }