...
1
2
3
4
5 package zlib
6
7 import (
8 "compress/flate"
9 "encoding/binary"
10 "fmt"
11 "hash"
12 "hash/adler32"
13 "io"
14 )
15
16
17
18 const (
19 NoCompression = flate.NoCompression
20 BestSpeed = flate.BestSpeed
21 BestCompression = flate.BestCompression
22 DefaultCompression = flate.DefaultCompression
23 HuffmanOnly = flate.HuffmanOnly
24 )
25
26
27
28 type Writer struct {
29 w io.Writer
30 level int
31 dict []byte
32 compressor *flate.Writer
33 digest hash.Hash32
34 err error
35 scratch [4]byte
36 wroteHeader bool
37 }
38
39
40
41
42
43
44 func NewWriter(w io.Writer) *Writer {
45 z, _ := NewWriterLevelDict(w, DefaultCompression, nil)
46 return z
47 }
48
49
50
51
52
53
54
55 func NewWriterLevel(w io.Writer, level int) (*Writer, error) {
56 return NewWriterLevelDict(w, level, nil)
57 }
58
59
60
61
62
63
64 func NewWriterLevelDict(w io.Writer, level int, dict []byte) (*Writer, error) {
65 if level < HuffmanOnly || level > BestCompression {
66 return nil, fmt.Errorf("zlib: invalid compression level: %d", level)
67 }
68 return &Writer{
69 w: w,
70 level: level,
71 dict: dict,
72 }, nil
73 }
74
75
76
77
78 func (z *Writer) Reset(w io.Writer) {
79 z.w = w
80
81 if z.compressor != nil {
82 z.compressor.Reset(w)
83 }
84 if z.digest != nil {
85 z.digest.Reset()
86 }
87 z.err = nil
88 z.scratch = [4]byte{}
89 z.wroteHeader = false
90 }
91
92
93 func (z *Writer) writeHeader() (err error) {
94 z.wroteHeader = true
95
96
97
98 z.scratch[0] = 0x78
99
100
101
102
103 switch z.level {
104 case -2, 0, 1:
105 z.scratch[1] = 0 << 6
106 case 2, 3, 4, 5:
107 z.scratch[1] = 1 << 6
108 case 6, -1:
109 z.scratch[1] = 2 << 6
110 case 7, 8, 9:
111 z.scratch[1] = 3 << 6
112 default:
113 panic("unreachable")
114 }
115 if z.dict != nil {
116 z.scratch[1] |= 1 << 5
117 }
118 z.scratch[1] += uint8(31 - binary.BigEndian.Uint16(z.scratch[:2])%31)
119 if _, err = z.w.Write(z.scratch[0:2]); err != nil {
120 return err
121 }
122 if z.dict != nil {
123
124 binary.BigEndian.PutUint32(z.scratch[:], adler32.Checksum(z.dict))
125 if _, err = z.w.Write(z.scratch[0:4]); err != nil {
126 return err
127 }
128 }
129 if z.compressor == nil {
130
131
132 z.compressor, err = flate.NewWriterDict(z.w, z.level, z.dict)
133 if err != nil {
134 return err
135 }
136 z.digest = adler32.New()
137 }
138 return nil
139 }
140
141
142
143
144 func (z *Writer) Write(p []byte) (n int, err error) {
145 if !z.wroteHeader {
146 z.err = z.writeHeader()
147 }
148 if z.err != nil {
149 return 0, z.err
150 }
151 if len(p) == 0 {
152 return 0, nil
153 }
154 n, err = z.compressor.Write(p)
155 if err != nil {
156 z.err = err
157 return
158 }
159 z.digest.Write(p)
160 return
161 }
162
163
164 func (z *Writer) Flush() error {
165 if !z.wroteHeader {
166 z.err = z.writeHeader()
167 }
168 if z.err != nil {
169 return z.err
170 }
171 z.err = z.compressor.Flush()
172 return z.err
173 }
174
175
176
177 func (z *Writer) Close() error {
178 if !z.wroteHeader {
179 z.err = z.writeHeader()
180 }
181 if z.err != nil {
182 return z.err
183 }
184 z.err = z.compressor.Close()
185 if z.err != nil {
186 return z.err
187 }
188 checksum := z.digest.Sum32()
189
190 binary.BigEndian.PutUint32(z.scratch[:], checksum)
191 _, z.err = z.w.Write(z.scratch[0:4])
192 return z.err
193 }
194
View as plain text