huffman.gno
12.85 Kb · 439 lines
1// Package huffman compresses a string with canonical Huffman codes: the
2// classic frequency-based prefix code, in the one form that is safe to put on
3// a chain.
4//
5// Two properties matter here and neither is about compression.
6//
7// The codes are CANONICAL. A plain Huffman tree is not unique: two equally
8// optimal trees give the same total bit count and different codes, and which
9// one you get depends on how ties are broken while merging. Canonical codes
10// throw the tree away and keep only the code LENGTHS, then re-derive the bits
11// by sorting on (length, symbol). The same input therefore always produces the
12// same table, which is what a deterministic Render and a reproducible payload
13// both need, and the table a decoder has to carry shrinks to one length per
14// symbol.
15//
16// Decoding VALIDATES. [Encoded] is a plain struct with exported fields, so a
17// caller can build one by hand or read one out of storage, and a decoder that
18// trusts what it is handed will happily return a string nobody ever encoded.
19// [Encoded.Decode] rebuilds the canonical table from the lengths it was given,
20// refuses a payload whose own Bits disagree with it, and checks the Kraft
21// equality, the bit budget and the decoded size before returning.
22//
23// A live demo of this package is at
24// [r/moul/x/daily/huffmandemo](/r/moul/x/daily/huffmandemo/v0).
25package huffman
26
27import (
28 "errors"
29 "strings"
30)
31
32// MaxLen bounds input so gas stays predictable.
33const MaxLen = 4096
34
35// maxCodeLen is the longest code this package will emit or accept.
36//
37// It is not arbitrary. The worst case for Huffman is a Fibonacci frequency
38// distribution, where n symbols need a code of n-1 bits and the input has to
39// be at least F(n+2) bytes long. F(34) is 5,702,887, so reaching 32 bits needs
40// megabytes of input and [MaxLen] stops at 4096, which bounds a code here to
41// 18 bits. The ceiling exists so the overflow is a returned error rather than
42// a shift nobody checked.
43const maxCodeLen = 32
44
45var (
46 // ErrEmpty is returned for an input with no bytes in it: there is no such
47 // thing as a code for an empty alphabet.
48 ErrEmpty = errors.New("huffman: empty input")
49 // ErrTooLong is returned when input exceeds MaxLen.
50 ErrTooLong = errors.New("huffman: input too long")
51 // ErrTooDeep is returned when the optimal code would be longer than 32
52 // bits. Unreachable at MaxLen, and returned rather than panicked because a
53 // realm's Render can never be fixed once the path is published.
54 ErrTooDeep = errors.New("huffman: code longer than 32 bits")
55 // ErrCorrupt is returned by Decode when the payload is not one this
56 // package could have produced.
57 ErrCorrupt = errors.New("huffman: corrupt payload")
58)
59
60// Code is one symbol's canonical prefix code.
61type Code struct {
62 Symbol byte
63 Len int // length in bits, 1 to 32
64 Bits uint32 // the code, right aligned, most significant bit first
65}
66
67// String renders the code as its bit pattern, "110" and not "6".
68func (c Code) String() string {
69 var b strings.Builder
70 for i := c.Len - 1; i >= 0; i-- {
71 if (c.Bits>>uint(i))&1 == 1 {
72 b.WriteByte('1')
73 } else {
74 b.WriteByte('0')
75 }
76 }
77 return b.String()
78}
79
80// Encoded is a compressed payload: the canonical table, which is the whole
81// header a decoder needs, and the packed bit stream.
82type Encoded struct {
83 Codes []Code // sorted by (Len, Symbol)
84 Data []byte // the bit stream, most significant bit of each byte first
85 NBits int // significant bits in Data; the rest of the last byte is padding
86 Size int // length of the original input, in bytes
87}
88
89// Saved reports how many bits the encoding removed, which is negative when the
90// alphabet is flat enough that the codes cost more than the bytes did.
91func (e *Encoded) Saved() int { return e.Size*8 - e.NBits }
92
93// BitString renders the stream as '0' and '1', padding excluded. It is meant
94// for demos and error messages: the result is NBits characters long.
95func (e *Encoded) BitString() string {
96 var b strings.Builder
97 for i := 0; i < e.NBits && i/8 < len(e.Data); i++ {
98 if (e.Data[i/8]>>uint(7-i%8))&1 == 1 {
99 b.WriteByte('1')
100 } else {
101 b.WriteByte('0')
102 }
103 }
104 return b.String()
105}
106
107// Table returns the canonical code table for s, sorted by (Len, Symbol),
108// without encoding anything.
109func Table(s string) ([]Code, error) {
110 if len(s) > MaxLen {
111 return nil, ErrTooLong
112 }
113 if len(s) == 0 {
114 return nil, ErrEmpty
115 }
116
117 var freq [256]int
118 for i := 0; i < len(s); i++ {
119 freq[s[i]]++
120 }
121
122 lengths, err := codeLengths(freq)
123 if err != nil {
124 return nil, err
125 }
126
127 codes := make([]Code, 0, len(lengths))
128 for _, l := range lengths {
129 codes = append(codes, Code{Symbol: l.Symbol, Len: l.Len})
130 }
131 sortCodes(codes)
132 assignCanonical(codes)
133 return codes, nil
134}
135
136// Encode builds the canonical table for s and packs s with it.
137func Encode(s string) (*Encoded, error) {
138 codes, err := Table(s)
139 if err != nil {
140 return nil, err
141 }
142
143 var (
144 byLen [256]int
145 byBits [256]uint32
146 declared [256]bool
147 )
148 for _, c := range codes {
149 byLen[c.Symbol] = c.Len
150 byBits[c.Symbol] = c.Bits
151 declared[c.Symbol] = true
152 }
153
154 e := &Encoded{Codes: codes, Size: len(s)}
155 for i := 0; i < len(s); i++ {
156 sym := s[i]
157 if !declared[sym] {
158 // Unreachable: the table is built from this very string.
159 return nil, ErrCorrupt
160 }
161 e.writeBits(byBits[sym], byLen[sym])
162 }
163 return e, nil
164}
165
166// writeBits appends the low n bits of v to the stream, most significant first.
167func (e *Encoded) writeBits(v uint32, n int) {
168 for i := n - 1; i >= 0; i-- {
169 if e.NBits%8 == 0 {
170 e.Data = append(e.Data, 0)
171 }
172 if (v>>uint(i))&1 == 1 {
173 e.Data[e.NBits/8] |= 1 << uint(7-e.NBits%8)
174 }
175 e.NBits++
176 }
177}
178
179// Decode validates the payload and returns the original string.
180//
181// Everything it checks is something a hand-built or storage-read Encoded can
182// get wrong: an empty or over-long table, a zero length, a table out of
183// canonical order, a repeated symbol, Bits that disagree with the lengths, a
184// code set that is not a complete prefix code, a zero-bit or zero-size
185// payload, a Data slice that is not exactly the bytes NBits needs, nonzero
186// padding in the last byte, a stream ending mid-code, and a decoded length
187// that is not Size.
188func (e *Encoded) Decode() (string, error) {
189 if len(e.Codes) == 0 || len(e.Codes) > 256 || e.NBits < 1 || e.Size < 1 {
190 // Encode refuses an empty input, so a zero-bit or zero-size payload
191 // is not one this package produced. Accepting it would make "" a
192 // valid decoding of a table that was never used.
193 return "", ErrCorrupt
194 }
195 if e.Size > MaxLen {
196 return "", ErrCorrupt
197 }
198 // Data holds exactly the bytes NBits needs, and the padding in the last
199 // one is zero. Neither is cosmetic: spare bytes and dirty padding are two
200 // ways to carry a second payload that decodes to the same string.
201 if len(e.Data) != (e.NBits+7)/8 {
202 return "", ErrCorrupt
203 }
204 if pad := len(e.Data)*8 - e.NBits; pad > 0 {
205 if e.Data[len(e.Data)-1]&byte((1<<uint(pad))-1) != 0 {
206 return "", ErrCorrupt
207 }
208 }
209
210 var (
211 count [maxCodeLen + 1]int
212 seen [256]bool
213 )
214 for i, c := range e.Codes {
215 if c.Len < 1 || c.Len > maxCodeLen {
216 return "", ErrCorrupt
217 }
218 if i > 0 {
219 prev := e.Codes[i-1]
220 if c.Len < prev.Len || (c.Len == prev.Len && c.Symbol <= prev.Symbol) {
221 return "", ErrCorrupt
222 }
223 }
224 // The order check only compares neighbours, so it catches a repeat at
225 // the same length and misses one at two different lengths. Encode
226 // never emits two codes for one byte, and a decoder that accepted it
227 // would map two streams to the same string.
228 if seen[c.Symbol] {
229 return "", ErrCorrupt
230 }
231 seen[c.Symbol] = true
232 count[c.Len]++
233 }
234
235 // Kraft equality: a complete prefix code fills the unit interval exactly.
236 // The single-symbol alphabet is the one exception this package emits, and
237 // it is incomplete by construction: its leaf has no sibling.
238 if !(len(e.Codes) == 1 && e.Codes[0].Len == 1) {
239 kraft := uint64(0)
240 for _, c := range e.Codes {
241 kraft += uint64(1) << uint(maxCodeLen-c.Len)
242 }
243 if kraft != uint64(1)<<uint(maxCodeLen) {
244 return "", ErrCorrupt
245 }
246 }
247
248 // The lengths are the header. Re-derive the bits from them and refuse a
249 // payload that carries different ones.
250 want := make([]Code, len(e.Codes))
251 copy(want, e.Codes)
252 assignCanonical(want)
253 for i, c := range e.Codes {
254 if c.Bits != want[i].Bits {
255 return "", ErrCorrupt
256 }
257 }
258
259 var (
260 first [maxCodeLen + 2]uint32
261 firstIdx [maxCodeLen + 2]int
262 code uint32
263 idx int
264 )
265 for l := 1; l <= maxCodeLen; l++ {
266 first[l] = code
267 firstIdx[l] = idx
268 code = (code + uint32(count[l])) << 1
269 idx += count[l]
270 }
271
272 var (
273 out strings.Builder
274 cur uint32
275 curLen int
276 )
277 for i := 0; i < e.NBits; i++ {
278 cur = cur<<1 | uint32((e.Data[i/8]>>uint(7-i%8))&1)
279 curLen++
280 if curLen > maxCodeLen {
281 return "", ErrCorrupt
282 }
283 if n := count[curLen]; n > 0 && cur >= first[curLen] && cur-first[curLen] < uint32(n) {
284 out.WriteByte(e.Codes[firstIdx[curLen]+int(cur-first[curLen])].Symbol)
285 cur, curLen = 0, 0
286 }
287 }
288 if curLen != 0 {
289 return "", ErrCorrupt // the stream ends in the middle of a code
290 }
291 got := out.String()
292 if len(got) != e.Size {
293 return "", ErrCorrupt
294 }
295 return got, nil
296}
297
298// node is one entry in the merge forest. A leaf carries a symbol, an internal
299// node carries the two indices it was merged from.
300type node struct {
301 weight int
302 symbol int // -1 for an internal node
303 left, right int
304}
305
306// codeLengths returns the optimal length per present symbol, deterministically.
307//
308// It is the two-queue variant of Huffman's algorithm: leaves pre-sorted by
309// (frequency, symbol) in one queue, merged nodes appended to a second in the
310// order they are created, which is already non-decreasing. Both queues are
311// then read from the front, so no priority queue and no map is involved and
312// nothing depends on heap tie-breaking. Ties prefer the leaf queue, which
313// keeps the shallower codes on the more frequent symbols.
314func codeLengths(freq [256]int) ([]Code, error) {
315 nodes := make([]node, 0, 511)
316 leaves := make([]int, 0, 256)
317 for sym := 0; sym < 256; sym++ {
318 if freq[sym] == 0 {
319 continue
320 }
321 nodes = append(nodes, node{weight: freq[sym], symbol: sym, left: -1, right: -1})
322 leaves = append(leaves, len(nodes)-1)
323 }
324 sortLeaves(nodes, leaves)
325
326 // One distinct symbol: the leaf has nothing to pair with, so the tree walk
327 // below would give it a zero-bit code and no stream at all.
328 if len(leaves) == 1 {
329 return []Code{{Symbol: byte(nodes[leaves[0]].symbol), Len: 1}}, nil
330 }
331
332 merged := make([]int, 0, 256)
333 headLeaf, headMerged := 0, 0
334 pop := func() int {
335 takeLeaf := headLeaf < len(leaves)
336 if takeLeaf && headMerged < len(merged) {
337 takeLeaf = nodes[leaves[headLeaf]].weight <= nodes[merged[headMerged]].weight
338 }
339 if takeLeaf {
340 i := leaves[headLeaf]
341 headLeaf++
342 return i
343 }
344 i := merged[headMerged]
345 headMerged++
346 return i
347 }
348 for (len(leaves)-headLeaf)+(len(merged)-headMerged) > 1 {
349 a := pop()
350 b := pop()
351 nodes = append(nodes, node{
352 weight: nodes[a].weight + nodes[b].weight,
353 symbol: -1,
354 left: a,
355 right: b,
356 })
357 merged = append(merged, len(nodes)-1)
358 }
359 root := merged[len(merged)-1]
360
361 // Depth-first walk, explicit stack: a chain's call depth is not a place to
362 // be clever, and the tree can be 18 levels deep at MaxLen.
363 out := make([]Code, 0, len(leaves))
364 type frame struct{ idx, depth int }
365 stack := []frame{{root, 0}}
366 for len(stack) > 0 {
367 f := stack[len(stack)-1]
368 stack = stack[:len(stack)-1]
369 n := nodes[f.idx]
370 if n.symbol >= 0 {
371 if f.depth > maxCodeLen {
372 return nil, ErrTooDeep
373 }
374 out = append(out, Code{Symbol: byte(n.symbol), Len: f.depth})
375 continue
376 }
377 stack = append(stack, frame{n.left, f.depth + 1}, frame{n.right, f.depth + 1})
378 }
379 return out, nil
380}
381
382// sortLeaves orders leaf indices by (weight, symbol), ascending. Insertion
383// sort: at most 256 entries, and gno has no sort.Slice.
384func sortLeaves(nodes []node, leaves []int) {
385 for i := 1; i < len(leaves); i++ {
386 v := leaves[i]
387 j := i - 1
388 for j >= 0 && leafLess(nodes[v], nodes[leaves[j]]) {
389 leaves[j+1] = leaves[j]
390 j--
391 }
392 leaves[j+1] = v
393 }
394}
395
396func leafLess(a, b node) bool {
397 if a.weight != b.weight {
398 return a.weight < b.weight
399 }
400 return a.symbol < b.symbol
401}
402
403// sortCodes orders a table by (Len, Symbol), ascending, which is the canonical
404// order every other function here assumes.
405func sortCodes(codes []Code) {
406 for i := 1; i < len(codes); i++ {
407 v := codes[i]
408 j := i - 1
409 for j >= 0 && codeLess(v, codes[j]) {
410 codes[j+1] = codes[j]
411 j--
412 }
413 codes[j+1] = v
414 }
415}
416
417func codeLess(a, b Code) bool {
418 if a.Len != b.Len {
419 return a.Len < b.Len
420 }
421 return a.Symbol < b.Symbol
422}
423
424// assignCanonical fills in Bits from Len alone: start at zero, add one per
425// code, shift left on every increase in length. The table must already be
426// sorted by (Len, Symbol).
427func assignCanonical(codes []Code) {
428 if len(codes) == 0 {
429 return
430 }
431 code := uint32(0)
432 prev := codes[0].Len
433 for i := range codes {
434 code <<= uint(codes[i].Len - prev)
435 prev = codes[i].Len
436 codes[i].Bits = code
437 code++
438 }
439}
