summaryrefslogtreecommitdiff
path: root/src/pkg/compress/lzw/writer.go
blob: 961b25f94f5fa6d03bcbdd549534548d69486fcb (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package lzw

import (
	"bufio"
	"errors"
	"fmt"
	"io"
)

// A writer is a buffered, flushable writer.
type writer interface {
	io.ByteWriter
	Flush() error
}

// An errWriteCloser is an io.WriteCloser that always returns a given error.
type errWriteCloser struct {
	err error
}

func (e *errWriteCloser) Write([]byte) (int, error) {
	return 0, e.err
}

func (e *errWriteCloser) Close() error {
	return e.err
}

const (
	// A code is a 12 bit value, stored as a uint32 when encoding to avoid
	// type conversions when shifting bits.
	maxCode     = 1<<12 - 1
	invalidCode = 1<<32 - 1
	// There are 1<<12 possible codes, which is an upper bound on the number of
	// valid hash table entries at any given point in time. tableSize is 4x that.
	tableSize = 4 * 1 << 12
	tableMask = tableSize - 1
	// A hash table entry is a uint32. Zero is an invalid entry since the
	// lower 12 bits of a valid entry must be a non-literal code.
	invalidEntry = 0
)

// encoder is LZW compressor.
type encoder struct {
	// w is the writer that compressed bytes are written to.
	w writer
	// order, write, bits, nBits and width are the state for
	// converting a code stream into a byte stream.
	order Order
	write func(*encoder, uint32) error
	bits  uint32
	nBits uint
	width uint
	// litWidth is the width in bits of literal codes.
	litWidth uint
	// hi is the code implied by the next code emission.
	// overflow is the code at which hi overflows the code width.
	hi, overflow uint32
	// savedCode is the accumulated code at the end of the most recent Write
	// call. It is equal to invalidCode if there was no such call.
	savedCode uint32
	// err is the first error encountered during writing. Closing the encoder
	// will make any future Write calls return errClosed
	err error
	// table is the hash table from 20-bit keys to 12-bit values. Each table
	// entry contains key<<12|val and collisions resolve by linear probing.
	// The keys consist of a 12-bit code prefix and an 8-bit byte suffix.
	// The values are a 12-bit code.
	table [tableSize]uint32
}

// writeLSB writes the code c for "Least Significant Bits first" data.
func (e *encoder) writeLSB(c uint32) error {
	e.bits |= c << e.nBits
	e.nBits += e.width
	for e.nBits >= 8 {
		if err := e.w.WriteByte(uint8(e.bits)); err != nil {
			return err
		}
		e.bits >>= 8
		e.nBits -= 8
	}
	return nil
}

// writeMSB writes the code c for "Most Significant Bits first" data.
func (e *encoder) writeMSB(c uint32) error {
	e.bits |= c << (32 - e.width - e.nBits)
	e.nBits += e.width
	for e.nBits >= 8 {
		if err := e.w.WriteByte(uint8(e.bits >> 24)); err != nil {
			return err
		}
		e.bits <<= 8
		e.nBits -= 8
	}
	return nil
}

// errOutOfCodes is an internal error that means that the encoder has run out
// of unused codes and a clear code needs to be sent next.
var errOutOfCodes = errors.New("lzw: out of codes")

// incHi increments e.hi and checks for both overflow and running out of
// unused codes. In the latter case, incHi sends a clear code, resets the
// encoder state and returns errOutOfCodes.
func (e *encoder) incHi() error {
	e.hi++
	if e.hi == e.overflow {
		e.width++
		e.overflow <<= 1
	}
	if e.hi == maxCode {
		clear := uint32(1) << e.litWidth
		if err := e.write(e, clear); err != nil {
			return err
		}
		e.width = uint(e.litWidth) + 1
		e.hi = clear + 1
		e.overflow = clear << 1
		for i := range e.table {
			e.table[i] = invalidEntry
		}
		return errOutOfCodes
	}
	return nil
}

// Write writes a compressed representation of p to e's underlying writer.
func (e *encoder) Write(p []byte) (n int, err error) {
	if e.err != nil {
		return 0, e.err
	}
	if len(p) == 0 {
		return 0, nil
	}
	n = len(p)
	litMask := uint32(1<<e.litWidth - 1)
	code := e.savedCode
	if code == invalidCode {
		// The first code sent is always a literal code.
		code, p = uint32(p[0])&litMask, p[1:]
	}
loop:
	for _, x := range p {
		literal := uint32(x) & litMask
		key := code<<8 | literal
		// If there is a hash table hit for this key then we continue the loop
		// and do not emit a code yet.
		hash := (key>>12 ^ key) & tableMask
		for h, t := hash, e.table[hash]; t != invalidEntry; {
			if key == t>>12 {
				code = t & maxCode
				continue loop
			}
			h = (h + 1) & tableMask
			t = e.table[h]
		}
		// Otherwise, write the current code, and literal becomes the start of
		// the next emitted code.
		if e.err = e.write(e, code); e.err != nil {
			return 0, e.err
		}
		code = literal
		// Increment e.hi, the next implied code. If we run out of codes, reset
		// the encoder state (including clearing the hash table) and continue.
		if err1 := e.incHi(); err1 != nil {
			if err1 == errOutOfCodes {
				continue
			}
			e.err = err1
			return 0, e.err
		}
		// Otherwise, insert key -> e.hi into the map that e.table represents.
		for {
			if e.table[hash] == invalidEntry {
				e.table[hash] = (key << 12) | e.hi
				break
			}
			hash = (hash + 1) & tableMask
		}
	}
	e.savedCode = code
	return n, nil
}

// Close closes the encoder, flushing any pending output. It does not close or
// flush e's underlying writer.
func (e *encoder) Close() error {
	if e.err != nil {
		if e.err == errClosed {
			return nil
		}
		return e.err
	}
	// Make any future calls to Write return errClosed.
	e.err = errClosed
	// Write the savedCode if valid.
	if e.savedCode != invalidCode {
		if err := e.write(e, e.savedCode); err != nil {
			return err
		}
		if err := e.incHi(); err != nil && err != errOutOfCodes {
			return err
		}
	}
	// Write the eof code.
	eof := uint32(1)<<e.litWidth + 1
	if err := e.write(e, eof); err != nil {
		return err
	}
	// Write the final bits.
	if e.nBits > 0 {
		if e.order == MSB {
			e.bits >>= 24
		}
		if err := e.w.WriteByte(uint8(e.bits)); err != nil {
			return err
		}
	}
	return e.w.Flush()
}

// NewWriter creates a new io.WriteCloser.
// Writes to the returned io.WriteCloser are compressed and written to w.
// It is the caller's responsibility to call Close on the WriteCloser when
// finished writing.
// The number of bits to use for literal codes, litWidth, must be in the
// range [2,8] and is typically 8.
func NewWriter(w io.Writer, order Order, litWidth int) io.WriteCloser {
	var write func(*encoder, uint32) error
	switch order {
	case LSB:
		write = (*encoder).writeLSB
	case MSB:
		write = (*encoder).writeMSB
	default:
		return &errWriteCloser{errors.New("lzw: unknown order")}
	}
	if litWidth < 2 || 8 < litWidth {
		return &errWriteCloser{fmt.Errorf("lzw: litWidth %d out of range", litWidth)}
	}
	bw, ok := w.(writer)
	if !ok {
		bw = bufio.NewWriter(w)
	}
	lw := uint(litWidth)
	return &encoder{
		w:         bw,
		order:     order,
		write:     write,
		width:     1 + lw,
		litWidth:  lw,
		hi:        1<<lw + 1,
		overflow:  1 << (lw + 1),
		savedCode: invalidCode,
	}
}