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