code review 4435051: image/jpeg: add an encoder.

src/pkg/image/jpeg/writer.go

+// 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 jpeg
+
+import (
+        "bufio"
+        "image"
+        "image/ycbcr"
+        "io"
+        "os"
+)
+
+// min returns the minimum of two integers.
+func min(x, y int) int {
+        if x < y {
+                return x
+        }
+        return y
+}
+
+// div returns a/b rounded to the nearest integer, instead of rounded to zero.
+func div(a int, b int) int {
+        if a >= 0 {
+                return (a + (b >> 1)) / b
+        }
+        return -((-a + (b >> 1)) / b)
+}
+
+// bitCount counts the number of bits needed to hold an integer.
+var bitCount = [256]byte{
+        0, 1, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4,
+        5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
+        6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6,
+        6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6,
+        7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,
+        7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,
+        7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,
+        7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,
+        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
+        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
+        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
+        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
+        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
+        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
+        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
+        8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8,
+}
+
+type quantIndex int
+
+const (
+        quantIndexLuminance quantIndex = iota
+        quantIndexChrominance
+        nQuantIndex
+)
+
+// unscaledQuant are the unscaled quantization tables. Each encoder copies and
+// scales the tables according to its quality parameter.
+var unscaledQuant = [nQuantIndex][blockSize]byte{
+        // Luminance.
+        {
+                16, 11, 10, 16, 24, 40, 51, 61,
+                12, 12, 14, 19, 26, 58, 60, 55,
+                14, 13, 16, 24, 40, 57, 69, 56,
+                14, 17, 22, 29, 51, 87, 80, 62,
+                18, 22, 37, 56, 68, 109, 103, 77,
+                24, 35, 55, 64, 81, 104, 113, 92,
+                49, 64, 78, 87, 103, 121, 120, 101,
+                72, 92, 95, 98, 112, 100, 103, 99,
+        },
+        // Chrominance.
+        {
+                17, 18, 24, 47, 99, 99, 99, 99,
+                18, 21, 26, 66, 99, 99, 99, 99,
+                24, 26, 56, 99, 99, 99, 99, 99,
+                47, 66, 99, 99, 99, 99, 99, 99,
+                99, 99, 99, 99, 99, 99, 99, 99,
+                99, 99, 99, 99, 99, 99, 99, 99,
+                99, 99, 99, 99, 99, 99, 99, 99,
+                99, 99, 99, 99, 99, 99, 99, 99,
+        },
+}
+
+type huffIndex int
+
+const (
+        huffIndexLuminanceDC huffIndex = iota
+        huffIndexLuminanceAC
+        huffIndexChrominanceDC
+        huffIndexChrominanceAC
+        nHuffIndex
+)
+
+// huffmanSpec specifies a Huffman encoding.
+type huffmanSpec struct {
+        // count[i] is the number of codes of length i bits.
+        count [16]byte
+        // value[i] is the decoded value of the i'th codeword.
+        value []byte
+}
+
+// theHuffmanSpec is the Huffman encoding specifications.
+// This encoder uses the same Huffman encoding for all images.
+var theHuffmanSpec = [nHuffIndex]huffmanSpec{
+        // Luminance DC.
+        {
+                [16]byte{0, 1, 5, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0},
+                []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11},
+        },
+        // Luminance AC.
+        {
+                [16]byte{0, 2, 1, 3, 3, 2, 4, 3, 5, 5, 4, 4, 0, 0, 1, 125},
+                []byte{
+                        0x01, 0x02, 0x03, 0x00, 0x04, 0x11, 0x05, 0x12,
+                        0x21, 0x31, 0x41, 0x06, 0x13, 0x51, 0x61, 0x07,
+                        0x22, 0x71, 0x14, 0x32, 0x81, 0x91, 0xa1, 0x08,
+                        0x23, 0x42, 0xb1, 0xc1, 0x15, 0x52, 0xd1, 0xf0,
+                        0x24, 0x33, 0x62, 0x72, 0x82, 0x09, 0x0a, 0x16,
+                        0x17, 0x18, 0x19, 0x1a, 0x25, 0x26, 0x27, 0x28,
+                        0x29, 0x2a, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39,
+                        0x3a, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49,
+                        0x4a, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59,
+                        0x5a, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69,
+                        0x6a, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79,
+                        0x7a, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89,
+                        0x8a, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98,
+                        0x99, 0x9a, 0xa2, 0xa3, 0xa4, 0xa5, 0xa6, 0xa7,
+                        0xa8, 0xa9, 0xaa, 0xb2, 0xb3, 0xb4, 0xb5, 0xb6,
+                        0xb7, 0xb8, 0xb9, 0xba, 0xc2, 0xc3, 0xc4, 0xc5,
+                        0xc6, 0xc7, 0xc8, 0xc9, 0xca, 0xd2, 0xd3, 0xd4,
+                        0xd5, 0xd6, 0xd7, 0xd8, 0xd9, 0xda, 0xe1, 0xe2,
+                        0xe3, 0xe4, 0xe5, 0xe6, 0xe7, 0xe8, 0xe9, 0xea,
+                        0xf1, 0xf2, 0xf3, 0xf4, 0xf5, 0xf6, 0xf7, 0xf8,
+                        0xf9, 0xfa,
+                },
+        },
+        // Chrominance DC.
+        {
+                [16]byte{0, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0},
+                []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11},
+        },
+        // Chrominance AC.
+        {
+                [16]byte{0, 2, 1, 2, 4, 4, 3, 4, 7, 5, 4, 4, 0, 1, 2, 119},
+                []byte{
+                        0x00, 0x01, 0x02, 0x03, 0x11, 0x04, 0x05, 0x21,
+                        0x31, 0x06, 0x12, 0x41, 0x51, 0x07, 0x61, 0x71,
+                        0x13, 0x22, 0x32, 0x81, 0x08, 0x14, 0x42, 0x91,
+                        0xa1, 0xb1, 0xc1, 0x09, 0x23, 0x33, 0x52, 0xf0,
+                        0x15, 0x62, 0x72, 0xd1, 0x0a, 0x16, 0x24, 0x34,
+                        0xe1, 0x25, 0xf1, 0x17, 0x18, 0x19, 0x1a, 0x26,
+                        0x27, 0x28, 0x29, 0x2a, 0x35, 0x36, 0x37, 0x38,
+                        0x39, 0x3a, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48,
+                        0x49, 0x4a, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58,
+                        0x59, 0x5a, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68,
+                        0x69, 0x6a, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78,
+                        0x79, 0x7a, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87,
+                        0x88, 0x89, 0x8a, 0x92, 0x93, 0x94, 0x95, 0x96,
+                        0x97, 0x98, 0x99, 0x9a, 0xa2, 0xa3, 0xa4, 0xa5,
+                        0xa6, 0xa7, 0xa8, 0xa9, 0xaa, 0xb2, 0xb3, 0xb4,
+                        0xb5, 0xb6, 0xb7, 0xb8, 0xb9, 0xba, 0xc2, 0xc3,
+                        0xc4, 0xc5, 0xc6, 0xc7, 0xc8, 0xc9, 0xca, 0xd2,
+                        0xd3, 0xd4, 0xd5, 0xd6, 0xd7, 0xd8, 0xd9, 0xda,
+                        0xe2, 0xe3, 0xe4, 0xe5, 0xe6, 0xe7, 0xe8, 0xe9,
+                        0xea, 0xf2, 0xf3, 0xf4, 0xf5, 0xf6, 0xf7, 0xf8,
+                        0xf9, 0xfa,
+                },
+        },
+}
+
+// huffmanLUT is a compiled look-up table representation of a huffmanSpec.
+// Each value maps to an uint32 of which the 8 most significant bits hold the

