1// Copyright 2013 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
5package gif
6
7import (
8	"bufio"
9	"bytes"
10	"compress/lzw"
11	"errors"
12	"image"
13	"image/color"
14	"image/color/palette"
15	"image/draw"
16	"internal/byteorder"
17	"io"
18)
19
20// Graphic control extension fields.
21const (
22	gcLabel     = 0xF9
23	gcBlockSize = 0x04
24)
25
26var log2Lookup = [8]int{2, 4, 8, 16, 32, 64, 128, 256}
27
28func log2(x int) int {
29	for i, v := range log2Lookup {
30		if x <= v {
31			return i
32		}
33	}
34	return -1
35}
36
37// writer is a buffered writer.
38type writer interface {
39	Flush() error
40	io.Writer
41	io.ByteWriter
42}
43
44// encoder encodes an image to the GIF format.
45type encoder struct {
46	// w is the writer to write to. err is the first error encountered during
47	// writing. All attempted writes after the first error become no-ops.
48	w   writer
49	err error
50	// g is a reference to the data that is being encoded.
51	g GIF
52	// globalCT is the size in bytes of the global color table.
53	globalCT int
54	// buf is a scratch buffer. It must be at least 256 for the blockWriter.
55	buf              [256]byte
56	globalColorTable [3 * 256]byte
57	localColorTable  [3 * 256]byte
58}
59
60// blockWriter writes the block structure of GIF image data, which
61// comprises (n, (n bytes)) blocks, with 1 <= n <= 255. It is the
62// writer given to the LZW encoder, which is thus immune to the
63// blocking.
64type blockWriter struct {
65	e *encoder
66}
67
68func (b blockWriter) setup() {
69	b.e.buf[0] = 0
70}
71
72func (b blockWriter) Flush() error {
73	return b.e.err
74}
75
76func (b blockWriter) WriteByte(c byte) error {
77	if b.e.err != nil {
78		return b.e.err
79	}
80
81	// Append c to buffered sub-block.
82	b.e.buf[0]++
83	b.e.buf[b.e.buf[0]] = c
84	if b.e.buf[0] < 255 {
85		return nil
86	}
87
88	// Flush block
89	b.e.write(b.e.buf[:256])
90	b.e.buf[0] = 0
91	return b.e.err
92}
93
94// blockWriter must be an io.Writer for lzw.NewWriter, but this is never
95// actually called.
96func (b blockWriter) Write(data []byte) (int, error) {
97	for i, c := range data {
98		if err := b.WriteByte(c); err != nil {
99			return i, err
100		}
101	}
102	return len(data), nil
103}
104
105func (b blockWriter) close() {
106	// Write the block terminator (0x00), either by itself, or along with a
107	// pending sub-block.
108	if b.e.buf[0] == 0 {
109		b.e.writeByte(0)
110	} else {
111		n := uint(b.e.buf[0])
112		b.e.buf[n+1] = 0
113		b.e.write(b.e.buf[:n+2])
114	}
115	b.e.flush()
116}
117
118func (e *encoder) flush() {
119	if e.err != nil {
120		return
121	}
122	e.err = e.w.Flush()
123}
124
125func (e *encoder) write(p []byte) {
126	if e.err != nil {
127		return
128	}
129	_, e.err = e.w.Write(p)
130}
131
132func (e *encoder) writeByte(b byte) {
133	if e.err != nil {
134		return
135	}
136	e.err = e.w.WriteByte(b)
137}
138
139func (e *encoder) writeHeader() {
140	if e.err != nil {
141		return
142	}
143	_, e.err = io.WriteString(e.w, "GIF89a")
144	if e.err != nil {
145		return
146	}
147
148	// Logical screen width and height.
149	byteorder.LePutUint16(e.buf[0:2], uint16(e.g.Config.Width))
150	byteorder.LePutUint16(e.buf[2:4], uint16(e.g.Config.Height))
151	e.write(e.buf[:4])
152
153	if p, ok := e.g.Config.ColorModel.(color.Palette); ok && len(p) > 0 {
154		paddedSize := log2(len(p)) // Size of Global Color Table: 2^(1+n).
155		e.buf[0] = fColorTable | uint8(paddedSize)
156		e.buf[1] = e.g.BackgroundIndex
157		e.buf[2] = 0x00 // Pixel Aspect Ratio.
158		e.write(e.buf[:3])
159		var err error
160		e.globalCT, err = encodeColorTable(e.globalColorTable[:], p, paddedSize)
161		if err != nil && e.err == nil {
162			e.err = err
163			return
164		}
165		e.write(e.globalColorTable[:e.globalCT])
166	} else {
167		// All frames have a local color table, so a global color table
168		// is not needed.
169		e.buf[0] = 0x00
170		e.buf[1] = 0x00 // Background Color Index.
171		e.buf[2] = 0x00 // Pixel Aspect Ratio.
172		e.write(e.buf[:3])
173	}
174
175	// Add animation info if necessary.
176	if len(e.g.Image) > 1 && e.g.LoopCount >= 0 {
177		e.buf[0] = 0x21 // Extension Introducer.
178		e.buf[1] = 0xff // Application Label.
179		e.buf[2] = 0x0b // Block Size.
180		e.write(e.buf[:3])
181		_, err := io.WriteString(e.w, "NETSCAPE2.0") // Application Identifier.
182		if err != nil && e.err == nil {
183			e.err = err
184			return
185		}
186		e.buf[0] = 0x03 // Block Size.
187		e.buf[1] = 0x01 // Sub-block Index.
188		byteorder.LePutUint16(e.buf[2:4], uint16(e.g.LoopCount))
189		e.buf[4] = 0x00 // Block Terminator.
190		e.write(e.buf[:5])
191	}
192}
193
194func encodeColorTable(dst []byte, p color.Palette, size int) (int, error) {
195	if uint(size) >= uint(len(log2Lookup)) {
196		return 0, errors.New("gif: cannot encode color table with more than 256 entries")
197	}
198	for i, c := range p {
199		if c == nil {
200			return 0, errors.New("gif: cannot encode color table with nil entries")
201		}
202		var r, g, b uint8
203		// It is most likely that the palette is full of color.RGBAs, so they
204		// get a fast path.
205		if rgba, ok := c.(color.RGBA); ok {
206			r, g, b = rgba.R, rgba.G, rgba.B
207		} else {
208			rr, gg, bb, _ := c.RGBA()
209			r, g, b = uint8(rr>>8), uint8(gg>>8), uint8(bb>>8)
210		}
211		dst[3*i+0] = r
212		dst[3*i+1] = g
213		dst[3*i+2] = b
214	}
215	n := log2Lookup[size]
216	if n > len(p) {
217		// Pad with black.
218		clear(dst[3*len(p) : 3*n])
219	}
220	return 3 * n, nil
221}
222
223func (e *encoder) colorTablesMatch(localLen, transparentIndex int) bool {
224	localSize := 3 * localLen
225	if transparentIndex >= 0 {
226		trOff := 3 * transparentIndex
227		return bytes.Equal(e.globalColorTable[:trOff], e.localColorTable[:trOff]) &&
228			bytes.Equal(e.globalColorTable[trOff+3:localSize], e.localColorTable[trOff+3:localSize])
229	}
230	return bytes.Equal(e.globalColorTable[:localSize], e.localColorTable[:localSize])
231}
232
233func (e *encoder) writeImageBlock(pm *image.Paletted, delay int, disposal byte) {
234	if e.err != nil {
235		return
236	}
237
238	if len(pm.Palette) == 0 {
239		e.err = errors.New("gif: cannot encode image block with empty palette")
240		return
241	}
242
243	b := pm.Bounds()
244	if b.Min.X < 0 || b.Max.X >= 1<<16 || b.Min.Y < 0 || b.Max.Y >= 1<<16 {
245		e.err = errors.New("gif: image block is too large to encode")
246		return
247	}
248	if !b.In(image.Rectangle{Max: image.Point{e.g.Config.Width, e.g.Config.Height}}) {
249		e.err = errors.New("gif: image block is out of bounds")
250		return
251	}
252
253	transparentIndex := -1
254	for i, c := range pm.Palette {
255		if c == nil {
256			e.err = errors.New("gif: cannot encode color table with nil entries")
257			return
258		}
259		if _, _, _, a := c.RGBA(); a == 0 {
260			transparentIndex = i
261			break
262		}
263	}
264
265	if delay > 0 || disposal != 0 || transparentIndex != -1 {
266		e.buf[0] = sExtension  // Extension Introducer.
267		e.buf[1] = gcLabel     // Graphic Control Label.
268		e.buf[2] = gcBlockSize // Block Size.
269		if transparentIndex != -1 {
270			e.buf[3] = 0x01 | disposal<<2
271		} else {
272			e.buf[3] = 0x00 | disposal<<2
273		}
274		byteorder.LePutUint16(e.buf[4:6], uint16(delay)) // Delay Time (1/100ths of a second)
275
276		// Transparent color index.
277		if transparentIndex != -1 {
278			e.buf[6] = uint8(transparentIndex)
279		} else {
280			e.buf[6] = 0x00
281		}
282		e.buf[7] = 0x00 // Block Terminator.
283		e.write(e.buf[:8])
284	}
285	e.buf[0] = sImageDescriptor
286	byteorder.LePutUint16(e.buf[1:3], uint16(b.Min.X))
287	byteorder.LePutUint16(e.buf[3:5], uint16(b.Min.Y))
288	byteorder.LePutUint16(e.buf[5:7], uint16(b.Dx()))
289	byteorder.LePutUint16(e.buf[7:9], uint16(b.Dy()))
290	e.write(e.buf[:9])
291
292	// To determine whether or not this frame's palette is the same as the
293	// global palette, we can check a couple things. First, do they actually
294	// point to the same []color.Color? If so, they are equal so long as the
295	// frame's palette is not longer than the global palette...
296	paddedSize := log2(len(pm.Palette)) // Size of Local Color Table: 2^(1+n).
297	if gp, ok := e.g.Config.ColorModel.(color.Palette); ok && len(pm.Palette) <= len(gp) && &gp[0] == &pm.Palette[0] {
298		e.writeByte(0) // Use the global color table.
299	} else {
300		ct, err := encodeColorTable(e.localColorTable[:], pm.Palette, paddedSize)
301		if err != nil {
302			if e.err == nil {
303				e.err = err
304			}
305			return
306		}
307		// This frame's palette is not the very same slice as the global
308		// palette, but it might be a copy, possibly with one value turned into
309		// transparency by DecodeAll.
310		if ct <= e.globalCT && e.colorTablesMatch(len(pm.Palette), transparentIndex) {
311			e.writeByte(0) // Use the global color table.
312		} else {
313			// Use a local color table.
314			e.writeByte(fColorTable | uint8(paddedSize))
315			e.write(e.localColorTable[:ct])
316		}
317	}
318
319	litWidth := paddedSize + 1
320	if litWidth < 2 {
321		litWidth = 2
322	}
323	e.writeByte(uint8(litWidth)) // LZW Minimum Code Size.
324
325	bw := blockWriter{e: e}
326	bw.setup()
327	lzww := lzw.NewWriter(bw, lzw.LSB, litWidth)
328	if dx := b.Dx(); dx == pm.Stride {
329		_, e.err = lzww.Write(pm.Pix[:dx*b.Dy()])
330		if e.err != nil {
331			lzww.Close()
332			return
333		}
334	} else {
335		for i, y := 0, b.Min.Y; y < b.Max.Y; i, y = i+pm.Stride, y+1 {
336			_, e.err = lzww.Write(pm.Pix[i : i+dx])
337			if e.err != nil {
338				lzww.Close()
339				return
340			}
341		}
342	}
343	lzww.Close() // flush to bw
344	bw.close()   // flush to e.w
345}
346
347// Options are the encoding parameters.
348type Options struct {
349	// NumColors is the maximum number of colors used in the image.
350	// It ranges from 1 to 256.
351	NumColors int
352
353	// Quantizer is used to produce a palette with size NumColors.
354	// palette.Plan9 is used in place of a nil Quantizer.
355	Quantizer draw.Quantizer
356
357	// Drawer is used to convert the source image to the desired palette.
358	// draw.FloydSteinberg is used in place of a nil Drawer.
359	Drawer draw.Drawer
360}
361
362// EncodeAll writes the images in g to w in GIF format with the
363// given loop count and delay between frames.
364func EncodeAll(w io.Writer, g *GIF) error {
365	if len(g.Image) == 0 {
366		return errors.New("gif: must provide at least one image")
367	}
368
369	if len(g.Image) != len(g.Delay) {
370		return errors.New("gif: mismatched image and delay lengths")
371	}
372
373	e := encoder{g: *g}
374	// The GIF.Disposal, GIF.Config and GIF.BackgroundIndex fields were added
375	// in Go 1.5. Valid Go 1.4 code, such as when the Disposal field is omitted
376	// in a GIF struct literal, should still produce valid GIFs.
377	if e.g.Disposal != nil && len(e.g.Image) != len(e.g.Disposal) {
378		return errors.New("gif: mismatched image and disposal lengths")
379	}
380	if e.g.Config == (image.Config{}) {
381		p := g.Image[0].Bounds().Max
382		e.g.Config.Width = p.X
383		e.g.Config.Height = p.Y
384	} else if e.g.Config.ColorModel != nil {
385		if _, ok := e.g.Config.ColorModel.(color.Palette); !ok {
386			return errors.New("gif: GIF color model must be a color.Palette")
387		}
388	}
389
390	if ww, ok := w.(writer); ok {
391		e.w = ww
392	} else {
393		e.w = bufio.NewWriter(w)
394	}
395
396	e.writeHeader()
397	for i, pm := range g.Image {
398		disposal := uint8(0)
399		if g.Disposal != nil {
400			disposal = g.Disposal[i]
401		}
402		e.writeImageBlock(pm, g.Delay[i], disposal)
403	}
404	e.writeByte(sTrailer)
405	e.flush()
406	return e.err
407}
408
409// Encode writes the Image m to w in GIF format.
410func Encode(w io.Writer, m image.Image, o *Options) error {
411	// Check for bounds and size restrictions.
412	b := m.Bounds()
413	if b.Dx() >= 1<<16 || b.Dy() >= 1<<16 {
414		return errors.New("gif: image is too large to encode")
415	}
416
417	opts := Options{}
418	if o != nil {
419		opts = *o
420	}
421	if opts.NumColors < 1 || 256 < opts.NumColors {
422		opts.NumColors = 256
423	}
424	if opts.Drawer == nil {
425		opts.Drawer = draw.FloydSteinberg
426	}
427
428	pm, _ := m.(*image.Paletted)
429	if pm == nil {
430		if cp, ok := m.ColorModel().(color.Palette); ok {
431			pm = image.NewPaletted(b, cp)
432			for y := b.Min.Y; y < b.Max.Y; y++ {
433				for x := b.Min.X; x < b.Max.X; x++ {
434					pm.Set(x, y, cp.Convert(m.At(x, y)))
435				}
436			}
437		}
438	}
439	if pm == nil || len(pm.Palette) > opts.NumColors {
440		// Set pm to be a palettedized copy of m, including its bounds, which
441		// might not start at (0, 0).
442		//
443		// TODO: Pick a better sub-sample of the Plan 9 palette.
444		pm = image.NewPaletted(b, palette.Plan9[:opts.NumColors])
445		if opts.Quantizer != nil {
446			pm.Palette = opts.Quantizer.Quantize(make(color.Palette, 0, opts.NumColors), m)
447		}
448		opts.Drawer.Draw(pm, b, m, b.Min)
449	}
450
451	// When calling Encode instead of EncodeAll, the single-frame image is
452	// translated such that its top-left corner is (0, 0), so that the single
453	// frame completely fills the overall GIF's bounds.
454	if pm.Rect.Min != (image.Point{}) {
455		dup := *pm
456		dup.Rect = dup.Rect.Sub(dup.Rect.Min)
457		pm = &dup
458	}
459
460	return EncodeAll(w, &GIF{
461		Image: []*image.Paletted{pm},
462		Delay: []int{0},
463		Config: image.Config{
464			ColorModel: pm.Palette,
465			Width:      b.Dx(),
466			Height:     b.Dy(),
467		},
468	})
469}
470