~bigbes/sr-ht-dolt

ref: bbaaa1f7b2e833698b0db403de50fc6f160c2bae sr-ht-dolt/third_party/gozstd-purego/gozstd.go -rw-r--r-- 7.0 KiB
bbaaa1f7 — Eugene Blikh web: the milestones page says when its rollup is partial 5 days ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
// Package gozstd is a pure-Go, drop-in replacement for the subset of
// github.com/dolthub/gozstd that github.com/dolthub/dolt/go references, backed
// by github.com/klauspost/compress/zstd. It exists so this project can build
// with CGO_ENABLED=0: upstream gozstd bundles the zstd C source and requires
// cgo for every symbol, which is the last hard cgo dependency in the dolt
// import graph (ICU is already avoided via the gms_pure_go build tag).
//
// It is wired in via a `replace github.com/dolthub/gozstd => ./third_party/
// gozstd-purego` directive; dolt itself is not modified.
//
// # Scope and correctness
//
// dolt uses gozstd only from its NBS "archive" subsystem (store/nbs/archive_*),
// and references exactly nine symbols: Compress, CompressDict, Decompress,
// DecompressDict, BuildDict, NewCDict, NewDDict, and the CDict/DDict types. This
// file reproduces those with identical signatures. Streaming (Writer/Reader),
// the Stream* helpers, CompressLevel and NewCDictLevel are part of upstream
// gozstd's API but are unused by dolt, so they are intentionally omitted.
//
// zstd is a standard frame format (RFC 8878) and dictionaries carry an embedded
// dict-ID, so frames and dictionaries produced by upstream libzstd (e.g. by the
// real `dolt` CLI) are decodable here and vice-versa. This project only ever
// runs the DECOMPRESS side at runtime (it serves chunks over remotesapi and
// browses read-only; it never builds archives or runs gc), so the compress and
// dictionary-training paths are compiled but never executed by the server. They
// are implemented faithfully regardless, so the swap stays correct if a future
// caller does reach them.
//
// Concurrency: klauspost's Encoder.EncodeAll and Decoder.DecodeAll are safe for
// concurrent use, matching gozstd's Compress*/Decompress* which dolt calls from
// multiple goroutines.
package gozstd

import (
	"fmt"
	"hash/fnv"
	"sync"

	"github.com/klauspost/compress/zstd"
)

// DefaultCompressionLevel mirrors upstream gozstd's default (zstd level 3).
const DefaultCompressionLevel = 3

// Shared, dictionary-less coders. klauspost's zstd Encoder/Decoder built over a
// nil stream are the documented one-shot EncodeAll/DecodeAll coders, and both
// are safe for concurrent use, so a single lazily-initialized instance backs
// all dictionary-less Compress/Decompress calls.
var (
	encOnce sync.Once
	enc     *zstd.Encoder
	decOnce sync.Once
	dec     *zstd.Decoder
)

func sharedEncoder() *zstd.Encoder {
	encOnce.Do(func() {
		e, err := zstd.NewWriter(nil,
			zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(DefaultCompressionLevel)))
		if err != nil {
			panic(fmt.Errorf("gozstd-purego: new encoder: %w", err))
		}
		enc = e
	})
	return enc
}

func sharedDecoder() *zstd.Decoder {
	decOnce.Do(func() {
		d, err := zstd.NewReader(nil)
		if err != nil {
			panic(fmt.Errorf("gozstd-purego: new decoder: %w", err))
		}
		dec = d
	})
	return dec
}

// Compress appends the zstd-compressed form of src to dst and returns dst.
// Signature and append semantics match upstream gozstd.Compress. Note klauspost
// takes (src, dst) where gozstd takes (dst, src); both append to dst.
func Compress(dst, src []byte) []byte {
	return sharedEncoder().EncodeAll(src, dst)
}

// Decompress appends the decompressed form of src to dst and returns dst.
func Decompress(dst, src []byte) ([]byte, error) {
	return sharedDecoder().DecodeAll(src, dst)
}

