Source file src/compress/flate/huffman_bit_writer.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 flate
     6  
     7  import (
     8  	"io"
     9  	"math"
    10  	"sync"
    11  )
    12  
    13  const (
    14  	// The largest offset code.
    15  	offsetCodeCount = 30
    16  
    17  	// The special code used to mark the end of a block.
    18  	endBlockMarker = 256
    19  
    20  	// The first length code.
    21  	lengthCodesStart = 257
    22  
    23  	// The number of codegen codes.
    24  	codegenCodeCount = 19
    25  	badCode          = 255
    26  
    27  	// maxPredefinedTokens is the maximum number of tokens
    28  	// where we check if fixed size is smaller.
    29  	maxPredefinedTokens = 250
    30  
    31  	// bufferFlushSize indicates the buffer size
    32  	// after which bytes are flushed to the writer.
    33  	// Should preferably be a multiple of 6, since
    34  	// we accumulate 6 bytes between writes to the buffer.
    35  	bufferFlushSize = 246
    36  )
    37  
    38  // lengthExtraBitsMinCode is the minimum length code that emits extra bits.
    39  const lengthExtraBitsMinCode = 8
    40  
    41  // lengthExtraBits[i] is the number of extra bits needed by
    42  // length code i + lengthCodesStart.
    43  var lengthExtraBits = [32]uint8{
    44  	/* 257 */ 0, 0, 0,
    45  	/* 260 */ 0, 0, 0, 0, 0, 1, 1, 1, 1, 2,
    46  	/* 270 */ 2, 2, 2, 3, 3, 3, 3, 4, 4, 4,
    47  	/* 280 */ 4, 5, 5, 5, 5, 0,
    48  }
    49  
    50  // lengthBase[i] is the length indicated by length code i + lengthCodesStart.
    51  var lengthBase = [32]uint8{
    52  	0, 1, 2, 3, 4, 5, 6, 7, 8, 10,
    53  	12, 14, 16, 20, 24, 28, 32, 40, 48, 56,
    54  	64, 80, 96, 112, 128, 160, 192, 224, 255,
    55  }
    56  
    57  // offsetExtraBitsMinCode is the minimum offset code that emits extra bits.
    58  const offsetExtraBitsMinCode = 4
    59  
    60  // offsetExtraBits[i] is the number of extra bits for offset code i.
    61  var offsetExtraBits = [32]int8{
    62  	0, 0, 0, 0, 1, 1, 2, 2, 3, 3,
    63  	4, 4, 5, 5, 6, 6, 7, 7, 8, 8,
    64  	9, 9, 10, 10, 11, 11, 12, 12, 13, 13,
    65  	/* extended window */
    66  	14, 14,
    67  }
    68  
    69  // offsetCombined combines offset lookup of extra bits and offset code in a single table.
    70  var offsetCombined = [32]uint32{
    71  	0x0, 0x0, 0x0, 0x0, 0x401, 0x601, 0x802, 0xc02,
    72  	0x1003, 0x1803, 0x2004, 0x3004, 0x4005, 0x6005,
    73  	0x8006, 0xc006, 0x10007, 0x18007, 0x20008, 0x30008,
    74  	0x40009, 0x60009, 0x8000a, 0xc000a, 0x10000b, 0x18000b,
    75  	0x20000c, 0x30000c, 0x40000d, 0x60000d, 0x0, 0x0}
    76  
    77  /*
    78  Generated with:
    79  
    80  func genOffsetCombined() {
    81  	var offsetBase = [32]uint32{
    82  		0x000000, 0x000001, 0x000002, 0x000003, 0x000004,
    83  		0x000006, 0x000008, 0x00000c, 0x000010, 0x000018,
    84  		0x000020, 0x000030, 0x000040, 0x000060, 0x000080,
    85  		0x0000c0, 0x000100, 0x000180, 0x000200, 0x000300,
    86  		0x000400, 0x000600, 0x000800, 0x000c00, 0x001000,
    87  		0x001800, 0x002000, 0x003000, 0x004000, 0x006000,
    88  
    89  		0x008000, 0x00c000,
    90  	}
    91  
    92  	for i := range offsetCombined[:] {
    93  		// Don't use extended window values...
    94  		if offsetExtraBits[i] == 0 || offsetBase[i] > 0x006000 {
    95  			continue
    96  		}
    97  		offsetCombined[i] = uint32(offsetExtraBits[i]) | (offsetBase[i] << 8)
    98  	}
    99  	fmt.Printf("offsetCombined = %#v\n", offsetCombined)
   100  }
   101  */
   102  
   103  // codegenOrder is the order in which codegen code sizes are written.
   104  var codegenOrder = []uint32{16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15}
   105  
   106  // huffmanBitWriter encodes tokens and values to a stream.
   107  // The huffmanBitWriter supports reusing huffman tables and will combine
   108  // blocks, if compression is less than creating a new table.
   109  //
   110  // An incoming block estimates the output size of a new table using a
   111  // 'fresh' by calculating the optimal size and adding a penalty.
   112  // A Huffman table is not optimal, which is why we add a penalty,
   113  // and generating a new table is slower for both compression and decompression.
   114  type huffmanBitWriter struct {
   115  	// writer is the underlying writer.
   116  	// Do not use it directly; use the write method, which ensures
   117  	// that Write errors are sticky.
   118  	writer io.Writer
   119  
   120  	// Data waiting to be written is bytes[0:nbytes]
   121  	// and then the low nbits of bits.
   122  	bits   uint64
   123  	nbits  uint8
   124  	nbytes uint8
   125  
   126  	// If wroteHuffman is set, a table for outputting only literals
   127  	// has been generated and offsets are invalid.
   128  	wroteHuffman    bool
   129  	literalEncoding *huffmanEncoder
   130  	tmpLitEncoding  *huffmanEncoder
   131  	offsetEncoding  *huffmanEncoder
   132  	codegenEncoding *huffmanEncoder
   133  	err             error
   134  
   135  	// If prevHeader is non-zero the Huffman table can be reused.
   136  	// It also indicates that an EOB has not yet been emitted, so if a new table
   137  	// is generated, an EOB with the previous table must be written.
   138  	prevHeader int
   139  
   140  	// logNewTablePenalty is a log2 penalty reduction for creating new tables.
   141  	// The initial penalty is 100%.
   142  	// Adding 1 will cut the penalty in half.
   143  	logNewTablePenalty uint
   144  	bytes              [256 + 8]byte
   145  	literalFreq        [lengthCodesStart + 32]uint16
   146  	offsetFreq         [32]uint16
   147  	codegenFreq        [codegenCodeCount]uint16
   148  
   149  	// codegen must have an extra space for the final symbol.
   150  	codegen [literalCount + offsetCodeCount + 1]uint8
   151  }
   152  
   153  // newHuffmanBitWriter creates a new huffmanBitWriter that will write to w.
   154  func newHuffmanBitWriter(w io.Writer) *huffmanBitWriter {
   155  	return &huffmanBitWriter{
   156  		writer:          w,
   157  		literalEncoding: newHuffmanEncoder(literalCount),
   158  		tmpLitEncoding:  newHuffmanEncoder(literalCount),
   159  		codegenEncoding: newHuffmanEncoder(codegenCodeCount),
   160  		offsetEncoding:  newHuffmanEncoder(offsetCodeCount),
   161  	}
   162  }
   163  
   164  // reset the huffmanBitWriter state and replace the output.
   165  func (w *huffmanBitWriter) reset(writer io.Writer) {
   166  	w.writer = writer
   167  	w.bits, w.nbits, w.nbytes, w.err = 0, 0, 0, nil
   168  	w.prevHeader = 0
   169  	w.wroteHuffman = false
   170  }
   171  
   172  // canReuse checks if the current generated tables can be
   173  // reused for the provided tokens.
   174  func (w *huffmanBitWriter) canReuse(t *tokens) (ok bool) {
   175  	a := t.offHist[:offsetCodeCount]
   176  	b := w.offsetEncoding.codes
   177  	b = b[:len(a)]
   178  	for i, v := range a {
   179  		if v != 0 && b[i].zero() {
   180  			return false
   181  		}
   182  	}
   183  
   184  	a = t.extraHist[:literalCount-256]
   185  	b = w.literalEncoding.codes[256:literalCount]
   186  	b = b[:len(a)]
   187  	for i, v := range a {
   188  		if v != 0 && b[i].zero() {
   189  			return false
   190  		}
   191  	}
   192  
   193  	a = t.litHist[:256]
   194  	b = w.literalEncoding.codes[:len(a)]
   195  	for i, v := range a {
   196  		if v != 0 && b[i].zero() {
   197  			return false
   198  		}
   199  	}
   200  	return true
   201  }
   202  
   203  // flush flushes the currently encoded data.
   204  // An EOB will be written if the current block hasn't been ended.
   205  func (w *huffmanBitWriter) flush() {
   206  	if w.err != nil {
   207  		w.nbits = 0
   208  		return
   209  	}
   210  	if w.prevHeader > 0 {
   211  		// We owe an EOB
   212  		w.writeCode(w.literalEncoding.codes[endBlockMarker])
   213  		w.prevHeader = 0
   214  	}
   215  	n := w.nbytes
   216  	for w.nbits != 0 {
   217  		w.bytes[n] = byte(w.bits)
   218  		w.bits >>= 8
   219  		if w.nbits > 8 { // Avoid underflow
   220  			w.nbits -= 8
   221  		} else {
   222  			w.nbits = 0
   223  		}
   224  		n++
   225  	}
   226  	w.bits = 0
   227  	if n > 0 {
   228  		w.write(w.bytes[:n])
   229  	}
   230  	w.nbytes = 0
   231  }
   232  
   233  // write writes the provided bytes directly to the output,
   234  // ignoring all queued bytes.
   235  func (w *huffmanBitWriter) write(b []byte) {
   236  	if w.err != nil {
   237  		return
   238  	}
   239  	_, w.err = w.writer.Write(b)
   240  }
   241  
   242  // writeBits writes nb bits from b to the stream.
   243  func (w *huffmanBitWriter) writeBits(b int32, nb uint8) {
   244  	w.bits |= uint64(b) << (w.nbits & 63)
   245  	w.nbits += nb
   246  	if w.nbits >= 48 {
   247  		w.flushBits()
   248  	}
   249  }
   250  
   251  // writeBytes writes the provided bytes to the stream.
   252  func (w *huffmanBitWriter) writeBytes(bytes []byte) {
   253  	if w.err != nil {
   254  		return
   255  	}
   256  	n := w.nbytes
   257  	if w.nbits&7 != 0 {
   258  		w.err = InternalError("writeBytes with unfinished bits")
   259  		return
   260  	}
   261  	for w.nbits != 0 {
   262  		w.bytes[n] = byte(w.bits)
   263  		w.bits >>= 8
   264  		w.nbits -= 8
   265  		n++
   266  	}
   267  	if n != 0 {
   268  		w.write(w.bytes[:n])
   269  	}
   270  	w.nbytes = 0
   271  	w.write(bytes)
   272  }
   273  
   274  // RFC 1951 3.2.7 specifies a special run-length encoding for specifying
   275  // the literal and offset lengths arrays (which are concatenated into a single
   276  // array).  This method generates that run-length encoding.
   277  //
   278  // The result is written into the codegen array, and the frequencies
   279  // of each code is written into the codegenFreq array.
   280  // Codes 0-15 are single byte codes. Codes 16-18 are followed by additional
   281  // information. Code badCode is an end marker
   282  //
   283  //	numLiterals      The number of literals in literalEncoding
   284  //	numOffsets       The number of offsets in offsetEncoding
   285  //	litenc, offenc   The literal and offset encoder to use
   286  func (w *huffmanBitWriter) generateCodegen(numLiterals int, numOffsets int, litEnc, offEnc *huffmanEncoder) {
   287  	clear(w.codegenFreq[:])
   288  	// Note that we are using codegen both as a temporary variable for holding
   289  	// a copy of the frequencies, and as the place where we put the result.
   290  	// This is fine because the output is always shorter than the input used
   291  	// so far.
   292  	codegen := w.codegen[:] // cache
   293  	// Copy the concatenated code sizes to codegen. Put a marker at the end.
   294  	cgnl := codegen[:numLiterals]
   295  	for i := range cgnl {
   296  		cgnl[i] = litEnc.codes[i].len()
   297  	}
   298  
   299  	cgnl = codegen[numLiterals : numLiterals+numOffsets]
   300  	for i := range cgnl {
   301  		cgnl[i] = offEnc.codes[i].len()
   302  	}
   303  	codegen[numLiterals+numOffsets] = badCode
   304  
   305  	size := codegen[0]
   306  	count := 1
   307  	outIndex := 0
   308  	for inIndex := 1; size != badCode; inIndex++ {
   309  		// INVARIANT: We have seen "count" copies of size that have not yet
   310  		// had output generated for them.
   311  		nextSize := codegen[inIndex]
   312  		if nextSize == size {
   313  			count++
   314  			continue
   315  		}
   316  		// We need to generate codegen indicating "count" of size.
   317  		if size != 0 {
   318  			codegen[outIndex] = size
   319  			outIndex++
   320  			w.codegenFreq[size]++
   321  			count--
   322  			for count >= 3 {
   323  				n := min(6, count)
   324  				codegen[outIndex] = 16
   325  				outIndex++
   326  				codegen[outIndex] = uint8(n - 3)
   327  				outIndex++
   328  				w.codegenFreq[16]++
   329  				count -= n
   330  			}
   331  		} else {
   332  			for count >= 11 {
   333  				n := min(138, count)
   334  				codegen[outIndex] = 18
   335  				outIndex++
   336  				codegen[outIndex] = uint8(n - 11)
   337  				outIndex++
   338  				w.codegenFreq[18]++
   339  				count -= n
   340  			}
   341  			if count >= 3 {
   342  				// count >= 3 && count <= 10
   343  				codegen[outIndex] = 17
   344  				outIndex++
   345  				codegen[outIndex] = uint8(count - 3)
   346  				outIndex++
   347  				w.codegenFreq[17]++
   348  				count = 0
   349  			}
   350  		}
   351  		count--
   352  		for ; count >= 0; count-- {
   353  			codegen[outIndex] = size
   354  			outIndex++
   355  			w.codegenFreq[size]++
   356  		}
   357  		// Set up invariant for next time through the loop.
   358  		size = nextSize
   359  		count = 1
   360  	}
   361  	// Marker indicating the end of the codegen.
   362  	codegen[outIndex] = badCode
   363  }
   364  
   365  // codegens returns current number of non-zero codegens.
   366  func (w *huffmanBitWriter) codegens() int {
   367  	numCodegens := len(w.codegenFreq)
   368  	for numCodegens > 4 && w.codegenFreq[codegenOrder[numCodegens-1]] == 0 {
   369  		numCodegens--
   370  	}
   371  	return numCodegens
   372  }
   373  
   374  // headerSize returns the size of the header with the current encodings.
   375  func (w *huffmanBitWriter) headerSize() (size, numCodegens int) {
   376  	numCodegens = len(w.codegenFreq)
   377  	for numCodegens > 4 && w.codegenFreq[codegenOrder[numCodegens-1]] == 0 {
   378  		numCodegens--
   379  	}
   380  	return 3 + 5 + 5 + 4 + (3 * numCodegens) +
   381  		w.codegenEncoding.bitLength(w.codegenFreq[:]) +
   382  		int(w.codegenFreq[16])*2 +
   383  		int(w.codegenFreq[17])*3 +
   384  		int(w.codegenFreq[18])*7, numCodegens
   385  }
   386  
   387  // dynamicSize returns the size of dynamically encoded data in bits.
   388  func (w *huffmanBitWriter) dynamicReuseSize(litEnc, offEnc *huffmanEncoder) (size int) {
   389  	size = litEnc.bitLength(w.literalFreq[:]) +
   390  		offEnc.bitLength(w.offsetFreq[:])
   391  	return size
   392  }
   393  
   394  // dynamicSize returns the size of dynamically encoded data in bits.
   395  func (w *huffmanBitWriter) dynamicSize(litEnc, offEnc *huffmanEncoder, extraBits int) (size, numCodegens int) {
   396  	header, numCodegens := w.headerSize()
   397  	size = header +
   398  		litEnc.bitLength(w.literalFreq[:]) +
   399  		offEnc.bitLength(w.offsetFreq[:]) +
   400  		extraBits
   401  	return size, numCodegens
   402  }
   403  
   404  // extraBitSize returns the number of bits that will be written
   405  // as "extra" bits on matches.
   406  func (w *huffmanBitWriter) extraBitSize() int {
   407  	total := 0
   408  	for i, n := range w.literalFreq[257:literalCount] {
   409  		total += int(n) * int(lengthExtraBits[i&31])
   410  	}
   411  	for i, n := range w.offsetFreq[:offsetCodeCount] {
   412  		total += int(n) * int(offsetExtraBits[i&31])
   413  	}
   414  	return total
   415  }
   416  
   417  // fixedSize returns the size of dynamically encoded data in bits.
   418  func (w *huffmanBitWriter) fixedSize(extraBits int) int {
   419  	return 3 +
   420  		fixedLiteralEncoding().bitLength(w.literalFreq[:]) +
   421  		fixedOffsetEncoding().bitLength(w.offsetFreq[:]) +
   422  		extraBits
   423  }
   424  
   425  // storedSize calculates the stored size, including header.
   426  // The function returns the size in bits and whether the block
   427  // fits inside a single block.
   428  func (w *huffmanBitWriter) storedSize(in []byte) (int, bool) {
   429  	if in == nil {
   430  		return 0, false
   431  	}
   432  	if len(in) <= maxStoreBlockSize {
   433  		return (len(in) + 5) * 8, true
   434  	}
   435  	return 0, false
   436  }
   437  
   438  // writeCode writes 'c' to the stream.
   439  func (w *huffmanBitWriter) writeCode(c hcode) {
   440  	w.bits |= c.code64() << (w.nbits & reg8SizeMask64)
   441  	w.nbits += c.len()
   442  	if w.nbits >= 48 {
   443  		w.flushBits()
   444  	}
   445  }
   446  
   447  // flushBits writes accumulated bits to the byte buffer.
   448  func (w *huffmanBitWriter) flushBits() {
   449  	bits := w.bits
   450  	w.bits >>= 48
   451  	w.nbits -= 48
   452  	n := w.nbytes
   453  
   454  	// We overwrite, but faster...
   455  	storeLE64(w.bytes[n:], bits)
   456  	n += 6
   457  
   458  	if n >= bufferFlushSize {
   459  		if w.err != nil {
   460  			n = 0
   461  			return
   462  		}
   463  		w.write(w.bytes[:n])
   464  		n = 0
   465  	}
   466  
   467  	w.nbytes = n
   468  }
   469  
   470  // writeDynamicHeader writes the header of a dynamic Huffman block to the output stream.
   471  //
   472  // numLiterals is the number of literals specified in codegen.
   473  // numOffsets is the number of offsets specified in codegen.
   474  // numCodegens is the number of codegens used in codegen.
   475  func (w *huffmanBitWriter) writeDynamicHeader(numLiterals int, numOffsets int, numCodegens int, isEof bool) {
   476  	if w.err != nil {
   477  		return
   478  	}
   479  	var firstBits int32 = 4
   480  	if isEof {
   481  		firstBits = 5
   482  	}
   483  	w.writeBits(firstBits, 3)
   484  	w.writeBits(int32(numLiterals-257), 5)
   485  	w.writeBits(int32(numOffsets-1), 5)
   486  	w.writeBits(int32(numCodegens-4), 4)
   487  
   488  	for i := range numCodegens {
   489  		value := uint(w.codegenEncoding.codes[codegenOrder[i]].len())
   490  		w.writeBits(int32(value), 3)
   491  	}
   492  
   493  	i := 0
   494  	for {
   495  		var codeWord = uint32(w.codegen[i])
   496  		i++
   497  		if codeWord == badCode {
   498  			break
   499  		}
   500  		w.writeCode(w.codegenEncoding.codes[codeWord])
   501  
   502  		switch codeWord {
   503  		case 16:
   504  			w.writeBits(int32(w.codegen[i]), 2)
   505  			i++
   506  		case 17:
   507  			w.writeBits(int32(w.codegen[i]), 3)
   508  			i++
   509  		case 18:
   510  			w.writeBits(int32(w.codegen[i]), 7)
   511  			i++
   512  		}
   513  	}
   514  }
   515  
   516  // writeStoredHeader writes a stored header.
   517  // If the stored block is only used for EOF,
   518  // it is replaced with a fixed huffman block.
   519  func (w *huffmanBitWriter) writeStoredHeader(length int, isEof bool) {
   520  	if w.err != nil {
   521  		return
   522  	}
   523  	if w.prevHeader > 0 {
   524  		// We owe an EOB
   525  		w.writeCode(w.literalEncoding.codes[endBlockMarker])
   526  		w.prevHeader = 0
   527  	}
   528  
   529  	// To write EOF, use a fixed encoding block. 10 bits instead of 5 bytes.
   530  	if length == 0 && isEof {
   531  		w.writeFixedHeader(isEof)
   532  		// EOB: 7 bits, value: 0
   533  		w.writeBits(0, 7)
   534  		w.flush()
   535  		return
   536  	}
   537  
   538  	var flag int32
   539  	if isEof {
   540  		flag = 1
   541  	}
   542  	w.writeBits(flag, 3)
   543  	w.flush()
   544  	w.writeBits(int32(length), 16)
   545  	w.writeBits(int32(^uint16(length)), 16)
   546  }
   547  
   548  // writeFixedHeader writes a fixed encoding header to the output stream.
   549  func (w *huffmanBitWriter) writeFixedHeader(isEof bool) {
   550  	if w.err != nil {
   551  		return
   552  	}
   553  	if w.prevHeader > 0 {
   554  		// We owe an EOB
   555  		w.writeCode(w.literalEncoding.codes[endBlockMarker])
   556  		w.prevHeader = 0
   557  	}
   558  
   559  	// Indicate that we are a fixed Huffman block
   560  	var value int32 = 2
   561  	if isEof {
   562  		value = 3
   563  	}
   564  	w.writeBits(value, 3)
   565  }
   566  
   567  // writeBlock writes a block of tokens using the smallest encoding.
   568  // The original input can be supplied, and if the Huffman-encoded data
   569  // is larger than the original bytes, the data will be written as a
   570  // stored block.
   571  // If the input is nil, the tokens will always be Huffman encoded.
   572  func (w *huffmanBitWriter) writeBlock(tokens *tokens, eof bool, input []byte) {
   573  	if w.err != nil {
   574  		return
   575  	}
   576  
   577  	tokens.AddEOB()
   578  	if w.prevHeader > 0 {
   579  		// We owe an EOB
   580  		w.writeCode(w.literalEncoding.codes[endBlockMarker])
   581  		w.prevHeader = 0
   582  	}
   583  	numLiterals, numOffsets := w.indexTokens(tokens)
   584  	w.generate()
   585  	var extraBits int
   586  	storedSize, storable := w.storedSize(input)
   587  	if storable {
   588  		extraBits = w.extraBitSize()
   589  	}
   590  
   591  	// Figure out smallest code.
   592  	// Fixed Huffman baseline.
   593  	var literalEncoding = fixedLiteralEncoding()
   594  	var offsetEncoding = fixedOffsetEncoding()
   595  	var size = math.MaxInt32
   596  	if tokens.n < maxPredefinedTokens {
   597  		size = w.fixedSize(extraBits)
   598  	}
   599  
   600  	// Dynamic Huffman?
   601  	var numCodegens int
   602  
   603  	// Generate codegen and codegenFrequencies, which indicates how to encode
   604  	// the literalEncoding and the offsetEncoding.
   605  	w.generateCodegen(numLiterals, numOffsets, w.literalEncoding, w.offsetEncoding)
   606  	w.codegenEncoding.generate(w.codegenFreq[:], 7)
   607  	dynamicSize, numCodegens := w.dynamicSize(w.literalEncoding, w.offsetEncoding, extraBits)
   608  
   609  	if dynamicSize < size {
   610  		size = dynamicSize
   611  		literalEncoding = w.literalEncoding
   612  		offsetEncoding = w.offsetEncoding
   613  	}
   614  
   615  	// Stored bytes?
   616  	if storable && storedSize <= size {
   617  		w.writeStoredHeader(len(input), eof)
   618  		w.writeBytes(input)
   619  		return
   620  	}
   621  
   622  	// Huffman.
   623  	if literalEncoding == fixedLiteralEncoding() {
   624  		w.writeFixedHeader(eof)
   625  	} else {
   626  		w.writeDynamicHeader(numLiterals, numOffsets, numCodegens, eof)
   627  	}
   628  
   629  	// Write the tokens.
   630  	w.writeTokens(tokens.Slice(), literalEncoding.codes, offsetEncoding.codes)
   631  }
   632  
   633  // writeBlockDynamic encodes a block using a dynamic Huffman table.
   634  // This should be used if the symbols used have a disproportionate
   635  // histogram distribution.
   636  func (w *huffmanBitWriter) writeBlockDynamic(tokens *tokens, eof bool, input []byte, sync bool) {
   637  	if w.err != nil {
   638  		return
   639  	}
   640  
   641  	sync = sync || eof
   642  	if sync {
   643  		tokens.AddEOB()
   644  	} else {
   645  		// Ensure we can always write EOB.
   646  		tokens.extraHist[0] = 1
   647  	}
   648  
   649  	// We cannot reuse pure Huffman table, and must mark as EOF.
   650  	if (w.wroteHuffman || eof) && w.prevHeader > 0 {
   651  		// We will not try to reuse.
   652  		w.writeCode(w.literalEncoding.codes[endBlockMarker])
   653  		w.prevHeader = 0
   654  		w.wroteHuffman = false
   655  	}
   656  
   657  	if w.prevHeader > 0 && !w.canReuse(tokens) {
   658  		w.writeCode(w.literalEncoding.codes[endBlockMarker])
   659  		w.prevHeader = 0
   660  	}
   661  
   662  	numLiterals, numOffsets := w.indexTokens(tokens)
   663  	extraBits := 0
   664  	ssize, storable := w.storedSize(input)
   665  
   666  	if storable || w.prevHeader > 0 {
   667  		extraBits = w.extraBitSize()
   668  	}
   669  
   670  	var size int
   671  
   672  	// Check whether we should reuse the previous Huffman table.
   673  	if w.prevHeader > 0 {
   674  		// Estimate size for using a new table.
   675  		// Use the previous header size as the best estimate.
   676  		newSize := w.prevHeader + tokens.EstimatedBits()
   677  
   678  		// The estimated size is calculated as an optimal table.
   679  		// We add a penalty to make it more realistic and re-use a bit more.
   680  		newSize += int(w.literalEncoding.codes[endBlockMarker].len()) + newSize>>w.logNewTablePenalty
   681  
   682  		// Calculate the size for reusing the current table.
   683  		reuseSize := w.dynamicReuseSize(w.literalEncoding, w.offsetEncoding) + extraBits
   684  
   685  		// Check if a new table is better.
   686  		if newSize < reuseSize {
   687  			// Write the EOB we owe.
   688  			w.writeCode(w.literalEncoding.codes[endBlockMarker])
   689  			size = newSize
   690  			w.prevHeader = 0
   691  		} else {
   692  			size = reuseSize
   693  		}
   694  
   695  		// Small blocks can be more efficient with fixed encoding.
   696  		if tokens.n < maxPredefinedTokens {
   697  			if preSize := w.fixedSize(extraBits) + 7; preSize < size {
   698  				// Check if we get a reasonable size decrease.
   699  				if storable && ssize <= size {
   700  					w.writeStoredHeader(len(input), eof)
   701  					w.writeBytes(input)
   702  					return
   703  				}
   704  				w.writeFixedHeader(eof)
   705  				if !sync {
   706  					tokens.AddEOB()
   707  				}
   708  				w.writeTokens(tokens.Slice(), fixedLiteralEncoding().codes, fixedOffsetEncoding().codes)
   709  				return
   710  			}
   711  		}
   712  
   713  		// Check if we get a reasonable size decrease.
   714  		if storable && ssize <= size {
   715  			w.writeStoredHeader(len(input), eof)
   716  			w.writeBytes(input)
   717  			return
   718  		}
   719  	}
   720  
   721  	// We want a new block/table
   722  	if w.prevHeader == 0 {
   723  		w.literalFreq[endBlockMarker] = 1
   724  
   725  		w.generate()
   726  		// Generate codegen and codegenFrequencies, which indicates how to encode
   727  		// the literalEncoding and the offsetEncoding.
   728  		w.generateCodegen(numLiterals, numOffsets, w.literalEncoding, w.offsetEncoding)
   729  		w.codegenEncoding.generate(w.codegenFreq[:], 7)
   730  
   731  		var numCodegens int
   732  		size, numCodegens = w.dynamicSize(w.literalEncoding, w.offsetEncoding, extraBits)
   733  
   734  		// Store predefined or raw, if we don't get a reasonable improvement.
   735  		if tokens.n < maxPredefinedTokens {
   736  			if preSize := w.fixedSize(extraBits); preSize <= size {
   737  				// Store bytes, if we don't get an improvement.
   738  				if storable && ssize <= preSize {
   739  					w.writeStoredHeader(len(input), eof)
   740  					w.writeBytes(input)
   741  					return
   742  				}
   743  				w.writeFixedHeader(eof)
   744  				if !sync {
   745  					tokens.AddEOB()
   746  				}
   747  				w.writeTokens(tokens.Slice(), fixedLiteralEncoding().codes, fixedOffsetEncoding().codes)
   748  				return
   749  			}
   750  		}
   751  
   752  		if storable && ssize <= size {
   753  			// Store bytes, if we don't get an improvement.
   754  			w.writeStoredHeader(len(input), eof)
   755  			w.writeBytes(input)
   756  			return
   757  		}
   758  
   759  		// Write Huffman table.
   760  		w.writeDynamicHeader(numLiterals, numOffsets, numCodegens, eof)
   761  		if !sync {
   762  			w.prevHeader, _ = w.headerSize()
   763  		}
   764  		w.wroteHuffman = false
   765  	}
   766  
   767  	if sync {
   768  		w.prevHeader = 0
   769  	}
   770  	// Write the tokens.
   771  	w.writeTokens(tokens.Slice(), w.literalEncoding.codes, w.offsetEncoding.codes)
   772  }
   773  
   774  // indexTokens indexes a slice of tokens, updates literalFreq and offsetFreq,
   775  // and generates literalEncoding and offsetEncoding.
   776  // It returns the number of literal and offset tokens.
   777  func (w *huffmanBitWriter) indexTokens(t *tokens) (numLiterals, numOffsets int) {
   778  	*(*[256]uint16)(w.literalFreq[:]) = t.litHist
   779  	*(*[32]uint16)(w.literalFreq[256:]) = t.extraHist
   780  	w.offsetFreq = t.offHist
   781  
   782  	if t.n == 0 {
   783  		return
   784  	}
   785  	// get the number of literals
   786  	numLiterals = len(w.literalFreq)
   787  	for w.literalFreq[numLiterals-1] == 0 {
   788  		numLiterals--
   789  	}
   790  	// get the number of offsets
   791  	numOffsets = len(w.offsetFreq)
   792  	for numOffsets > 0 && w.offsetFreq[numOffsets-1] == 0 {
   793  		numOffsets--
   794  	}
   795  	if numOffsets == 0 {
   796  		// We haven't found a single match. If we want to go with the dynamic encoding,
   797  		// we should count at least one offset to be sure that the offset huffman tree could be encoded.
   798  		w.offsetFreq[0] = 1
   799  		numOffsets = 1
   800  	}
   801  	return
   802  }
   803  
   804  // generate literalEncoding and offsetEncoding based on respective histograms.
   805  func (w *huffmanBitWriter) generate() {
   806  	w.literalEncoding.generate(w.literalFreq[:literalCount], 15)
   807  	w.offsetEncoding.generate(w.offsetFreq[:offsetCodeCount], 15)
   808  }
   809  
   810  // writeTokens writes a slice of tokens to the output.
   811  // Codes for literal and offset encoding must be supplied.
   812  func (w *huffmanBitWriter) writeTokens(tokens []token, lenCodes, offCodes []hcode) {
   813  	if w.err != nil {
   814  		return
   815  	}
   816  	if len(tokens) == 0 {
   817  		return
   818  	}
   819  
   820  	// Only last token should be endBlockMarker.
   821  	var deferEOB bool
   822  	if tokens[len(tokens)-1] == endBlockMarker {
   823  		tokens = tokens[:len(tokens)-1]
   824  		deferEOB = true
   825  	}
   826  
   827  	// Create slices up to the next power of two to avoid bounds checks.
   828  	lits := lenCodes[:256]
   829  	offs := offCodes[:32]
   830  	lengths := lenCodes[lengthCodesStart:]
   831  	lengths = lengths[:32]
   832  
   833  	// Go 1.16 LOVES having these on stack.
   834  	bits, nbits, nbytes := w.bits, w.nbits, w.nbytes
   835  
   836  	for _, t := range tokens {
   837  		if t < 256 {
   838  			c := lits[t]
   839  			bits |= c.code64() << (nbits & 63)
   840  			nbits += c.len()
   841  			if nbits >= 48 {
   842  				storeLE64(w.bytes[nbytes:], bits)
   843  				bits >>= 48
   844  				nbits -= 48
   845  				nbytes += 6
   846  				if nbytes >= bufferFlushSize {
   847  					if w.err != nil {
   848  						nbytes = 0
   849  						return
   850  					}
   851  					_, w.err = w.writer.Write(w.bytes[:nbytes])
   852  					nbytes = 0
   853  				}
   854  			}
   855  			continue
   856  		}
   857  
   858  		// Write the length
   859  		length := t.length()
   860  		lenCode := lengthCode(length) & 31
   861  		// inlined 'w.writeCode(lengths[lengthCode])'
   862  		c := lengths[lenCode]
   863  		bits |= c.code64() << (nbits & 63)
   864  		nbits += c.len()
   865  		if nbits >= 48 {
   866  			storeLE64(w.bytes[nbytes:], bits)
   867  			bits >>= 48
   868  			nbits -= 48
   869  			nbytes += 6
   870  			if nbytes >= bufferFlushSize {
   871  				if w.err != nil {
   872  					nbytes = 0
   873  					return
   874  				}
   875  				_, w.err = w.writer.Write(w.bytes[:nbytes])
   876  				nbytes = 0
   877  			}
   878  		}
   879  
   880  		if lenCode >= lengthExtraBitsMinCode {
   881  			extraLengthBits := lengthExtraBits[lenCode]
   882  			//w.writeBits(extraLength, extraLengthBits)
   883  			extraLength := int32(length - lengthBase[lenCode])
   884  			bits |= uint64(extraLength) << (nbits & 63)
   885  			nbits += extraLengthBits
   886  			if nbits >= 48 {
   887  				storeLE64(w.bytes[nbytes:], bits)
   888  				bits >>= 48
   889  				nbits -= 48
   890  				nbytes += 6
   891  				if nbytes >= bufferFlushSize {
   892  					if w.err != nil {
   893  						nbytes = 0
   894  						return
   895  					}
   896  					_, w.err = w.writer.Write(w.bytes[:nbytes])
   897  					nbytes = 0
   898  				}
   899  			}
   900  		}
   901  		// Write the offset
   902  		offset := t.offset()
   903  		offCode := (offset >> 16) & 31
   904  		// inlined 'w.writeCode(offs[offCode])'
   905  		c = offs[offCode]
   906  		bits |= c.code64() << (nbits & 63)
   907  		nbits += c.len()
   908  		if nbits >= 48 {
   909  			storeLE64(w.bytes[nbytes:], bits)
   910  			bits >>= 48
   911  			nbits -= 48
   912  			nbytes += 6
   913  			if nbytes >= bufferFlushSize {
   914  				if w.err != nil {
   915  					nbytes = 0
   916  					return
   917  				}
   918  				_, w.err = w.writer.Write(w.bytes[:nbytes])
   919  				nbytes = 0
   920  			}
   921  		}
   922  
   923  		if offCode >= offsetExtraBitsMinCode {
   924  			offsetComb := offsetCombined[offCode]
   925  			bits |= uint64((offset-(offsetComb>>8))&matchOffsetOnlyMask) << (nbits & 63)
   926  			nbits += uint8(offsetComb)
   927  			if nbits >= 48 {
   928  				storeLE64(w.bytes[nbytes:], bits)
   929  				bits >>= 48
   930  				nbits -= 48
   931  				nbytes += 6
   932  				if nbytes >= bufferFlushSize {
   933  					if w.err != nil {
   934  						nbytes = 0
   935  						return
   936  					}
   937  					_, w.err = w.writer.Write(w.bytes[:nbytes])
   938  					nbytes = 0
   939  				}
   940  			}
   941  		}
   942  	}
   943  	// Restore...
   944  	w.bits, w.nbits, w.nbytes = bits, nbits, nbytes
   945  
   946  	if deferEOB {
   947  		w.writeCode(lenCodes[endBlockMarker])
   948  	}
   949  }
   950  
   951  // huffOffset is a static offset encoder used for Huffman-only encoding.
   952  // It can be reused since we will not be encoding offset values.
   953  var huffOffset = sync.OnceValue(func() *huffmanEncoder {
   954  	w := newHuffmanBitWriter(nil)
   955  	w.offsetFreq[0] = 1
   956  	h := newHuffmanEncoder(offsetCodeCount)
   957  	h.generate(w.offsetFreq[:offsetCodeCount], 15)
   958  	return h
   959  })
   960  
   961  // writeBlockHuff encodes a block of bytes as either
   962  // Huffman-encoded literals or uncompressed bytes if the
   963  // results gain very little from compression.
   964  func (w *huffmanBitWriter) writeBlockHuff(eof bool, input []byte, sync bool) {
   965  	if w.err != nil {
   966  		return
   967  	}
   968  
   969  	// Clear histogram
   970  	clear(w.literalFreq[:])
   971  	if !w.wroteHuffman {
   972  		clear(w.offsetFreq[:])
   973  	}
   974  
   975  	const numLiterals = endBlockMarker + 1
   976  	const numOffsets = 1
   977  
   978  	// Estimate size of literal encoding.
   979  	const guessHeaderSizeBits = 70 * 8 // 70 bytes; see https://stackoverflow.com/a/25454430
   980  	histogram(input, w.literalFreq[:numLiterals])
   981  	ssize, storable := w.storedSize(input)
   982  	if storable && len(input) > 1024 {
   983  		// Quick check for incompressible content.
   984  		// The following checks if all frequencies lie
   985  		// close to the average frequency.
   986  		// If so, we quickly store the data uncompressed.
   987  		// This will typically only trigger on random data.
   988  		// Most other data will typically exit after only a few iterations.
   989  		abs := float64(0)
   990  		avg := float64(len(input)) / 256
   991  		max := float64(len(input) * 2)
   992  		for _, v := range w.literalFreq[:256] {
   993  			diff := float64(v) - avg
   994  			abs += diff * diff
   995  			if abs >= max {
   996  				break
   997  			}
   998  		}
   999  		if abs < max {
  1000  			// No chance we can compress this...
  1001  			w.writeStoredHeader(len(input), eof)
  1002  			w.writeBytes(input)
  1003  			return
  1004  		}
  1005  	}
  1006  	w.literalFreq[endBlockMarker] = 1
  1007  	w.tmpLitEncoding.generate(w.literalFreq[:numLiterals], 15)
  1008  	estBits := w.tmpLitEncoding.canEncodeLen(w.literalFreq[:numLiterals])
  1009  	if estBits < math.MaxInt32 {
  1010  		estBits += w.prevHeader
  1011  		if w.prevHeader == 0 {
  1012  			estBits += guessHeaderSizeBits
  1013  		}
  1014  		estBits += estBits >> w.logNewTablePenalty
  1015  	}
  1016  
  1017  	// Store bytes, if we don't get a reasonable improvement.
  1018  	if storable && ssize <= estBits {
  1019  		w.writeStoredHeader(len(input), eof)
  1020  		w.writeBytes(input)
  1021  		return
  1022  	}
  1023  
  1024  	if w.prevHeader > 0 {
  1025  		reuseSize := w.literalEncoding.canEncodeLen(w.literalFreq[:256])
  1026  		if estBits < reuseSize {
  1027  			// We owe an EOB
  1028  			w.writeCode(w.literalEncoding.codes[endBlockMarker])
  1029  			w.prevHeader = 0
  1030  		}
  1031  	}
  1032  
  1033  	if w.prevHeader == 0 {
  1034  		// Use the temp encoding, so swap.
  1035  		w.literalEncoding, w.tmpLitEncoding = w.tmpLitEncoding, w.literalEncoding
  1036  		// Generate codegen and codegenFrequencies, which indicates how to encode
  1037  		// the literalEncoding and the offsetEncoding.
  1038  		w.generateCodegen(numLiterals, numOffsets, w.literalEncoding, huffOffset())
  1039  		w.codegenEncoding.generate(w.codegenFreq[:], 7)
  1040  		numCodegens := w.codegens()
  1041  
  1042  		// Huffman.
  1043  		w.writeDynamicHeader(numLiterals, numOffsets, numCodegens, eof)
  1044  		w.wroteHuffman = true
  1045  		w.prevHeader, _ = w.headerSize()
  1046  	}
  1047  
  1048  	encoding := w.literalEncoding.codes[:256]
  1049  	// Go 1.16 LOVES having these on stack. At least 1.5x the speed.
  1050  	bits, nbits, nbytes := w.bits, w.nbits, w.nbytes
  1051  
  1052  	// Unroll, write 3 codes/loop.
  1053  	// Fastest number of unrolls.
  1054  	for len(input) > 3 {
  1055  		// We must have at least 48 bits free.
  1056  		if nbits >= 8 {
  1057  			n := nbits >> 3
  1058  			storeLE64(w.bytes[nbytes:], bits)
  1059  			bits >>= (n * 8) & 63
  1060  			nbits -= n * 8
  1061  			nbytes += n
  1062  		}
  1063  		if nbytes >= bufferFlushSize {
  1064  			if w.err != nil {
  1065  				nbytes = 0
  1066  				return
  1067  			}
  1068  			_, w.err = w.writer.Write(w.bytes[:nbytes])
  1069  			nbytes = 0
  1070  		}
  1071  		a, b := encoding[input[0]], encoding[input[1]]
  1072  		bits |= a.code64() << (nbits & 63)
  1073  		bits |= b.code64() << ((nbits + a.len()) & 63)
  1074  		c := encoding[input[2]]
  1075  		nbits += b.len() + a.len()
  1076  		bits |= c.code64() << (nbits & 63)
  1077  		nbits += c.len()
  1078  		input = input[3:]
  1079  	}
  1080  
  1081  	// Remaining...
  1082  	for _, t := range input {
  1083  		if nbits >= 48 {
  1084  			storeLE64(w.bytes[nbytes:], bits)
  1085  			bits >>= 48
  1086  			nbits -= 48
  1087  			nbytes += 6
  1088  			if nbytes >= bufferFlushSize {
  1089  				if w.err != nil {
  1090  					nbytes = 0
  1091  					return
  1092  				}
  1093  				_, w.err = w.writer.Write(w.bytes[:nbytes])
  1094  				nbytes = 0
  1095  			}
  1096  		}
  1097  		// Bitwriting inlined, ~30% speedup
  1098  		c := encoding[t]
  1099  		bits |= c.code64() << (nbits & 63)
  1100  
  1101  		nbits += c.len()
  1102  	}
  1103  	// Restore...
  1104  	w.bits, w.nbits, w.nbytes = bits, nbits, nbytes
  1105  
  1106  	// Flush if needed to have space.
  1107  	if w.nbits >= 48 {
  1108  		w.flushBits()
  1109  	}
  1110  
  1111  	if eof || sync {
  1112  		w.writeCode(w.literalEncoding.codes[endBlockMarker])
  1113  		w.prevHeader = 0
  1114  		w.wroteHuffman = false
  1115  	}
  1116  }
  1117  

View as plain text