[r, 2011/04/18 19:30:42]

s/an/a/

[nigeltao, 2011/04/19 00:47:06]

Done.

+// codeword size in bits and the 24 least significant bits hold the codeword.
+// The maximum codeword size is 16 bits.
+type huffmanLUT []uint32
+
+func (h *huffmanLUT) init(s huffmanSpec) {
+        maxValue := 0
+        for _, v := range s.value {
+                if int(v) > maxValue {
+                        maxValue = int(v)
+                }
+        }
+        *h = make([]uint32, maxValue+1)
+        code, k := uint32(0), 0
+        for i := 0; i < len(s.count); i++ {
+                nBits := uint32(i+1) << 24
+                for j := uint8(0); j < s.count[i]; j++ {
+                        (*h)[s.value[k]] = nBits | code
+                        code++
+                        k++
+                }
+                code <<= 1
+        }
+}
+
+// theHuffmanLUT are compiled representations of theHuffmanSpec.
+var theHuffmanLUT [4]huffmanLUT
+
+func init() {
+        for i, s := range theHuffmanSpec {
+                theHuffmanLUT[i].init(s)
+        }
+}
+
+// writer is a buffered writer.
+type writer interface {
+        Flush() os.Error
+        Write([]byte) (int, os.Error)
+        WriteByte(byte) os.Error
+        WriteString(string) (int, os.Error)
+}
+
+// An encoder wraps a writer's methods to save the first returned error as
+// e.err and return void.

