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

View as plain text