Source file src/pkg/compress/lzw/writer.go
1 // Copyright 2011 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 lzw
6
7 import (
8 "bufio"
9 "errors"
10 "fmt"
11 "io"
12 )
13
14 // A writer is a buffered, flushable writer.
15 type writer interface {
16 WriteByte(byte) error
17 Flush() error
18 }
19
20 // An errWriteCloser is an io.WriteCloser that always returns a given error.
21 type errWriteCloser struct {
22 err error
23 }
24
25 func (e *errWriteCloser) Write([]byte) (int, error) {
26 return 0, e.err
27 }
28
29 func (e *errWriteCloser) Close() error {
30 return e.err
31 }
32
33 const (
34 // A code is a 12 bit value, stored as a uint32 when encoding to avoid
35 // type conversions when shifting bits.
36 maxCode = 1<<12 - 1
37 invalidCode = 1<<32 - 1
38 // There are 1<<12 possible codes, which is an upper bound on the number of
39 // valid hash table entries at any given point in time. tableSize is 4x that.
40 tableSize = 4 * 1 << 12
41 tableMask = tableSize - 1
42 // A hash table entry is a uint32. Zero is an invalid entry since the
43 // lower 12 bits of a valid entry must be a non-literal code.
44 invalidEntry = 0
45 )
46
47 // encoder is LZW compressor.
48 type encoder struct {
49 // w is the writer that compressed bytes are written to.
50 w writer
51 // order, write, bits, nBits and width are the state for
52 // converting a code stream into a byte stream.
53 order Order
54 write func(*encoder, uint32) error
55 bits uint32
56 nBits uint
57 width uint
58 // litWidth is the width in bits of literal codes.
59 litWidth uint
60 // hi is the code implied by the next code emission.
61 // overflow is the code at which hi overflows the code width.
62 hi, overflow uint32
63 // savedCode is the accumulated code at the end of the most recent Write
64 // call. It is equal to invalidCode if there was no such call.
65 savedCode uint32
66 // err is the first error encountered during writing. Closing the encoder
67 // will make any future Write calls return errClosed
68 err error
69 // table is the hash table from 20-bit keys to 12-bit values. Each table
70 // entry contains key<<12|val and collisions resolve by linear probing.
71 // The keys consist of a 12-bit code prefix and an 8-bit byte suffix.
72 // The values are a 12-bit code.
73 table [tableSize]uint32
74 }
75
76 // writeLSB writes the code c for "Least Significant Bits first" data.
77 func (e *encoder) writeLSB(c uint32) error {
78 e.bits |= c << e.nBits
79 e.nBits += e.width
80 for e.nBits >= 8 {
81 if err := e.w.WriteByte(uint8(e.bits)); err != nil {
82 return err
83 }
84 e.bits >>= 8
85 e.nBits -= 8
86 }
87 return nil
88 }
89
90 // writeMSB writes the code c for "Most Significant Bits first" data.
91 func (e *encoder) writeMSB(c uint32) error {
92 e.bits |= c << (32 - e.width - e.nBits)
93 e.nBits += e.width
94 for e.nBits >= 8 {
95 if err := e.w.WriteByte(uint8(e.bits >> 24)); err != nil {
96 return err
97 }
98 e.bits <<= 8
99 e.nBits -= 8
100 }
101 return nil
102 }
103
104 // errOutOfCodes is an internal error that means that the encoder has run out
105 // of unused codes and a clear code needs to be sent next.
106 var errOutOfCodes = errors.New("lzw: out of codes")
107
108 // incHi increments e.hi and checks for both overflow and running out of
109 // unused codes. In the latter case, incHi sends a clear code, resets the
110 // encoder state and returns errOutOfCodes.
111 func (e *encoder) incHi() error {
112 e.hi++
113 if e.hi == e.overflow {
114 e.width++
115 e.overflow <<= 1
116 }
117 if e.hi == maxCode {
118 clear := uint32(1) << e.litWidth
119 if err := e.write(e, clear); err != nil {
120 return err
121 }
122 e.width = uint(e.litWidth) + 1
123 e.hi = clear + 1
124 e.overflow = clear << 1
125 for i := range e.table {
126 e.table[i] = invalidEntry
127 }
128 return errOutOfCodes
129 }
130 return nil
131 }
132
133 // Write writes a compressed representation of p to e's underlying writer.
134 func (e *encoder) Write(p []byte) (int, error) {
135 if e.err != nil {
136 return 0, e.err
137 }
138 if len(p) == 0 {
139 return 0, nil
140 }
141 litMask := uint32(1<<e.litWidth - 1)
142 code := e.savedCode
143 if code == invalidCode {
144 // The first code sent is always a literal code.
145 code, p = uint32(p[0])&litMask, p[1:]
146 }
147 loop:
148 for _, x := range p {
149 literal := uint32(x) & litMask
150 key := code<<8 | literal
151 // If there is a hash table hit for this key then we continue the loop
152 // and do not emit a code yet.
153 hash := (key>>12 ^ key) & tableMask
154 for h, t := hash, e.table[hash]; t != invalidEntry; {
155 if key == t>>12 {
156 code = t & maxCode
157 continue loop
158 }
159 h = (h + 1) & tableMask
160 t = e.table[h]
161 }
162 // Otherwise, write the current code, and literal becomes the start of
163 // the next emitted code.
164 if e.err = e.write(e, code); e.err != nil {
165 return 0, e.err
166 }
167 code = literal
168 // Increment e.hi, the next implied code. If we run out of codes, reset
169 // the encoder state (including clearing the hash table) and continue.
170 if err := e.incHi(); err != nil {
171 if err == errOutOfCodes {
172 continue
173 }
174 e.err = err
175 return 0, e.err
176 }
177 // Otherwise, insert key -> e.hi into the map that e.table represents.
178 for {
179 if e.table[hash] == invalidEntry {
180 e.table[hash] = (key << 12) | e.hi
181 break
182 }
183 hash = (hash + 1) & tableMask
184 }
185 }
186 e.savedCode = code
187 return len(p), nil
188 }
189
190 // Close closes the encoder, flushing any pending output. It does not close or
191 // flush e's underlying writer.
192 func (e *encoder) Close() error {
193 if e.err != nil {
194 if e.err == errClosed {
195 return nil
196 }
197 return e.err
198 }
199 // Make any future calls to Write return errClosed.
200 e.err = errClosed
201 // Write the savedCode if valid.
202 if e.savedCode != invalidCode {
203 if err := e.write(e, e.savedCode); err != nil {
204 return err
205 }
206 if err := e.incHi(); err != nil && err != errOutOfCodes {
207 return err
208 }
209 }
210 // Write the eof code.
211 eof := uint32(1)<<e.litWidth + 1
212 if err := e.write(e, eof); err != nil {
213 return err
214 }
215 // Write the final bits.
216 if e.nBits > 0 {
217 if e.order == MSB {
218 e.bits >>= 24
219 }
220 if err := e.w.WriteByte(uint8(e.bits)); err != nil {
221 return err
222 }
223 }
224 return e.w.Flush()
225 }
226
227 // NewWriter creates a new io.WriteCloser that satisfies writes by compressing
228 // the data and writing it to w.
229 // It is the caller's responsibility to call Close on the WriteCloser when
230 // finished writing.
231 // The number of bits to use for literal codes, litWidth, must be in the
232 // range [2,8] and is typically 8.
233 func NewWriter(w io.Writer, order Order, litWidth int) io.WriteCloser {
234 var write func(*encoder, uint32) error
235 switch order {
236 case LSB:
237 write = (*encoder).writeLSB
238 case MSB:
239 write = (*encoder).writeMSB
240 default:
241 return &errWriteCloser{errors.New("lzw: unknown order")}
242 }
243 if litWidth < 2 || 8 < litWidth {
244 return &errWriteCloser{fmt.Errorf("lzw: litWidth %d out of range", litWidth)}
245 }
246 bw, ok := w.(writer)
247 if !ok {
248 bw = bufio.NewWriter(w)
249 }
250 lw := uint(litWidth)
251 return &encoder{
252 w: bw,
253 order: order,
254 write: write,
255 width: 1 + lw,
256 litWidth: lw,
257 hi: 1<<lw + 1,
258 overflow: 1 << (lw + 1),
259 savedCode: invalidCode,
260 }
261 }