Repository: klauspost/compress (Go)

Issue:
zstd block compression seems to ignore any dictionary for the first few jobs
Hi

We're using the zstd package to encode/decode blocks of memory using `EncodeAll` and `DecodeAll`, and using a dictionary to help improve the compression.

Some behaviour we've noticed in the encoder is that for the first `N` number of blocks that get compressed, the compression ratio is higher, where `N` is the concurrency value configured when calling `NewWriter` - it seems almost like the dictionary isn't being used for those first blocks. The result still decompresses fine, so it doesn't look to be any issue with the actual compression, it's just larger than it seems it could be.

I've written a short executable to demonstrate the behaviour. It uses a dictionary I generated by creating some random samples and calling `zstd --train`. It then repeatedly attempts to compress a record that matches the dictionary and prints the size of the compressed data each time.

```
package main

import (
	_ "embed"
	"flag"
	"fmt"

	"github.com/klauspost/compress/zstd"
)

//go:embed zstd.dict
var dict []byte

func main() {
	n := flag.Int("n", 4, "Encoder concurrency level")
	flag.Parse()

	enc, err := zstd.NewWriter(nil, zstd.WithEncoderConcurrency(*n), zstd.WithEncoderDict(dict))
	if err != nil {
		panic(err)
	}

	uncompressed := []byte(`{"FieldA":1,"FieldB":2,"FieldC":3}`)

	for i := 0; i < 5; i++ {
		compressed := enc.EncodeAll(uncompressed, nil)
		fmt.Printf("compressed record to %d\n", len(compressed))
	}
}
```

By varying the concurrency value you can see the difference:
```
$ go run . -n 1
compressed record to 51
compressed record to 36
compressed record to 36
compressed record to 36
compressed record to 36

$ go run . -n 2
compressed record to 51
compressed record to 51
compressed record to 36
compressed record to 36
compressed record to 36

$ go run . -n 3
compressed record to 51
compressed record to 51
compressed record to 51
compressed record to 36
compressed record to 36
```

We can workaround the issue by running some phony encodes with the encoder during initialisation to make sure everything is ready to go when we use it for real data.


Discussion:
Thanks I will check. Probably doesn't initialize it, but only resets correctly.



Patch:
diff --git a/zstd/enc_base.go b/zstd/enc_base.go
index e008b99298..5ca46038ad 100644
--- a/zstd/enc_base.go
+++ b/zstd/enc_base.go
@@ -144,6 +144,7 @@ func (e *fastBase) resetBase(d *dict, singleBlock bool) {
 	} else {
 		e.crc.Reset()
 	}
+	e.blk.dictLitEnc = nil
 	if d != nil {
 		low := e.lowMem
 		if singleBlock {
diff --git a/zstd/enc_dfast.go b/zstd/enc_dfast.go
index 7d425109ad..a154c18f74 100644
--- a/zstd/enc_dfast.go
+++ b/zstd/enc_dfast.go
@@ -1084,7 +1084,7 @@ func (e *doubleFastEncoderDict) Reset(d *dict, singleBlock bool) {
 			}
 		}
 		e.lastDictID = d.id
-		e.allDirty = true
+		allDirty = true
 	}
 	// Reset table to initial state
 	e.cur = e.maxMatchOff