[r, 2011/04/18 19:30:42]

void?

[nigeltao, 2011/04/19 00:47:06]

Rewritten.

+
+func (e *encoder) flush() {
+        if e.err != nil {
+                return
+        }
+        e.err = e.w.Flush()
+}
+
+func (e *encoder) write(p []byte) {
+        if e.err != nil {
+                return
+        }
+        _, e.err = e.w.Write(p)
+}
+
+func (e *encoder) writeByte(b byte) {
+        if e.err != nil {
+                return
+        }
+        e.err = e.w.WriteByte(b)
+}
+
+func (e *encoder) writeString(s string) {
+        if e.err != nil {
+                return
+        }
+        _, e.err = e.w.WriteString(s)
+}
+
+// encoder encodes an image to the JPEG format.

[r, 2011/04/18 19:30:42]

move this above its methods and combine the comments

[nigeltao, 2011/04/19 00:47:06]

Done.

+type encoder struct {
+        // w is the writer to write to. err is the first error encountered
+        // during writing.
+        w   writer
+        err os.Error
+        // bits and nBits are accumulated bits to write to w.
+        bits  uint32
+        nBits uint8
+        // quant is the scaled quantization tables.
+        quant [nQuantIndex][blockSize]byte
+}
+
+// emit emits the least significant nBits bits of bits to the bitstream.

[r, 2011/04/18 19:30:42]

nBits bits of bits to the bits. wow.

+// The precondition is bits < 1<<nBits && nBits <= 16.
+func (e *encoder) emit(bits uint32, nBits uint8) {
+        nBits += e.nBits
+        bits <<= 32 - nBits
+        bits |= e.bits
+        for nBits >= 8 {
+                b := uint8(bits >> 24)
+                e.writeByte(b)
+                if b == 255 {
+                        e.writeByte(0x00)
+                }
+                bits <<= 8
+                nBits -= 8
+        }
+        e.bits, e.nBits = bits, nBits
+}
+
+// emitHuff emits the given value with the given Huffman encoder.
+func (e *encoder) emitHuff(h huffIndex, value int) {
+        x := theHuffmanLUT[h][value]
+        e.emit(x&(1<<24-1), uint8(x>>24))
+}
+
+// emitHuffRLE emits a run of runLength copies of value encoded with the given
+// Huffman encoder.
+func (e *encoder) emitHuffRLE(h huffIndex, runLength, value int) {
+        a, b := value, value
+        if a < 0 {
+                a, b = -value, value-1
+        }
+        var nBits uint8
+        if a < 0x100 {
+                nBits = bitCount[a]
+        } else {
+                nBits = 8 + bitCount[a>>8]
+        }
+        e.emitHuff(h, runLength<<4|int(nBits))
+        if nBits > 0 {
+                e.emit(uint32(b)&(1<<nBits-1), nBits)
+        }
+}
+
+// writeMarkerHeader writes the header for a maker with the given length.

[r, 2011/04/18 19:30:42]

s/maker/marker/

[nigeltao, 2011/04/19 00:47:06]

Done.

