huffman.go

     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
View as plain text