Skip to content

Commit e29d4d3

Browse files
joechenrhclaude
joechenrh
and
claude
authored
flate: use Peek instead of ReadByte for the bufio.Reader decode path (#1169)
* flate: read the bufio decode path via Peek instead of per-byte ReadByte huffmanBufioReader filled the bit buffer with fr.ReadByte() once per input byte. bufio.Reader.ReadByte does not inline, so each call is a barrier that spills the fill loop's fb/fnb registers -- roughly 1.85x slower than the *bytes.Reader variant, whose ReadByte inlines to a slice index. Read from a Peek window instead: peek bufio's buffered slice and index it directly (pbuf[pos]), refilling when drained. Peek only views the buffer, so consumed bytes are Discard()ed before every return, keeping the reader position correct for nextBlock/moreBits and the gzip trailer. peekBufio refills from at most one underlying Read (never Peek(Size()), which would block a streaming/sync writer until the whole buffer fills). Only the *bufio.Reader variant changes; the other four generated variants are byte-identical. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * flate: drop stray blank lines in the generated read sites The $GETBYTE_E$ placeholder started with a newline the template already supplied, so every read site picked up a blank line before its return -- 25 in total, and across all five variants rather than only the bufio one. With that fixed the four non-bufio variants are byte-identical to the stock generated output again (verified by diffing each generated function). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent c0bbc84 commit e29d4d3

3 files changed

Lines changed: 225 additions & 68 deletions

File tree

‎flate/_gen/gen_inflate.go‎

Lines changed: 78 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ func (f *decompressor) $FUNCNAME$() {
4949
// but is smart enough to keep local variables in registers, so use nb and b,
5050
// inline call to moreBits and reassign b,nb back to f on return.
5151
fnb, fb, dict := f.nb, f.b, &f.dict
52-
52+
$PEEKTOP$
5353
switch f.stepState {
5454
case stateInit:
5555
goto readLiteral
@@ -70,12 +70,9 @@ readLiteral:
7070
n := uint(f.hl.maxRead)
7171
for {
7272
for fnb < n {
73-
c, err := fr.ReadByte()
74-
if err != nil {
75-
f.b, f.nb = fb, fnb
73+
$GETBYTE_S$
7674
f.err = noEOF(err)
77-
return
78-
}
75+
$GETBYTE_E$
7976
f.roffset++
8077
fb |= uint32(c) << (fnb & regSizeMaskUint32)
8178
fnb += 8
@@ -88,7 +85,7 @@ readLiteral:
8885
}
8986
if n <= fnb {
9087
if n == 0 {
91-
f.b, f.nb = fb, fnb
88+
$DISCARD$ f.b, f.nb = fb, fnb
9289
if debugDecode {
9390
fmt.Println("huffsym: n==0")
9491
}
@@ -108,15 +105,15 @@ readLiteral:
108105
case v < 256:
109106
dict.writeByte(byte(v))
110107
if dict.availWrite() == 0 {
111-
f.toRead = dict.readFlush()
108+
$DISCARD$ f.toRead = dict.readFlush()
112109
f.step = $FUNCNAME$
113110
f.stepState = stateInit
114111
f.b, f.nb = fb, fnb
115112
return
116113
}
117114
goto readLiteral
118115
case v == 256:
119-
f.b, f.nb = fb, fnb
116+
$DISCARD$ f.b, f.nb = fb, fnb
120117
f.finishBlock()
121118
return
122119
// otherwise, reference to older data
@@ -127,15 +124,12 @@ readLiteral:
127124
length = int(val.length) + 3
128125
n := uint(val.extra)
129126
for fnb < n {
130-
c, err := fr.ReadByte()
131-
if err != nil {
132-
f.b, f.nb = fb, fnb
127+
$GETBYTE_S$
133128
if debugDecode {
134129
fmt.Println("morebits n>0:", err)
135130
}
136131
f.err = err
137-
return
138-
}
132+
$GETBYTE_E$
139133
f.roffset++
140134
fb |= uint32(c) << (fnb&regSizeMaskUint32)
141135
fnb += 8
@@ -144,7 +138,7 @@ readLiteral:
144138
fb >>= n & regSizeMaskUint32
145139
fnb -= n
146140
default:
147-
if debugDecode {
141+
$DISCARD$ if debugDecode {
148142
fmt.Println(v, ">= maxNumLit")
149143
}
150144
f.err = CorruptInputError(f.roffset)
@@ -155,15 +149,12 @@ readLiteral:
155149
var dist uint32
156150
if f.hd == nil {
157151
for fnb < 5 {
158-
c, err := fr.ReadByte()
159-
if err != nil {
160-
f.b, f.nb = fb, fnb
152+
$GETBYTE_S$
161153
if debugDecode {
162154
fmt.Println("morebits f.nb<5:", err)
163155
}
164156
f.err = err
165-
return
166-
}
157+
$GETBYTE_E$
167158
f.roffset++
168159
fb |= uint32(c) << (fnb&regSizeMaskUint32)
169160
fnb += 8
@@ -182,12 +173,9 @@ readLiteral:
182173
// inline call to moreBits and reassign b,nb back to f on return.
183174
for {
184175
for fnb < n {
185-
c, err := fr.ReadByte()
186-
if err != nil {
187-
f.b, f.nb = fb, fnb
176+
$GETBYTE_S$
188177
f.err = noEOF(err)
189-
return
190-
}
178+
$GETBYTE_E$
191179
f.roffset++
192180
fb |= uint32(c) << (fnb & regSizeMaskUint32)
193181
fnb += 8
@@ -200,7 +188,7 @@ readLiteral:
200188
}
201189
if n <= fnb {
202190
if n == 0 {
203-
f.b, f.nb = fb, fnb
191+
$DISCARD$ f.b, f.nb = fb, fnb
204192
if debugDecode {
205193
fmt.Println("huffsym: n==0")
206194
}
@@ -223,15 +211,12 @@ readLiteral:
223211
// have 1 bit in bottom of dist, need nb more.
224212
extra := (dist & 1) << (nb & regSizeMaskUint32)
225213
for fnb < nb {
226-
c, err := fr.ReadByte()
227-
if err != nil {
228-
f.b, f.nb = fb, fnb
214+
$GETBYTE_S$
229215
if debugDecode {
230216
fmt.Println("morebits f.nb<nb:", err)
231217
}
232218
f.err = err
233-
return
234-
}
219+
$GETBYTE_E$
235220
f.roffset++
236221
fb |= uint32(c) << (fnb&regSizeMaskUint32)
237222
fnb += 8
@@ -242,7 +227,7 @@ readLiteral:
242227
dist = 1<<((nb+1)&regSizeMaskUint32) + 1 + extra
243228
// slower: dist = bitMask32[nb+1] + 2 + extra
244229
default:
245-
f.b, f.nb = fb, fnb
230+
$DISCARD$ f.b, f.nb = fb, fnb
246231
if debugDecode {
247232
fmt.Println("dist too big:", dist, maxNumDist)
248233
}
@@ -252,7 +237,7 @@ readLiteral:
252237
253238
// No check on length; encoding can be prescient.
254239
if dist > uint32(dict.histSize()) {
255-
f.b, f.nb = fb, fnb
240+
$DISCARD$ f.b, f.nb = fb, fnb
256241
if debugDecode {
257242
fmt.Println("dist > dict.histSize():", dist, dict.histSize())
258243
}
@@ -274,7 +259,7 @@ copyHistory:
274259
f.copyLen -= cnt
275260
276261
if dict.availWrite() == 0 || f.copyLen > 0 {
277-
f.toRead = dict.readFlush()
262+
$DISCARD$ f.toRead = dict.readFlush()
278263
f.step = $FUNCNAME$ // We need to continue this work
279264
f.stepState = stateDict
280265
f.b, f.nb = fb, fnb
@@ -286,11 +271,67 @@ copyHistory:
286271
}
287272
288273
`
274+
275+
f.WriteString(`
276+
func peekBufio(fr *bufio.Reader) ([]byte, error) {
277+
if fr.Buffered() == 0 {
278+
if _, err := fr.Peek(1); err != nil && fr.Buffered() == 0 {
279+
return nil, err
280+
}
281+
}
282+
return fr.Peek(fr.Buffered())
283+
}
284+
285+
`)
286+
// Each variant pulls the next input byte differently. The default reads via
287+
// fr.ReadByte(). The *bufio.Reader variant reads from a Peek window so the
288+
// per-byte fill inlines as a slice index; since Peek only views the buffer and
289+
// Discard advances it, it must Discard consumed bytes before every return.
290+
// getS/getE wrap each read site around that site's own EOF handling:
291+
// <getS> <site-specific error body> <getE>
292+
// (output indentation is normalized by the go fmt step, so these need not be
293+
// pre-indented).
294+
type readImpl struct{ top, getS, getE, discard string }
295+
296+
readByte := readImpl{
297+
getS: `c, err := fr.ReadByte()
298+
if err != nil {
299+
f.b, f.nb = fb, fnb`,
300+
getE: `return
301+
}`,
302+
}
303+
peek := readImpl{
304+
top: "pbuf, _ := fr.Peek(fr.Buffered())\npos := 0\n",
305+
getS: `if pos >= len(pbuf) {
306+
fr.Discard(pos)
307+
var err error
308+
pbuf, err = peekBufio(fr)
309+
pos = 0
310+
if len(pbuf) == 0 {
311+
f.b, f.nb = fb, fnb`,
312+
getE: `return
313+
}
314+
}
315+
c := pbuf[pos]
316+
pos++`,
317+
discard: "fr.Discard(pos)\n",
318+
}
319+
289320
for i, t := range types {
290-
s := strings.Replace(template, "$FUNCNAME$", "huffman"+names[i], -1)
291-
s = strings.Replace(s, "$TYPE$", t, -1)
292-
f.WriteString(s)
321+
r := readByte
322+
if t == "*bufio.Reader" {
323+
r = peek
324+
}
325+
f.WriteString(strings.NewReplacer(
326+
"$FUNCNAME$", "huffman"+names[i],
327+
"$TYPE$", t,
328+
"$PEEKTOP$", r.top,
329+
"$GETBYTE_S$", r.getS,
330+
"$GETBYTE_E$", r.getE,
331+
"$DISCARD$", r.discard,
332+
).Replace(template))
293333
}
334+
294335
f.WriteString("func (f *decompressor) huffmanBlockDecoder() {\n")
295336
f.WriteString("\tswitch f.r.(type) {\n")
296337
for i, t := range types {

‎flate/inflate_bufio_test.go‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package flate
2+
3+
import (
4+
"bufio"
5+
"bytes"
6+
"io"
7+
"testing"
8+
"testing/iotest"
9+
)
10+
11+
// roundtripBufio round-trips data through readers that force the *bufio.Reader
12+
// decode path (huffmanBufioReader, which reads via Peek and Discards consumed
13+
// bytes at each return), asserting byte-for-byte equality. The OneByte/Half
14+
// wrappers stress the Peek-refill boundaries.
15+
func roundtripBufio(t *testing.T, data []byte) {
16+
t.Helper()
17+
for _, lvl := range []int{0, 1, 6, 9} {
18+
var comp bytes.Buffer
19+
w, err := NewWriter(&comp, lvl)
20+
if err != nil {
21+
t.Fatal(err)
22+
}
23+
if _, err := w.Write(data); err != nil {
24+
t.Fatal(err)
25+
}
26+
if err := w.Close(); err != nil {
27+
t.Fatal(err)
28+
}
29+
c := comp.Bytes()
30+
for name, mk := range map[string]func() io.Reader{
31+
"onebyte": func() io.Reader { return iotest.OneByteReader(bytes.NewReader(c)) },
32+
"half": func() io.Reader { return iotest.HalfReader(bytes.NewReader(c)) },
33+
"bufio": func() io.Reader { return bufio.NewReader(bytes.NewReader(c)) },
34+
} {
35+
got, err := io.ReadAll(NewReader(mk()))
36+
if err != nil {
37+
t.Fatalf("level %d, %s reader: %v", lvl, name, err)
38+
}
39+
if !bytes.Equal(got, data) {
40+
t.Fatalf("level %d, %s reader: round-trip mismatch (%d vs %d bytes)", lvl, name, len(got), len(data))
41+
}
42+
}
43+
}
44+
}
45+
46+
func FuzzInflateBufio(f *testing.F) {
47+
f.Add([]byte(""))
48+
f.Add([]byte("a"))
49+
f.Add(bytes.Repeat([]byte("hello world "), 5000))
50+
f.Add(bytes.Repeat([]byte{0}, 100000))
51+
seq := make([]byte, 20000)
52+
for i := range seq {
53+
seq[i] = byte(i*31 + i/7)
54+
}
55+
f.Add(seq)
56+
f.Fuzz(func(t *testing.T, data []byte) {
57+
if len(data) > 1<<18 { // keep fuzz execs fast; larger inputs add no new paths here
58+
return
59+
}
60+
roundtripBufio(t, data)
61+
})
62+
}

0 commit comments

Comments
 (0)