diff --git a/zstd/enc_fast.go b/zstd/enc_fast.go
index cbc626eec6..f45a3da7da 100644
--- a/zstd/enc_fast.go
+++ b/zstd/enc_fast.go
@@ -829,13 +829,12 @@ func (e *fastEncoderDict) Reset(d *dict, singleBlock bool) {
 		}
 		if true {
 			end := e.maxMatchOff + int32(len(d.content)) - 8
-			for i := e.maxMatchOff; i < end; i += 3 {
+			for i := e.maxMatchOff; i < end; i += 2 {
 				const hashLog = tableBits
 
 				cv := load6432(d.content, i-e.maxMatchOff)
-				nextHash := hashLen(cv, hashLog, tableFastHashLen)      // 0 -> 5
-				nextHash1 := hashLen(cv>>8, hashLog, tableFastHashLen)  // 1 -> 6
-				nextHash2 := hashLen(cv>>16, hashLog, tableFastHashLen) // 2 -> 7
+				nextHash := hashLen(cv, hashLog, tableFastHashLen)     // 0 -> 6
+				nextHash1 := hashLen(cv>>8, hashLog, tableFastHashLen) // 1 -> 7
 				e.dictTable[nextHash] = tableEntry{
 					val:    uint32(cv),
 					offset: i,
@@ -844,10 +843,6 @@ func (e *fastEncoderDict) Reset(d *dict, singleBlock bool) {
 					val:    uint32(cv >> 8),
 					offset: i + 1,
 				}
-				e.dictTable[nextHash2] = tableEntry{
-					val:    uint32(cv >> 16),
-					offset: i + 2,
-				}
 			}
 		}
 		e.lastDictID = d.id


Test patch:
diff --git a/zstd/dict_test.go b/zstd/dict_test.go
index 28024d03a4..6f609a8f69 100644
--- a/zstd/dict_test.go
+++ b/zstd/dict_test.go
@@ -189,6 +189,154 @@ func TestEncoder_SmallDict(t *testing.T) {
 	}
 }
 
+func TestEncoder_SmallDictFresh(t *testing.T) {
+	// All files have CRC
+	zr := testCreateZipReader("testdata/dict-tests-small.zip", t)
+	var dicts [][]byte
+	var encs []func() *Encoder
+	var noDictEncs []*Encoder
+	var encNames []string
+
+	for _, tt := range zr.File {
+		if !strings.HasSuffix(tt.Name, ".dict") {
+			continue
+		}
+		func() {
+			r, err := tt.Open()
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer r.Close()
+			in, err := io.ReadAll(r)
+			if err != nil {
+				t.Fatal(err)
+			}
+			dicts = append(dicts, in)
+			for level := SpeedFastest; level < speedLast; level++ {
+				if isRaceTest && level >= SpeedBestCompression {
+					break
+				}
+				level := level
+				encs = append(encs, func() *Encoder {
+					enc, err := NewWriter(nil, WithEncoderConcurrency(1), WithEncoderDict(in), WithEncoderLevel(level), WithWindowSize(1<<17))
+					if err != nil {
+						t.Fatal(err)
+					}
+					return enc
+				})
+				encNames = append(encNames, fmt.Sprint("level-", level.String(), "-dict-", len(dicts)))
+
+				enc, err := NewWriter(nil, WithEncoderConcurrency(1), WithEncoderLevel(level), WithWindowSize(1<<17))
+				if err != nil {
+					t.Fatal(err)
+				}
+				noDictEncs = append(noDictEncs, enc)
+			}
+		}()
+	}
+	dec, err := NewReader(nil, WithDecoderConcurrency(1), WithDecoderDicts(dicts...))
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+	defer dec.Close()
+	for i, tt := range zr.File {
+		if testing.Short() && i > 100 {
+			break
+		}
+		if !strings.HasSuffix(tt.Name, ".zst") {
+			continue
+		}
+		r, err := tt.Open()
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer r.Close()
+		in, err := io.ReadAll(r)
+		if err != nil {
+			t.Fatal(err)
+		}
+		decoded, err := dec.DecodeAll(in, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if testing.Short() && len(decoded) > 1000 {
+			continue
+		}
+
+		t.Run("encodeall-"+tt.Name, func(t *testing.T) {
+			// Attempt to compress with all dicts
+			var b []byte
+			var tmp []byte
+			for i := range encs {
+				i := i
+				t.Run(encNames[i], func(t *testing.T) {
+					enc := encs[i]()
+					defer enc.Close()
+					b = enc.EncodeAll(decoded, b[:0])
+					tmp, err = dec.DecodeAll(in, tmp[:0])
+					if err != nil {
+						t.Fatal(err)
+					}
+					if !bytes.Equal(tmp, decoded) {
+						t.Fatal("output mismatch")
+					}
+
+					tmp = noDictEncs[i].EncodeAll(decoded, tmp[:0])
+
+					if strings.Contains(t.Name(), "dictplain") && strings.Contains(t.Name(), "dict-1") {
+						t.Log("reference:", len(in), "no dict:", len(tmp), "with dict:", len(b), "SAVED:", len(tmp)-len(b))
+						// Check that we reduced this significantly
+						if len(b) > 250 {
+							t.Error("output was bigger than expected")
+						}
+					}
+				})
+			}
+		})
+		t.Run("stream-"+tt.Name, func(t *testing.T) {
+			// Attempt to compress with all dicts
+			var tmp []byte
+			for i := range encs {
+				i := i
+				t.Run(encNames[i], func(t *testing.T) {
+					enc := encs[i]()
+					defer enc.Close()
+					var buf bytes.Buffer
+					enc.ResetContentSize(&buf, int64(len(decoded)))
+					_, err := enc.Write(decoded)
+					if err != nil {
+						t.Fatal(err)
+					}
+					err = enc.Close()
+					if err != nil {
+						t.Fatal(err)
+					}
+					tmp, err = dec.DecodeAll(buf.Bytes(), tmp[:0])
+					if err != nil {
+						t.Fatal(err)
+					}
+					if !bytes.Equal(tmp, decoded) {
+						t.Fatal("output mismatch")
+					}
+					var buf2 bytes.Buffer
+					noDictEncs[i].Reset(&buf2)
+					noDictEncs[i].Write(decoded)
+					noDictEncs[i].Close()
+
+					if strings.Contains(t.Name(), "dictplain") && strings.Contains(t.Name(), "dict-1") {
+						t.Log("reference:", len(in), "no dict:", buf2.Len(), "with dict:", buf.Len(), "SAVED:", buf2.Len()-buf.Len())
+						// Check that we reduced this significantly
+						if buf.Len() > 250 {
+							t.Error("output was bigger than expected")
+						}
+					}
+				})
+			}
+		})
+	}
+}
+
 func benchmarkEncodeAllLimitedBySize(b *testing.B, lowerLimit int, upperLimit int) {
 	zr := testCreateZipReader("testdata/dict-tests-small.zip", b)
 	t := testing.TB(b)