huffman.mx raw

   1  // Copyright 2009 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 jpeg
   6  
   7  import (
   8  	"io"
   9  )
  10  
  11  // maxCodeLength is the maximum (inclusive) number of bits in a Huffman code.
  12  const maxCodeLength = 16
  13  
  14  // maxNCodes is the maximum (inclusive) number of codes in a Huffman tree.
  15  const maxNCodes = 256
  16  
  17  // lutSize is the log-2 size of the Huffman decoder's look-up table.
  18  const lutSize = 8
  19  
  20  // huffman is a Huffman decoder, specified in section C.
  21  type huffman struct {
  22  	// length is the number of codes in the tree.
  23  	nCodes int32
  24  	// lut is the look-up table for the next lutSize bits in the bit-stream.
  25  	// The high 8 bits of the uint16 are the encoded value. The low 8 bits
  26  	// are 1 plus the code length, or 0 if the value is too large to fit in
  27  	// lutSize bits.
  28  	lut [1 << lutSize]uint16
  29  	// vals are the decoded values, sorted by their encoding.
  30  	vals [maxNCodes]uint8
  31  	// minCodes[i] is the minimum code of length i, or -1 if there are no
  32  	// codes of that length.
  33  	minCodes [maxCodeLength]int32
  34  	// maxCodes[i] is the maximum code of length i, or -1 if there are no
  35  	// codes of that length.
  36  	maxCodes [maxCodeLength]int32
  37  	// valsIndices[i] is the index into vals of minCodes[i].
  38  	valsIndices [maxCodeLength]int32
  39  }
  40  
  41  // errShortHuffmanData means that an unexpected EOF occurred while decoding
  42  // Huffman data.
  43  var errShortHuffmanData = FormatError("short Huffman data")
  44  
  45  // ensureNBits reads bytes from the byte buffer to ensure that d.bits.n is at
  46  // least n. For best performance (avoiding function calls inside hot loops),
  47  // the caller is the one responsible for first checking that d.bits.n < n.
  48  func (d *decoder) ensureNBits(n int32) error {
  49  	for {
  50  		c, err := d.readByteStuffedByte()
  51  		if err != nil {
  52  			if err == io.ErrUnexpectedEOF {
  53  				return errShortHuffmanData
  54  			}
  55  			return err
  56  		}
  57  		d.bits.a = d.bits.a<<8 | uint32(c)
  58  		d.bits.n += 8
  59  		if d.bits.m == 0 {
  60  			d.bits.m = 1 << 7
  61  		} else {
  62  			d.bits.m <<= 8
  63  		}
  64  		if d.bits.n >= n {
  65  			break
  66  		}
  67  	}
  68  	return nil
  69  }
  70  
  71  // receiveExtend is the composition of RECEIVE and EXTEND, specified in section
  72  // F.2.2.1.
  73  func (d *decoder) receiveExtend(t uint8) (int32, error) {
  74  	if d.bits.n < int32(t) {
  75  		if err := d.ensureNBits(int32(t)); err != nil {
  76  			return 0, err
  77  		}
  78  	}
  79  	d.bits.n -= int32(t)
  80  	d.bits.m >>= t
  81  	s := int32(1) << t
  82  	x := int32(d.bits.a>>uint8(d.bits.n)) & (s - 1)
  83  	if x < s>>1 {
  84  		x += ((-1) << t) + 1
  85  	}
  86  	return x, nil
  87  }
  88  
  89  // processDHT processes a Define Huffman Table marker, and initializes a huffman
  90  // struct from its contents. Specified in section B.2.4.2.
  91  func (d *decoder) processDHT(n int) error {
  92  	for n > 0 {
  93  		if n < 17 {
  94  			return FormatError("DHT has wrong length")
  95  		}
  96  		if err := d.readFull(d.tmp[:17]); err != nil {
  97  			return err
  98  		}
  99  		tc := d.tmp[0] >> 4
 100  		if tc > maxTc {
 101  			return FormatError("bad Tc value")
 102  		}
 103  		th := d.tmp[0] & 0x0f
 104  		// The baseline th <= 1 restriction is specified in table B.5.
 105  		if th > maxTh || (d.baseline && th > 1) {
 106  			return FormatError("bad Th value")
 107  		}
 108  		h := &d.huff[tc][th]
 109  
 110  		// Read nCodes and h.vals (and derive h.nCodes).
 111  		// nCodes[i] is the number of codes with code length i.
 112  		// h.nCodes is the total number of codes.
 113  		h.nCodes = 0
 114  		var nCodes [maxCodeLength]int32
 115  		for i := range nCodes {
 116  			nCodes[i] = int32(d.tmp[i+1])
 117  			h.nCodes += nCodes[i]
 118  		}
 119  		if h.nCodes == 0 {
 120  			return FormatError("Huffman table has zero length")
 121  		}
 122  		if h.nCodes > maxNCodes {
 123  			return FormatError("Huffman table has excessive length")
 124  		}
 125  		n -= int(h.nCodes) + 17
 126  		if n < 0 {
 127  			return FormatError("DHT has wrong length")
 128  		}
 129  		if err := d.readFull(h.vals[:h.nCodes]); err != nil {
 130  			return err
 131  		}
 132  
 133  		// Derive the look-up table.
 134  		clear(h.lut[:])
 135  		var x, code uint32
 136  		for i := uint32(0); i < lutSize; i++ {
 137  			code <<= 1
 138  			for j := int32(0); j < nCodes[i]; j++ {
 139  				// The codeLength is 1+i, so shift code by 8-(1+i) to
 140  				// calculate the high bits for every 8-bit sequence
 141  				// whose codeLength's high bits matches code.
 142  				// The high 8 bits of lutValue are the encoded value.
 143  				// The low 8 bits are 1 plus the codeLength.
 144  				base := uint8(code << (7 - i))
 145  				lutValue := uint16(h.vals[x])<<8 | uint16(2+i)
 146  				for k := uint8(0); k < 1<<(7-i); k++ {
 147  					h.lut[base|k] = lutValue
 148  				}
 149  				code++
 150  				x++
 151  			}
 152  		}
 153  
 154  		// Derive minCodes, maxCodes, and valsIndices.
 155  		var c, index int32
 156  		for i, n := range nCodes {
 157  			if n == 0 {
 158  				h.minCodes[i] = -1
 159  				h.maxCodes[i] = -1
 160  				h.valsIndices[i] = -1
 161  			} else {
 162  				h.minCodes[i] = c
 163  				h.maxCodes[i] = c + n - 1
 164  				h.valsIndices[i] = index
 165  				c += n
 166  				index += n
 167  			}
 168  			c <<= 1
 169  		}
 170  	}
 171  	return nil
 172  }
 173  
 174  // decodeHuffman returns the next Huffman-coded value from the bit-stream,
 175  // decoded according to h.
 176  func (d *decoder) decodeHuffman(h *huffman) (uint8, error) {
 177  	if h.nCodes == 0 {
 178  		return 0, FormatError("uninitialized Huffman table")
 179  	}
 180  
 181  	if d.bits.n < 8 {
 182  		if err := d.ensureNBits(8); err != nil {
 183  			if err != errMissingFF00 && err != errShortHuffmanData {
 184  				return 0, err
 185  			}
 186  			// There are no more bytes of data in this segment, but we may still
 187  			// be able to read the next symbol out of the previously read bits.
 188  			// First, undo the readByte that the ensureNBits call made.
 189  			if d.bytes.nUnreadable != 0 {
 190  				d.unreadByteStuffedByte()
 191  			}
 192  			goto slowPath
 193  		}
 194  	}
 195  	if v := h.lut[(d.bits.a>>uint32(d.bits.n-lutSize))&0xff]; v != 0 {
 196  		n := (v & 0xff) - 1
 197  		d.bits.n -= int32(n)
 198  		d.bits.m >>= n
 199  		return uint8(v >> 8), nil
 200  	}
 201  
 202  slowPath:
 203  	for i, code := 0, int32(0); i < maxCodeLength; i++ {
 204  		if d.bits.n == 0 {
 205  			if err := d.ensureNBits(1); err != nil {
 206  				return 0, err
 207  			}
 208  		}
 209  		if d.bits.a&d.bits.m != 0 {
 210  			code |= 1
 211  		}
 212  		d.bits.n--
 213  		d.bits.m >>= 1
 214  		if code <= h.maxCodes[i] {
 215  			return h.vals[h.valsIndices[i]+code-h.minCodes[i]], nil
 216  		}
 217  		code <<= 1
 218  	}
 219  	return 0, FormatError("bad Huffman code")
 220  }
 221  
 222  func (d *decoder) decodeBit() (bool, error) {
 223  	if d.bits.n == 0 {
 224  		if err := d.ensureNBits(1); err != nil {
 225  			return false, err
 226  		}
 227  	}
 228  	ret := d.bits.a&d.bits.m != 0
 229  	d.bits.n--
 230  	d.bits.m >>= 1
 231  	return ret, nil
 232  }
 233  
 234  func (d *decoder) decodeBits(n int32) (uint32, error) {
 235  	if d.bits.n < n {
 236  		if err := d.ensureNBits(n); err != nil {
 237  			return 0, err
 238  		}
 239  	}
 240  	ret := d.bits.a >> uint32(d.bits.n-n)
 241  	ret &= (1 << uint32(n)) - 1
 242  	d.bits.n -= n
 243  	d.bits.m >>= uint32(n)
 244  	return ret, nil
 245  }
 246