// CDict is a compression dictionary. Upstream gozstd wraps an opaque libzstd
// ZSTD_CDict; here it holds a klauspost Encoder pre-loaded with the dictionary,
// which embeds the dictionary's ID into every frame it writes (as libzstd does).
type CDict struct {
	enc *zstd.Encoder
}

// NewCDict builds a CDict from a raw zstd dictionary at the default level.
func NewCDict(dict []byte) (*CDict, error) {
	e, err := zstd.NewWriter(nil,
		zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(DefaultCompressionLevel)),
		zstd.WithEncoderDict(dict))
	if err != nil {
		return nil, fmt.Errorf("gozstd-purego: NewCDict: %w", err)
	}
	return &CDict{enc: e}, nil
}

// Release frees the dictionary's coder. Upstream gozstd requires this to free C
// memory; here it closes the klauspost Encoder. dolt does not currently call it,
// but it is provided for API parity and to release coder goroutines. Safe to
// call on nil.
func (cd *CDict) Release() {
	if cd != nil && cd.enc != nil {
		cd.enc.Close()
	}
}

// DDict is a decompression dictionary: a klauspost Decoder pre-loaded with the
// dictionary. DecodeAll selects it for frames whose embedded dict-ID matches.
type DDict struct {
	dec *zstd.Decoder
}

// NewDDict builds a DDict from a raw zstd dictionary. The dictionary is a
// standard structured (ZDICT) dictionary carrying magic + a dict-ID, which
// klauspost parses via WithDecoderDicts.
func NewDDict(dict []byte) (*DDict, error) {
	d, err := zstd.NewReader(nil, zstd.WithDecoderDicts(dict))
	if err != nil {
		return nil, fmt.Errorf("gozstd-purego: NewDDict: %w", err)
	}
	return &DDict{dec: d}, nil
}

// Release frees the dictionary's coder. Safe to call on nil.
func (dd *DDict) Release() {
	if dd != nil && dd.dec != nil {
		dd.dec.Close()
	}
}

// CompressDict appends the dictionary-compressed form of src to dst using cd.
func CompressDict(dst, src []byte, cd *CDict) []byte {
	return cd.enc.EncodeAll(src, dst)
}

// DecompressDict appends the dictionary-decompressed form of src to dst using
// dd. The frame's embedded dict-ID must match dd's dictionary.
func DecompressDict(dst, src []byte, dd *DDict) ([]byte, error) {
	return dd.dec.DecodeAll(src, dst)
}

// BuildDict trains a zstd dictionary from samples. Upstream gozstd wraps
// libzstd's ZDICT_trainFromBuffer; klauspost's trainer uses a different
// algorithm, so the produced dictionary bytes differ, but the result is a valid
// standard zstd dictionary. This function is only reached by dolt's archive-
// writer / gc code, which this project never runs, so its output is never
// consumed here; it is implemented faithfully so the swap stays correct if a
// future caller does train a dictionary. desiredDictLen has no direct klauspost
// analogue (its trainer sizes the dictionary from the content) and is ignored.
//
// It panics on failure to match upstream gozstd, which likewise cannot return
// an error from this signature; a training failure is a programmer/data error,
// not a recoverable runtime condition, and — being off this project's serving
// path — must surface loudly rather than yield a silently-bad dictionary.
func BuildDict(samples [][]byte, desiredDictLen int) []byte {
	_ = desiredDictLen
	// Derive a stable, non-zero dict-ID from the samples so distinct inputs get
	// distinct IDs (klauspost embeds this ID; libzstd would assign its own).
	h := fnv.New32a()
	for _, s := range samples {
		_, _ = h.Write(s)
	}
	id := h.Sum32()
	if id == 0 {
		id = 1
	}
	d, err := zstd.BuildDict(zstd.BuildDictOptions{
		ID:       id,
		Contents: samples,
		// libzstd's default repcodes; a zeroed offset set is rejected.
		Offsets: [3]int{1, 4, 8},
	})
	if err != nil {
		panic(fmt.Errorf("gozstd-purego: BuildDict: %w", err))
	}
	return d
}