internal/compress: Import flate and such from klauspost/compress

6 files changed, 950 insertions, 0 deletions

diff --git a/internal/compress/zlib/reader.go b/internal/compress/zlib/reader.go
new file mode 100644
index 00000000..2d009887
--- /dev/null
+++ b/internal/compress/zlib/reader.go
@@ -0,0 +1,205 @@
+// Copyright 2009 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+/*
+Package zlib implements reading and writing of zlib format compressed data,
+as specified in RFC 1950.
+
+This package differs from the standard library's compress/zlib package
+in that it pools readers and writers to reduce allocations.
+
+Note that closing a reader or writer causes it to be returned to a pool
+for reuse. Therefore, the caller must not retain references to a
+reader or writer after closing it; in the standard library's
+compress/zlib package, it is legal to Reset a closed reader or writer
+and continue using it; that is not allowed here, so there is simply no
+Resetter interface.
+
+The implementation provides filters that uncompress during reading
+and compress during writing.  For example, to write compressed data
+to a buffer:
+
+	var b bytes.Buffer
+	w := zlib.NewWriter(&b)
+	w.Write([]byte("hello, world\n"))
+	w.Close()
+
+and to read that data back:
+
+	r, err := zlib.NewReader(&b)
+	io.Copy(os.Stdout, r)
+	r.Close()
+*/
+package zlib
+
+import (
+	"encoding/binary"
+	"errors"
+	"hash"
+	"io"
+	"sync"
+
+	"codeberg.org/lindenii/furgit/internal/compress/flate"
+)
+
+const (
+	zlibDeflate   = 8
+	zlibMaxWindow = 7
+)
+
+var (
+	// ErrChecksum is returned when reading ZLIB data that has an invalid checksum.
+	ErrChecksum = errors.New("zlib: invalid checksum")
+	// ErrDictionary is returned when reading ZLIB data that has an invalid dictionary.
+	ErrDictionary = errors.New("zlib: invalid dictionary")
+	// ErrHeader is returned when reading ZLIB data that has an invalid header.
+	ErrHeader = errors.New("zlib: invalid header")
+)
+
+var readerPool = sync.Pool{
+	New: func() any {
+		r := new(Reader)
+
+		return r
+	},
+}
+
+// Reader reads and verifies one zlib stream.
+//
+// Reader implements io.ReadCloser.
+type Reader struct {
+	r            flate.Reader
+	decompressor io.ReadCloser
+	digest       hash.Hash32
+	counter      *countingFlateReader
+	err          error
+	scratch      [4]byte
+}
+
+// countingFlateReader wraps flate input and tracks consumed bytes.
+type countingFlateReader struct {
+	inner flate.Reader
+	read  uint64
+}
+
+// Read implements io.Reader.
+func (reader *countingFlateReader) Read(dst []byte) (int, error) {
+	n, err := reader.inner.Read(dst)
+	reader.read += uint64(n)
+
+	return n, err
+}
+
+// ReadByte implements io.ByteReader.
+func (reader *countingFlateReader) ReadByte() (byte, error) {
+	b, err := reader.inner.ReadByte()
+	if err == nil {
+		reader.read++
+	}
+
+	return b, err
+}
+
+// NewReader creates a new ReadCloser.
+// Reads from the returned ReadCloser read and decompress data from r.
+// If r does not implement [io.ByteReader], the decompressor may read more
+// data than necessary from r.
+// It is the caller's responsibility to call Close on the ReadCloser when done.
+func NewReader(r io.Reader) (*Reader, error) {
+	return NewReaderDict(r, nil)
+}
+
+// NewReaderDict is like [NewReader] but uses a preset dictionary.
+// NewReaderDict ignores the dictionary if the compressed data does not refer to it.
+// If the compressed data refers to a different dictionary, NewReaderDict returns [ErrDictionary].
+func NewReaderDict(r io.Reader, dict []byte) (*Reader, error) {
+	v := readerPool.Get()
+
+	z, ok := v.(*Reader)
+	if !ok {
+		panic("zlib: pool returned unexpected type")
+	}
+
+	err := z.reset(r, dict)
+	if err != nil {
+		return nil, err
+	}
+
+	return z, nil
+}
+
+// Read decompresses bytes from receiver into p.
+func (z *Reader) Read(p []byte) (int, error) {
+	if z.err != nil {
+		return 0, z.err
+	}
+
+	var n int
+
+	n, z.err = z.decompressor.Read(p)
+
+	_, err := z.digest.Write(p[0:n])
+	if err != nil {
+		z.err = err
+
+		return n, z.err
+	}
+
+	if !errors.Is(z.err, io.EOF) {
+		// In the normal case we return here.
+		return n, z.err
+	}
+
+	// Finished file; check checksum.
+	_, err = io.ReadFull(z.r, z.scratch[0:4])
+	if err != nil {
+		if errors.Is(err, io.EOF) {
+			err = io.ErrUnexpectedEOF
+		}
+
+		z.err = err
+
+		return n, z.err
+	}
+	// ZLIB (RFC 1950) is big-endian, unlike GZIP (RFC 1952).
+	checksum := binary.BigEndian.Uint32(z.scratch[:4])
+	if checksum != z.digest.Sum32() {
+		z.err = ErrChecksum
+
+		return n, z.err
+	}
+
+	return n, io.EOF
+}
+
+// InputConsumed returns compressed bytes consumed from stream input.
+//
+// This count includes the zlib header, deflate payload, and zlib checksum
+// trailer bytes read by the reader.
+func (z *Reader) InputConsumed() uint64 {
+	if z.counter == nil {
+		return 0
+	}
+
+	return z.counter.read
+}
+
+// Close does not close the wrapped [io.Reader] originally passed to [NewReader].
+// In order for the ZLIB checksum to be verified, the reader must be
+// fully consumed until the [io.EOF].
+// Close returns the instance to a global pool; you MUST NOT keep references after Close.
+func (z *Reader) Close() error {
+	if z.err != nil && !errors.Is(z.err, io.EOF) {
+		return z.err
+	}
+
+	z.err = z.decompressor.Close()
+	if z.err != nil {
+		return z.err
+	}
+
+	readerPool.Put(z)
+
+	return nil
+}
diff --git a/internal/compress/zlib/reader_reset.go b/internal/compress/zlib/reader_reset.go
new file mode 100644
index 00000000..f374111c
--- /dev/null
+++ b/internal/compress/zlib/reader_reset.go
@@ -0,0 +1,92 @@
+// Copyright 2009 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package zlib
+
+import (
+	"bufio"
+	"encoding/binary"
+	"errors"
+	"io"
+
+	"codeberg.org/lindenii/furgit/internal/adler32"
+	"codeberg.org/lindenii/furgit/internal/compress/flate"
+)
+
+// reset resets receiver to read a new zlib stream.
+func (z *Reader) reset(r io.Reader, dict []byte) error {
+	*z = Reader{decompressor: z.decompressor}
+
+	var input flate.Reader
+	if fr, ok := r.(flate.Reader); ok {
+		input = fr
+	} else {
+		input = bufio.NewReader(r)
+	}
+
+	z.counter = &countingFlateReader{inner: input}
+	z.r = z.counter
+
+	// Read the header (RFC 1950 section 2.2.).
+	_, z.err = io.ReadFull(z.r, z.scratch[0:2])
+	if z.err != nil {
+		if errors.Is(z.err, io.EOF) {
+			z.err = io.ErrUnexpectedEOF
+		}
+
+		return z.err
+	}
+
+	h := binary.BigEndian.Uint16(z.scratch[:2])
+	if (z.scratch[0]&0x0f != zlibDeflate) || (z.scratch[0]>>4 > zlibMaxWindow) || (h%31 != 0) {
+		z.err = ErrHeader
+
+		return z.err
+	}
+
+	haveDict := z.scratch[1]&0x20 != 0
+	if haveDict {
+		_, z.err = io.ReadFull(z.r, z.scratch[0:4])
+		if z.err != nil {
+			if errors.Is(z.err, io.EOF) {
+				z.err = io.ErrUnexpectedEOF
+			}
+
+			return z.err
+		}
+
+		checksum := binary.BigEndian.Uint32(z.scratch[:4])
+		if checksum != adler32.Checksum(dict) {
+			z.err = ErrDictionary
+
+			return z.err
+		}
+	}
+
+	if z.decompressor != nil {
+		resetter, ok := z.decompressor.(flate.Resetter)
+		if !ok {
+			panic("zlib: pooled decompressor does not implement flate.Resetter")
+		}
+
+		z.err = resetter.Reset(z.r, dict)
+		if z.err != nil {
+			return z.err
+		}
+
+		z.digest = adler32.New()
+
+		return nil
+	}
+
+	if haveDict {
+		z.decompressor = flate.NewReaderDict(z.r, dict)
+	} else {
+		z.decompressor = flate.NewReader(z.r)
+	}
+
+	z.digest = adler32.New()
+
+	return nil
+}
diff --git a/internal/compress/zlib/reader_test.go b/internal/compress/zlib/reader_test.go
new file mode 100644
index 00000000..9b534b1e
--- /dev/null
+++ b/internal/compress/zlib/reader_test.go
@@ -0,0 +1,193 @@
+// Copyright 2009 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package zlib
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"testing"
+)
+
+type zlibTest struct {
+	desc       string
+	raw        string
+	compressed []byte
+	dict       []byte
+	err        error
+}
+
+// Compare-to-golden test data was generated by the ZLIB example program at
+// https://www.zlib.net/zpipe.c
+
+var zlibTests = []zlibTest{
+	{
+		"truncated empty",
+		"",
+		[]byte{},
+		nil,
+		io.ErrUnexpectedEOF,
+	},
+	{
+		"truncated dict",
+		"",
+		[]byte{0x78, 0xbb},
+		[]byte{0x00},
+		io.ErrUnexpectedEOF,
+	},
+	{
+		"truncated checksum",
+		"",
+		[]byte{0x78, 0xbb, 0x00, 0x01, 0x00, 0x01, 0xca, 0x48,
+			0xcd, 0xc9, 0xc9, 0xd7, 0x51, 0x28, 0xcf, 0x2f,
+			0xca, 0x49, 0x01, 0x04, 0x00, 0x00, 0xff, 0xff,
+		},
+		[]byte{0x00},
+		io.ErrUnexpectedEOF,
+	},
+	{
+		"empty",
+		"",
+		[]byte{0x78, 0x9c, 0x03, 0x00, 0x00, 0x00, 0x00, 0x01},
+		nil,
+		nil,
+	},
+	{
+		"goodbye",
+		"goodbye, world",
+		[]byte{
+			0x78, 0x9c, 0x4b, 0xcf, 0xcf, 0x4f, 0x49, 0xaa,
+			0x4c, 0xd5, 0x51, 0x28, 0xcf, 0x2f, 0xca, 0x49,
+			0x01, 0x00, 0x28, 0xa5, 0x05, 0x5e,
+		},
+		nil,
+		nil,
+	},
+	{
+		"bad header (CINFO)",
+		"",
+		[]byte{0x88, 0x98, 0x03, 0x00, 0x00, 0x00, 0x00, 0x01},
+		nil,
+		ErrHeader,
+	},
+	{
+		"bad header (FCHECK)",
+		"",
+		[]byte{0x78, 0x9f, 0x03, 0x00, 0x00, 0x00, 0x00, 0x01},
+		nil,
+		ErrHeader,
+	},
+	{
+		"bad checksum",
+		"",
+		[]byte{0x78, 0x9c, 0x03, 0x00, 0x00, 0x00, 0x00, 0xff},
+		nil,
+		ErrChecksum,
+	},
+	{
+		"not enough data",
+		"",
+		[]byte{0x78, 0x9c, 0x03, 0x00, 0x00, 0x00},
+		nil,
+		io.ErrUnexpectedEOF,
+	},
+	{
+		"excess data is silently ignored",
+		"",
+		[]byte{
+			0x78, 0x9c, 0x03, 0x00, 0x00, 0x00, 0x00, 0x01,
+			0x78, 0x9c, 0xff,
+		},
+		nil,
+		nil,
+	},
+	{
+		"dictionary",
+		"Hello, World!\n",
+		[]byte{
+			0x78, 0xbb, 0x1c, 0x32, 0x04, 0x27, 0xf3, 0x00,
+			0xb1, 0x75, 0x20, 0x1c, 0x45, 0x2e, 0x00, 0x24,
+			0x12, 0x04, 0x74,
+		},
+		[]byte{
+			0x48, 0x65, 0x6c, 0x6c, 0x6f, 0x20, 0x57, 0x6f, 0x72, 0x6c, 0x64, 0x0a,
+		},
+		nil,
+	},
+	{
+		"wrong dictionary",
+		"",
+		[]byte{
+			0x78, 0xbb, 0x1c, 0x32, 0x04, 0x27, 0xf3, 0x00,
+			0xb1, 0x75, 0x20, 0x1c, 0x45, 0x2e, 0x00, 0x24,
+			0x12, 0x04, 0x74,
+		},
+		[]byte{
+			0x48, 0x65, 0x6c, 0x6c,
+		},
+		ErrDictionary,
+	},
+	{
+		"truncated zlib stream amid raw-block",
+		"hello",
+		[]byte{
+			0x78, 0x9c, 0x00, 0x0c, 0x00, 0xf3, 0xff, 0x68, 0x65, 0x6c, 0x6c, 0x6f,
+		},
+		nil,
+		io.ErrUnexpectedEOF,
+	},
+	{
+		"truncated zlib stream amid fixed-block",
+		"He",
+		[]byte{
+			0x78, 0x9c, 0xf2, 0x48, 0xcd,
+		},
+		nil,
+		io.ErrUnexpectedEOF,
+	},
+}
+
+func TestDecompressor(t *testing.T) {
+	b := new(bytes.Buffer)
+	for _, tt := range zlibTests {
+		in := bytes.NewReader(tt.compressed)
+
+		zr, err := NewReaderDict(in, tt.dict)
+		if err != nil {
+			if !errors.Is(err, tt.err) {
+				t.Errorf("%s: NewReader: %s", tt.desc, err)
+			}
+
+			continue
+		}
+		defer zr.Close()
+
+		// Read and verify correctness of data.
+		b.Reset()
+
+		n, err := io.Copy(b, zr)
+		if err != nil {
+			if !errors.Is(err, tt.err) {
+				t.Errorf("%s: io.Copy: %v want %v", tt.desc, err, tt.err)
+			}
+
+			continue
+		}
+
+		s := b.String()
+		if s != tt.raw {
+			t.Errorf("%s: got %d-byte %q want %d-byte %q", tt.desc, n, s, len(tt.raw), tt.raw)
+		}
+
+		// Check for sticky errors.
+		if n, err := zr.Read([]byte{0}); n != 0 || !errors.Is(err, io.EOF) {
+			t.Errorf("%s: Read() = (%d, %v), want (0, io.EOF)", tt.desc, n, err)
+		}
+
+		if err := zr.Close(); err != nil {
+			t.Errorf("%s: Close() = %v, want nil", tt.desc, err)
+		}
+	}
+}
diff --git a/internal/compress/zlib/writer.go b/internal/compress/zlib/writer.go
new file mode 100644
index 00000000..8a5562fb
--- /dev/null
+++ b/internal/compress/zlib/writer.go
@@ -0,0 +1,204 @@
+// Copyright 2009 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package zlib
+
+import (
+	"encoding/binary"
+	"fmt"
+	"hash"
+	"io"
+	"sync"
+
+	"codeberg.org/lindenii/furgit/internal/compress/flate"
+)
+
+// These constants are copied from the [flate] package, so that code that imports
+// [compress/zlib] does not also have to import [compress/flate].
+const (
+	NoCompression      = flate.NoCompression
+	BestSpeed          = flate.BestSpeed
+	BestCompression    = flate.BestCompression
+	DefaultCompression = flate.DefaultCompression
+	HuffmanOnly        = flate.HuffmanOnly
+)
+
+// A Writer takes data written to it and writes the compressed
+// form of that data to an underlying writer (see [NewWriter]).
+type Writer struct {
+	w           io.Writer
+	level       int
+	dict        []byte
+	compressor  *flate.Writer
+	digest      hash.Hash32
+	err         error
+	scratch     [4]byte
+	wroteHeader bool
+}
+
+var writerPool = sync.Pool{
+	New: func() any {
+		return new(Writer)
+	},
+}
+
+// NewWriter creates a new [Writer].
+// Writes to the returned Writer are compressed and written to w.
+//
+// It is the caller's responsibility to call Close on the Writer when done.
+// Writes may be buffered and not flushed until Close.
+func NewWriter(w io.Writer) *Writer {
+	z, _ := NewWriterLevelDict(w, DefaultCompression, nil)
+
+	return z
+}
+
+// NewWriterLevel is like [NewWriter] but specifies the compression level instead
+// of assuming [DefaultCompression].
+//
+// The compression level can be [DefaultCompression], [NoCompression], [HuffmanOnly]
+// or any integer value between [BestSpeed] and [BestCompression] inclusive.
+// The error returned will be nil if the level is valid.
+func NewWriterLevel(w io.Writer, level int) (*Writer, error) {
+	return NewWriterLevelDict(w, level, nil)
+}
+
+// NewWriterLevelDict is like [NewWriterLevel] but specifies a dictionary to
+// compress with.
+//
+// The dictionary may be nil. If not, its contents should not be modified until
+// the Writer is closed.
+func NewWriterLevelDict(w io.Writer, level int, dict []byte) (*Writer, error) {
+	if level < HuffmanOnly || level > BestCompression {
+		return nil, fmt.Errorf("zlib: invalid compression level: %d", level)
+	}
+
+	v := writerPool.Get()
+
+	z, ok := v.(*Writer)
+	if !ok {
+		panic("zlib: pool returned unexpected type")
+	}
+
+	// flate.Writer can only be Reset with the same level/dictionary mode.
+	// Reuse it only when the configuration is unchanged and dictionary-free.
+	reuseCompressor := z.compressor != nil && z.level == level && z.dict == nil && dict == nil
+	if !reuseCompressor {
+		z.compressor = nil
+	}
+
+	if z.digest != nil {
+		z.digest.Reset()
+	}
+
+	*z = Writer{
+		w:          w,
+		level:      level,
+		dict:       dict,
+		compressor: z.compressor,
+		digest:     z.digest,
+	}
+	if z.compressor != nil {
+		z.compressor.Reset(w)
+	}
+
+	return z, nil
+}
+
+// Reset clears the state of the [Writer] z such that it is equivalent to its
+// initial state from [NewWriterLevel] or [NewWriterLevelDict], but instead writing
+// to w.
+func (z *Writer) Reset(w io.Writer) {
+	z.w = w
+	// z.level and z.dict left unchanged.
+	if z.compressor != nil {
+		z.compressor.Reset(w)
+	}
+
+	if z.digest != nil {
+		z.digest.Reset()
+	}
+
+	z.err = nil
+	z.scratch = [4]byte{}
+	z.wroteHeader = false
+}
+
+// Write writes a compressed form of p to the underlying [io.Writer]. The
+// compressed bytes are not necessarily flushed until the [Writer] is closed or
+// explicitly flushed.
+func (z *Writer) Write(p []byte) (n int, err error) {
+	if !z.wroteHeader {
+		z.err = z.writeHeader()
+	}
+
+	if z.err != nil {
+		return 0, z.err
+	}
+
+	if len(p) == 0 {
+		return 0, nil
+	}
+
+	n, err = z.compressor.Write(p)
+	if err != nil {
+		z.err = err
+
+		return n, err
+	}
+
+	_, err = z.digest.Write(p)
+	if err != nil {
+		z.err = err
+
+		return 0, z.err
+	}
+
+	return n, err
+}
+
+// Flush flushes the Writer to its underlying [io.Writer].
+func (z *Writer) Flush() error {
+	if !z.wroteHeader {
+		z.err = z.writeHeader()
+	}
+
+	if z.err != nil {
+		return z.err
+	}
+
+	z.err = z.compressor.Flush()
+
+	return z.err
+}
+
+// Close closes the Writer, flushing any unwritten data to the underlying
+// [io.Writer], but does not close the underlying io.Writer.
+func (z *Writer) Close() error {
+	if !z.wroteHeader {
+		z.err = z.writeHeader()
+	}
+
+	if z.err != nil {
+		return z.err
+	}
+
+	z.err = z.compressor.Close()
+	if z.err != nil {
+		return z.err
+	}
+
+	checksum := z.digest.Sum32()
+	// ZLIB (RFC 1950) is big-endian, unlike GZIP (RFC 1952).
+	binary.BigEndian.PutUint32(z.scratch[:], checksum)
+
+	_, z.err = z.w.Write(z.scratch[0:4])
+	if z.err != nil {
+		return z.err
+	}
+
+	writerPool.Put(z)
+
+	return nil
+}
diff --git a/internal/compress/zlib/writer_header.go b/internal/compress/zlib/writer_header.go
new file mode 100644
index 00000000..43d3bdf5
--- /dev/null
+++ b/internal/compress/zlib/writer_header.go
@@ -0,0 +1,71 @@
+// Copyright 2009 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package zlib
+
+import (
+	"encoding/binary"
+
+	"codeberg.org/lindenii/furgit/internal/adler32"
+	"codeberg.org/lindenii/furgit/internal/compress/flate"
+)
+
+// writeHeader writes the ZLIB header.
+func (z *Writer) writeHeader() (err error) {
+	z.wroteHeader = true
+	// ZLIB has a two-byte header (as documented in RFC 1950).
+	// The first four bits is the CINFO (compression info), which is 7 for the default deflate window size.
+	// The next four bits is the CM (compression method), which is 8 for deflate.
+	z.scratch[0] = 0x78
+	// The next two bits is the FLEVEL (compression level). The four values are:
+	// 0=fastest, 1=fast, 2=default, 3=best.
+	// The next bit, FDICT, is set if a dictionary is given.
+	// The final five FCHECK bits form a mod-31 checksum.
+	switch z.level {
+	case -2, 0, 1:
+		z.scratch[1] = 0 << 6
+	case 2, 3, 4, 5:
+		z.scratch[1] = 1 << 6
+	case 6, -1:
+		z.scratch[1] = 2 << 6
+	case 7, 8, 9:
+		z.scratch[1] = 3 << 6
+	default:
+		panic("unreachable")
+	}
+
+	if z.dict != nil {
+		z.scratch[1] |= 1 << 5
+	}
+
+	z.scratch[1] += uint8(31 - binary.BigEndian.Uint16(z.scratch[:2])%31) //#nosec G115
+
+	_, err = z.w.Write(z.scratch[0:2])
+	if err != nil {
+		return err
+	}
+
+	if z.dict != nil {
+		// The next four bytes are the Adler-32 checksum of the dictionary.
+		binary.BigEndian.PutUint32(z.scratch[:], adler32.Checksum(z.dict))
+
+		_, err = z.w.Write(z.scratch[0:4])
+		if err != nil {
+			return err
+		}
+	}
+
+	if z.compressor == nil {
+		// Initialize deflater unless the Writer is being reused
+		// after a Reset call.
+		z.compressor, err = flate.NewWriterDict(z.w, z.level, z.dict)
+		if err != nil {
+			return err
+		}
+
+		z.digest = adler32.New()
+	}
+
+	return nil
+}
diff --git a/internal/compress/zlib/writer_test.go b/internal/compress/zlib/writer_test.go
new file mode 100644
index 00000000..6deffdb8
--- /dev/null
+++ b/internal/compress/zlib/writer_test.go
@@ -0,0 +1,185 @@
+// Copyright 2009 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package zlib
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"os"
+	"testing"
+)
+
+var filenames = []string{
+	"../testdata/gettysburg.txt",
+	"../testdata/e.txt",
+	"../testdata/pi.txt",
+}
+
+var data = []string{
+	"test a reasonable sized string that can be compressed",
+}
+
+// Tests that compressing and then decompressing the given file at the given compression level and dictionary
+// yields equivalent bytes to the original file.
+func testFileLevelDict(t *testing.T, fn string, level int, d string) {
+	// Read the file, as golden output.
+	golden, err := os.Open(fn)
+	if err != nil {
+		t.Errorf("%s (level=%d, dict=%q): %v", fn, level, d, err)
+
+		return
+	}
+	defer golden.Close()
+
+	b0, err0 := io.ReadAll(golden)
+	if err0 != nil {
+		t.Errorf("%s (level=%d, dict=%q): %v", fn, level, d, err0)
+
+		return
+	}
+
+	testLevelDict(t, fn, b0, level, d)
+}
+
+func testLevelDict(t *testing.T, fn string, b0 []byte, level int, d string) {
+	// Make dictionary, if given.
+	var dict []byte
+	if d != "" {
+		dict = []byte(d)
+	}
+
+	// Push data through a pipe that compresses at the write end, and decompresses at the read end.
+	piper, pipew := io.Pipe()
+	defer piper.Close()
+
+	go func() {
+		defer pipew.Close()
+
+		zlibw, err := NewWriterLevelDict(pipew, level, dict)
+		if err != nil {
+			t.Errorf("%s (level=%d, dict=%q): %v", fn, level, d, err)
+
+			return
+		}
+		defer zlibw.Close()
+
+		_, err = zlibw.Write(b0)
+		if err != nil {
+			t.Errorf("%s (level=%d, dict=%q): %v", fn, level, d, err)
+
+			return
+		}
+	}()
+
+	zlibr, err := NewReaderDict(piper, dict)
+	if err != nil {
+		t.Errorf("%s (level=%d, dict=%q): %v", fn, level, d, err)
+
+		return
+	}
+	defer zlibr.Close()
+
+	// Compare the decompressed data.
+	b1, err1 := io.ReadAll(zlibr)
+	if err1 != nil {
+		t.Errorf("%s (level=%d, dict=%q): %v", fn, level, d, err1)
+
+		return
+	}
+
+	if len(b0) != len(b1) {
+		t.Errorf("%s (level=%d, dict=%q): length mismatch %d versus %d", fn, level, d, len(b0), len(b1))
+
+		return
+	}
+
+	for i := range b0 {
+		if b0[i] != b1[i] {
+			t.Errorf("%s (level=%d, dict=%q): mismatch at %d, 0x%02x versus 0x%02x\n", fn, level, d, i, b0[i], b1[i])
+
+			return
+		}
+	}
+}
+
+func TestWriter(t *testing.T) {
+	for i, s := range data {
+		b := []byte(s)
+		tag := fmt.Sprintf("#%d", i)
+		testLevelDict(t, tag, b, DefaultCompression, "")
+		testLevelDict(t, tag, b, NoCompression, "")
+		testLevelDict(t, tag, b, HuffmanOnly, "")
+
+		for level := BestSpeed; level <= BestCompression; level++ {
+			testLevelDict(t, tag, b, level, "")
+		}
+	}
+}
+
+func TestWriterBig(t *testing.T) {
+	for i, fn := range filenames {
+		testFileLevelDict(t, fn, DefaultCompression, "")
+		testFileLevelDict(t, fn, NoCompression, "")
+		testFileLevelDict(t, fn, HuffmanOnly, "")
+
+		for level := BestSpeed; level <= BestCompression; level++ {
+			testFileLevelDict(t, fn, level, "")
+
+			if level >= 1 && testing.Short() {
+				break
+			}
+		}
+
+		if i == 0 && testing.Short() {
+			break
+		}
+	}
+}
+
+func TestWriterDict(t *testing.T) {
+	const dictionary = "0123456789."
+	for i, fn := range filenames {
+		testFileLevelDict(t, fn, DefaultCompression, dictionary)
+		testFileLevelDict(t, fn, NoCompression, dictionary)
+		testFileLevelDict(t, fn, HuffmanOnly, dictionary)
+
+		for level := BestSpeed; level <= BestCompression; level++ {
+			testFileLevelDict(t, fn, level, dictionary)
+
+			if level >= 1 && testing.Short() {
+				break
+			}
+		}
+
+		if i == 0 && testing.Short() {
+			break
+		}
+	}
+}
+
+func TestWriterDictIsUsed(t *testing.T) {
+	var (
+		input = []byte("Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.")
+		buf   bytes.Buffer
+	)
+
+	compressor, err := NewWriterLevelDict(&buf, BestCompression, input)
+	if err != nil {
+		t.Errorf("error in NewWriterLevelDict: %s", err)
+
+		return
+	}
+
+	compressor.Write(input)
+	compressor.Close()
+
+	const expectedMaxSize = 25
+
+	output := buf.Bytes()
+	if len(output) > expectedMaxSize {
+		t.Errorf("result too large (got %d, want <= %d bytes). Is the dictionary being used?", len(output), expectedMaxSize)
+	}
+}