From a19676f556d92fd671a646e80f866b8390a02086 Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Mon, 14 Feb 2022 14:34:07 +0100
Subject: [PATCH] flate: Skip predef checks if many tokens

It is extremely unlikely that predefined codes are better when there are multiple tokens (verified by tests)

Skip checks if there is 250 or more tokens.
---
 flate/huffman_bit_writer.go | 58 ++++++++++++++++++++++---------------
 1 file changed, 35 insertions(+), 23 deletions(-)

diff --git a/flate/huffman_bit_writer.go b/flate/huffman_bit_writer.go
index ff9f89f161..8c1bf42ff8 100644
--- a/flate/huffman_bit_writer.go
+++ b/flate/huffman_bit_writer.go
@@ -8,6 +8,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"io"
+	"math"
 )
 
 const (
@@ -24,6 +25,10 @@ const (
 	codegenCodeCount = 19
 	badCode          = 255
 
+	// maxPredefinedTokens is the maximum number of tokens
+	// where we check if fixed size is smaller.
+	maxPredefinedTokens = 250
+
 	// bufferFlushSize indicates the buffer size
 	// after which bytes are flushed to the writer.
 	// Should preferably be a multiple of 6, since
@@ -577,7 +582,10 @@ func (w *huffmanBitWriter) writeBlock(tokens *tokens, eof bool, input []byte) {
 	// Fixed Huffman baseline.
 	var literalEncoding = fixedLiteralEncoding
 	var offsetEncoding = fixedOffsetEncoding
-	var size = w.fixedSize(extraBits)
+	var size = math.MaxInt32
+	if tokens.n < maxPredefinedTokens {
+		size = w.fixedSize(extraBits)
+	}
 
 	// Dynamic Huffman?
 	var numCodegens int
@@ -678,19 +686,21 @@ func (w *huffmanBitWriter) writeBlockDynamic(tokens *tokens, eof bool, input []b
 			size = reuseSize
 		}
 
-		if preSize := w.fixedSize(extraBits) + 7; usePrefs && preSize < size {
-			// Check if we get a reasonable size decrease.
-			if storable && ssize <= size {
-				w.writeStoredHeader(len(input), eof)
-				w.writeBytes(input)
+		if tokens.n < maxPredefinedTokens {
+			if preSize := w.fixedSize(extraBits) + 7; usePrefs && preSize < size {
+				// Check if we get a reasonable size decrease.
+				if storable && ssize <= size {
+					w.writeStoredHeader(len(input), eof)
+					w.writeBytes(input)
+					return
+				}
+				w.writeFixedHeader(eof)
+				if !sync {
+					tokens.AddEOB()
+				}
+				w.writeTokens(tokens.Slice(), fixedLiteralEncoding.codes, fixedOffsetEncoding.codes)
 				return
 			}
-			w.writeFixedHeader(eof)
-			if !sync {
-				tokens.AddEOB()
-			}
-			w.writeTokens(tokens.Slice(), fixedLiteralEncoding.codes, fixedOffsetEncoding.codes)
-			return
 		}
 		// Check if we get a reasonable size decrease.
 		if storable && ssize <= size {
@@ -723,19 +733,21 @@ func (w *huffmanBitWriter) writeBlockDynamic(tokens *tokens, eof bool, input []b
 		size, numCodegens = w.dynamicSize(w.literalEncoding, w.offsetEncoding, extraBits)
 
 		// Store predefined, if we don't get a reasonable improvement.
-		if preSize := w.fixedSize(extraBits); usePrefs && preSize <= size {
-			// Store bytes, if we don't get an improvement.
-			if storable && ssize <= preSize {
-				w.writeStoredHeader(len(input), eof)
-				w.writeBytes(input)
+		if tokens.n < maxPredefinedTokens {
+			if preSize := w.fixedSize(extraBits); usePrefs && preSize <= size {
+				// Store bytes, if we don't get an improvement.
+				if storable && ssize <= preSize {
+					w.writeStoredHeader(len(input), eof)
+					w.writeBytes(input)
+					return
+				}
+				w.writeFixedHeader(eof)
+				if !sync {
+					tokens.AddEOB()
+				}
+				w.writeTokens(tokens.Slice(), fixedLiteralEncoding.codes, fixedOffsetEncoding.codes)
 				return
 			}
-			w.writeFixedHeader(eof)
-			if !sync {
-				tokens.AddEOB()
-			}
-			w.writeTokens(tokens.Slice(), fixedLiteralEncoding.codes, fixedOffsetEncoding.codes)
-			return
 		}
 
 		if storable && ssize <= size {