+func (e *encoder) writeMarkerHeader(marker uint8, markerlen int) {
+        e.writeByte(0xff)
+        e.writeByte(marker)
+        e.writeByte(uint8(markerlen >> 8))
+        e.writeByte(uint8(markerlen & 0xff))

[r, 2011/04/18 19:30:42]

this is a fair bit of overhead. we could put a [4]byte into encoder and do a
single Write call.

[nigeltao, 2011/04/19 00:47:06]

Done.

+}
+
+// writeDQT writes the Define Quantization Table marker.
+func (e *encoder) writeDQT() {
+        markerlen := 2
+        for _, q := range e.quant {
+                markerlen += 1 + len(q)
+        }
+        e.writeMarkerHeader(dqtMarker, markerlen)
+        for i, q := range e.quant {
+                e.writeByte(uint8(i))
+                e.write(q[:])
+        }
+}
+
+// writeSOF0 writes the Start Of Frame (Baseline) marker.
+func (e *encoder) writeSOF0(size image.Point) {
+        markerlen := 8 + 3*nComponent
+        e.writeMarkerHeader(sof0Marker, markerlen)
+        e.writeByte(8) // 8-bit color.
+        e.writeByte(uint8(size.Y >> 8))
+        e.writeByte(uint8(size.Y & 0xff))
+        e.writeByte(uint8(size.X >> 8))
+        e.writeByte(uint8(size.X & 0xff))

[r, 2011/04/18 19:30:42]

ditto.

[nigeltao, 2011/04/19 00:47:06]

Done.

+        e.writeByte(nComponent)
+        for i := 0; i < nComponent; i++ {
+                e.writeByte(uint8(i + 1))
+                // We use 4:2:0 chroma subsampling.
+                e.writeByte("\x22\x11\x11"[i])
+                e.writeByte("\x00\x01\x01"[i])

[r, 2011/04/18 19:30:42]

ditto

[nigeltao, 2011/04/19 00:47:06]

Done.

+        }
+}
+
+// writeDHT writes the Define Huffman Table marker.
+func (e *encoder) writeDHT() {
+        markerlen := 2
+        for _, s := range theHuffmanSpec {
+                markerlen += 1 + 16 + len(s.value)
+        }
+        e.writeMarkerHeader(dhtMarker, markerlen)
+        for i, s := range theHuffmanSpec {
+                e.writeByte("\x00\x10\x01\x11"[i])
+                e.write(s.count[:])
+                e.write(s.value)

[r, 2011/04/18 19:30:42]

ditto

[nigeltao, 2011/04/19 00:47:06]

I don't think this applies here; s.value is a slice. Besides, the underlying
writer e.w is presumably buffered (since it implements a WriteByte method), so
I'm not sure if adding additional buffering would be a significant win.

+        }
+}
+
+// writeBlock writes a block of pixel data using the given quantization table,
+// returning the post-quantized DC value of the DCT-transformed block.
+func (e *encoder) writeBlock(b *block, q quantIndex, prevDC int) int {
+        fdct(b)
+        // Emit the DC delta.
+        dc := div(b[0], (8 * int(e.quant[q][0])))
+        e.emitHuffRLE(huffIndex(2*q+0), 0, dc-prevDC)
+        // Emit the AC components.
+        h, runLength := huffIndex(2*q+1), 0
+        for k := 1; k < blockSize; k++ {
+                ac := div(b[unzig[k]], (8 * int(e.quant[q][k])))
+                if ac == 0 {
+                        runLength++
+                } else {
+                        for runLength > 15 {
+                                e.emitHuff(h, 0xf0)
+                                runLength -= 16
+                        }
+                        e.emitHuffRLE(h, runLength, ac)
+                        runLength = 0
+                }
+        }
+        if runLength > 0 {
+                e.emitHuff(h, 0x00)
+        }
+        return dc
+}
+
+// toYCbCr converts the 8x8 region of m whose top-left corner is p to its
+// YCbCr values.
+func toYCbCr(m image.Image, p image.Point, yBlock, cbBlock, crBlock *block) {
+        b := m.Bounds()
+        xmax := b.Max.X - 1
+        ymax := b.Max.Y - 1
+        for j := 0; j < 8; j++ {
+                for i := 0; i < 8; i++ {
+                        r, g, b, _ := m.At(min(p.X+i, xmax), min(p.Y+j, ymax)).RGBA()
+                        yy, cb, cr := ycbcr.RGBToYCbCr(uint8(r>>8), uint8(g>>8), uint8(b>>8))
+                        yBlock[8*j+i] = int(yy)
+                        cbBlock[8*j+i] = int(cb)
+                        crBlock[8*j+i] = int(cr)
+                }
+        }
+}
+
+// scale scales the 16x16 region represented by the 4 src blocks to the 8x8
+// dst block.
+func scale(dst *block, src *[4]block) {
+        for i := 0; i < 4; i++ {
+                dstOff := (i&2)<<4 | (i&1)<<2
+                for y := 0; y < 4; y++ {
+                        for x := 0; x < 4; x++ {
+                                j := 16*y + 2*x
+                                sum := src[i][j] + src[i][j+1] + src[i][j+8] + src[i][j+9]
+                                dst[8*y+x+dstOff] = (sum + 2) >> 2
+                        }
+                }
+        }
+}
+
+// writeSOS writes the StartOfScan marker.
+func (e *encoder) writeSOS(m image.Image) {
+        // Write the SOS marker "\xff\xda" followed by 12 bytes:
+        //      - the marker length "\x00\x0c",
+        //      - the number of components "\x03",
+        //      - component 1 uses DC table 0 and AC table 0 "\x01\x00",
+        //      - component 2 uses DC table 1 and AC table 1 "\x02\x11",
+        //      - component 3 uses DC table 1 and AC table 1 "\x03\x11",
+        //      - padding "\x00\x00\x00".
+        e.writeString("\xff\xda\x00\x0c\x03\x01\x00\x02\x11\x03\x11\x00\x00\x00")

[r, 2011/04/18 19:30:42]

depending on the implementation of the writer, writeString may copy.  better
than this approach is to have a global
var sosHeader = []byte("....")

also i'm not sure you're getting much readability from turning these bytes into
strings.

[nigeltao, 2011/04/19 00:47:06]

Done.

+        var (
+                // Scratch buffers to hold the YCbCr values.
+                yBlock  block
+                cbBlock [4]block
+                crBlock [4]block
+                cBlock  block

[r, 2011/04/18 19:30:42]

ok for now, but these are more unnecessary allocations.
probably not important, because we only do one SOS per image (right?)

[nigeltao, 2011/04/19 00:47:06]

There is only one SOS per image. Also, eventually, escape analysis should be
able to discover that these blocks are stack-allocatable.

+                // DC components are delta-encoded.
+                lastY, lastCb, lastCr int
+        )
+        bounds := m.Bounds()
+        for y := bounds.Min.Y; y < bounds.Max.Y; y += 16 {
+                for x := bounds.Min.X; x < bounds.Max.X; x += 16 {
+                        for i := 0; i < 4; i++ {
+                                xOff := (i & 1) * 8
+                                yOff := (i & 2) * 4
+                                p := image.Point{x + xOff, y + yOff}
+                                toYCbCr(m, p, &yBlock, &cbBlock[i], &crBlock[i])
+                                lastY = e.writeBlock(&yBlock, 0, lastY)
+                        }
+                        scale(&cBlock, &cbBlock)
+                        lastCb = e.writeBlock(&cBlock, 1, lastCb)
+                        scale(&cBlock, &crBlock)
+                        lastCr = e.writeBlock(&cBlock, 1, lastCr)
+                }
+        }
+        // Pad the last byte with 1's.
+        e.emit(0x7f, 7)
+}
+
+// DefaultQuality is the default quality encoding parameter.
+const DefaultQuality = 75
+
+// Options are the encoding parameters.
+// Quality ranges from 1 to 100 inclusive, higher is better.
+type Options struct {
+        Quality int
+}
+
+// Encode writes the Image m to w in JPEG 4:2:0 baseline format with the given
+// options. Default parameters are used if a nil *Options is passed.
+func Encode(w io.Writer, m image.Image, o *Options) os.Error {
+        b := m.Bounds()
+        if b.Dx() >= 1<<16 || b.Dy() >= 1<<16 {
+                return os.NewError("jpeg: image is too large to encode")
+        }
+        var e encoder
+        if ww, ok := w.(writer); ok {
+                e.w = ww
+        } else {
+                e.w = bufio.NewWriter(w)
+        }
+        // Clip quality to [1, 100].
+        quality := DefaultQuality
+        if o != nil {
+                quality = o.Quality
+                if quality < 1 {
+                        quality = 1
+                } else if quality > 100 {
+                        quality = 100
+                }
+        }
+        // Convert from a quality rating to a scaling factor.
+        var scale int
+        if quality < 50 {
+                scale = 5000 / quality
+        } else {
+                scale = 200 - quality*2
+        }
+        // Initialize the quantization tables.
+        for i := range e.quant {
+                for j := range e.quant[i] {
+                        x := int(unscaledQuant[i][j])
+                        x = (x*scale + 50) / 100
+                        if x < 1 {
+                                x = 1
+                        } else if x > 255 {
+                                x = 255
+                        }
+                        e.quant[i][j] = uint8(x)
+                }
+        }
+        // Write the Start Of Image marker.
+        e.writeString("\xff\xd8")
+        // Write the quantization tables.
+        e.writeDQT()
+        // Write the image dimensions.
+        e.writeSOF0(b.Size())
+        // Write the Huffman tables.
+        e.writeDHT()
+        // Write the image data.
+        e.writeSOS(m)
+        // Write the End Of Image marker.
+        e.writeString("\xff\xd9")
+        e.flush()
+        return e.err
+}