decode: Reuse read buffer per decode to speed things up

wader · Sep 14, 2021 · 776a6b3 · 776a6b3
1 parent 798141a
commit 776a6b3
Show file tree

Hide file tree

Showing 13 changed files with 53 additions and 36 deletions.
diff --git a/format/bzip2/bzip2.go b/format/bzip2/bzip2.go
@@ -91,7 +91,7 @@ func gzDecode(d *decode.D, in interface{}) interface{} {
 	deflateR := bzip2.NewReader(compressedBB)
 	uncompressed := &bytes.Buffer{}
 	crc32W := crc32.NewIEEE()
-	if _, err := io.Copy(io.MultiWriter(uncompressed, crc32W), deflateR); err != nil { //nolint:gosec
+	if _, err := decode.Copy(d, io.MultiWriter(uncompressed, crc32W), deflateR); err != nil { //nolint:gosec
 		d.Invalid(err.Error())
 	}
 	// calculatedCRC32 := crc32W.Sum(nil)

diff --git a/format/flac/flac.go b/format/flac/flac.go
@@ -73,7 +73,7 @@ func flacDecode(d *decode.D, in interface{}) interface{} {
 				framesNDecodedSamples += ffo.Samples
 			}
 
-			decode.MustCopy(md5Samples, bytes.NewReader(frameStreamSamplesBuf))
+			decode.MustCopy(d, md5Samples, bytes.NewReader(frameStreamSamplesBuf))
 			streamDecodedSamples += ffo.Samples
 
 			// reuse buffer if possible

diff --git a/format/flac/flac_frame.go b/format/flac/flac_frame.go
@@ -337,12 +337,13 @@ func frameDecode(d *decode.D, in interface{}) interface{} {
 		})
 
 		headerCRC := &crc.CRC{Bits: 8, Table: crc.ATM8Table}
-		decode.MustCopy(headerCRC, d.BitBufRange(frameStart, d.Pos()-frameStart))
+		decode.MustCopy(d, headerCRC, d.BitBufRange(frameStart, d.Pos()-frameStart))
 		d.FieldChecksumLen("crc", 8, headerCRC.Sum(nil), decode.BigEndian)
 
 	})
 
 	var channelSamples [][]int64
+	rs := make([]int64, 0, blockSize)
 	d.FieldArrayFn("subframes", func(d *decode.D) {
 		for channelIndex := 0; channelIndex < int(channels); channelIndex++ {
 			d.FieldStructFn("subframe", func(d *decode.D) {
@@ -418,7 +419,8 @@ func frameDecode(d *decode.D, in interface{}) interface{} {
 
 				decodeResiduals := func() []int64 {
 					// is less than blockSize
-					rs := make([]int64, 0, blockSize)
+					// reset array
+					rs = rs[:0]
 
 					// <2> Residual coding method:
 					// 00 : partitioned Rice coding with 4-bit Rice parameter; RESIDUAL_CODING_METHOD_PARTITIONED_RICE follows
@@ -583,7 +585,7 @@ func frameDecode(d *decode.D, in interface{}) interface{} {
 	d.FieldValidateUFn("byte_align", 0, func() uint64 { return d.U(d.ByteAlignBits()) })
 	// <16> CRC-16 (polynomial = x^16 + x^15 + x^2 + x^0, initialized with 0) of everything before the crc, back to and including the frame header sync code
 	footerCRC := &crc.CRC{Bits: 16, Table: crc.ANSI16Table}
-	decode.MustCopy(footerCRC, d.BitBufRange(frameStart, d.Pos()-frameStart))
+	decode.MustCopy(d, footerCRC, d.BitBufRange(frameStart, d.Pos()-frameStart))
 	d.FieldChecksumLen("footer_crc", 16, footerCRC.Sum(nil), decode.BigEndian)
 
 	streamSamples := len(channelSamples[0])
@@ -623,6 +625,7 @@ func frameDecode(d *decode.D, in interface{}) interface{} {
 
 	interleavedSamplesBuf := ffi.SamplesBuf
 	interleavedSamplesBufLen := len(channelSamples) * streamSamples * bytesPerSample
+	// TODO: decode read buffer?
 	// reuse buffer if possible
 	if interleavedSamplesBuf == nil || len(interleavedSamplesBuf) < interleavedSamplesBufLen {
 		interleavedSamplesBuf = make([]byte, interleavedSamplesBufLen)

diff --git a/format/gif/gif.go b/format/gif/gif.go
@@ -95,7 +95,7 @@ func gifDecode(d *decode.D, in interface{}) interface{} {
 									d.FieldU8("terminator")
 									seenTerminator = true
 								}
-								decode.MustCopy(dataBytes, b.Copy())
+								decode.MustCopy(d, dataBytes, b.Copy())
 							})
 						}
 					})

diff --git a/format/id3/id3v2.go b/format/id3/id3v2.go
@@ -560,7 +560,7 @@ func decodeFrame(d *decode.D, version int) uint64 {
 	if unsyncFlag {
 		// TODO: DecodeFn
 		// TODO: unknown after frame decode
-		unsyncedBb := decode.MustNewBitBufFromReader(unsyncReader{Reader: d.BitBufRange(d.Pos(), int64(dataSize)*8)})
+		unsyncedBb := decode.MustNewBitBufFromReader(d, unsyncReader{Reader: d.BitBufRange(d.Pos(), int64(dataSize)*8)})
 		d.FieldFormatBitBuf("unsync", unsyncedBb, decode.FormatFn(func(d *decode.D, in interface{}) interface{} {
 			if fn, ok := frames[idNormalized]; ok {
 				fn(d)

diff --git a/format/mpeg/avc_nalu.go b/format/mpeg/avc_nalu.go
@@ -109,7 +109,7 @@ func avcNALUDecode(d *decode.D, in interface{}) interface{} {
 	d.FieldBool("forbidden_zero_bit")
 	d.FieldU2("nal_ref_idc")
 	nalType, _ := d.FieldSymbolMapFn("nal_unit_type", avcNALNames, decode.Symbol{Desc: "Unknown"}, d.U5)
-	unescapedBb := decode.MustNewBitBufFromReader(decode.NALUnescapeReader{Reader: d.BitBufRange(d.Pos(), d.BitsLeft())})
+	unescapedBb := decode.MustNewBitBufFromReader(d, decode.NALUnescapeReader{Reader: d.BitBufRange(d.Pos(), d.BitsLeft())})
 
 	switch nalType {
 	case avcNALCodedSliceNonIDR,

diff --git a/format/mpeg/hevc_nalu.go b/format/mpeg/hevc_nalu.go
@@ -71,7 +71,7 @@ func hevcNALUDecode(d *decode.D, in interface{}) interface{} {
 	nalType, _ := d.FieldStringMapFn("nal_unit_type", hevcNALNames, "Unknown", d.U6, decode.NumberDecimal)
 	d.FieldU6("nuh_layer_id")
 	d.FieldU3("nuh_temporal_id_plus1")
-	unescapedBb := decode.MustNewBitBufFromReader(decode.NALUnescapeReader{Reader: d.BitBufRange(d.Pos(), d.BitsLeft())})
+	unescapedBb := decode.MustNewBitBufFromReader(d, decode.NALUnescapeReader{Reader: d.BitBufRange(d.Pos(), d.BitsLeft())})
 
 	_ = unescapedBb
 	_ = nalType

diff --git a/format/mpeg/mp3_frame.go b/format/mpeg/mp3_frame.go
@@ -358,8 +358,8 @@ func frameDecode(d *decode.D, in interface{}) interface{} {
 
 	crcHash := &crc.CRC{Bits: 16, Current: 0xffff, Table: crc.ANSI16Table}
 	// 2 bytes after sync and some other fields + all of side info
-	decode.MustCopy(crcHash, d.BitBufRange(2*8, 2*8))
-	decode.MustCopy(crcHash, d.BitBufRange(6*8, sideInfoBytes*8))
+	decode.MustCopy(d, crcHash, d.BitBufRange(2*8, 2*8))
+	decode.MustCopy(d, crcHash, d.BitBufRange(6*8, sideInfoBytes*8))
 	crcValue := d.FieldGet("crc")
 	if crcValue != nil {
 		d.FieldRemove("crc")

diff --git a/format/ogg/ogg_page.go b/format/ogg/ogg_page.go
@@ -48,9 +48,9 @@ func pageDecode(d *decode.D, in interface{}) interface{} {
 
 	pageChecksum := d.FieldMustRemove("page_checksum")
 	pageCRC := &crc.CRC{Bits: 32, Table: crc.Poly04c11db7Table}
-	decode.MustCopy(pageCRC, d.BitBufRange(startPos, pageChecksum.Range.Start-startPos))                 // header before checksum
-	decode.MustCopy(pageCRC, bytes.NewReader([]byte{0, 0, 0, 0}))                                        // zero checksum bits
-	decode.MustCopy(pageCRC, d.BitBufRange(pageChecksum.Range.Stop(), endPos-pageChecksum.Range.Stop())) // rest of page
+	decode.MustCopy(d, pageCRC, d.BitBufRange(startPos, pageChecksum.Range.Start-startPos))                 // header before checksum
+	decode.MustCopy(d, pageCRC, bytes.NewReader([]byte{0, 0, 0, 0}))                                        // zero checksum bits
+	decode.MustCopy(d, pageCRC, d.BitBufRange(pageChecksum.Range.Stop(), endPos-pageChecksum.Range.Stop())) // rest of page
 	d.FieldChecksumRange("page_checksum", pageChecksum.Range.Start, pageChecksum.Range.Len, pageCRC.Sum(nil), decode.LittleEndian)
 
 	return p

diff --git a/format/png/png.go b/format/png/png.go
@@ -189,7 +189,7 @@ func pngDecode(d *decode.D, in interface{}) interface{} {
 		})
 
 		chunkCRC := crc32.NewIEEE()
-		decode.MustCopy(chunkCRC, d.BitBufRange(crcStartPos, d.Pos()-crcStartPos))
+		decode.MustCopy(d, chunkCRC, d.BitBufRange(crcStartPos, d.Pos()-crcStartPos))
 		d.FieldChecksumLen("crc", 32, chunkCRC.Sum(nil), decode.BigEndian)
 	})
 

diff --git a/pkg/decode/decode.go b/pkg/decode/decode.go
@@ -104,6 +104,7 @@ type DecodeOptions struct {
 	StartOffset   int64
 	FormatOptions map[string]interface{}
 	FormatInArg   interface{}
+	ReadBuf       *[]byte
 }
 
 func (DecodeOptions) decodeOptions() {}
@@ -208,7 +209,7 @@ type D struct {
 
 	bitBuf *bitio.Buffer
 
-	bitsBuf []byte
+	readBuf *[]byte
 }
 
 // TODO: new struct decoder?
@@ -228,10 +229,21 @@ func NewDecoder(name string, description string, format *Format, bb *bitio.Buffe
 		},
 		Options: opts.FormatOptions,
 
-		bitBuf: cbb,
+		bitBuf:  cbb,
+		readBuf: opts.ReadBuf,
 	}
 }
 
+func (d *D) AllocReadBuf(n int) []byte {
+	if d.readBuf == nil {
+		d.readBuf = new([]byte)
+	}
+	if cap(*d.readBuf) < n {
+		*d.readBuf = make([]byte, n)
+	}
+	return *d.readBuf
+}
+
 func (d *D) FillGaps(namePrefix string) {
 	// TODO: d.Value is array?
 
@@ -446,7 +458,8 @@ func (d *D) FieldFormatr(name string, bitBuf *bitio.Buffer, v interface{}) *D {
 		},
 		Options: d.Options,
 
-		bitBuf: bitBuf,
+		bitBuf:  bitBuf,
+		readBuf: d.readBuf,
 	}
 }
 
@@ -842,7 +855,7 @@ func (d *D) DecodeRangeFn(firstBit int64, nBits int64, fn func(d *D)) {
 
 func (d *D) Format(formats []*Format, opts ...Options) interface{} {
 	bb := d.BitBufRange(d.Pos(), d.BitsLeft())
-	opts = append(opts, DecodeOptions{IsRoot: false, StartOffset: d.Pos()})
+	opts = append(opts, DecodeOptions{ReadBuf: d.readBuf, IsRoot: false, StartOffset: d.Pos()})
 	dv, v, err := decode("", "", bb, formats, opts)
 	if dv == nil || dv.Errors() != nil {
 		panic(err)
@@ -870,7 +883,7 @@ func (d *D) Format(formats []*Format, opts ...Options) interface{} {
 
 func (d *D) FieldTryFormat(name string, formats []*Format, opts ...Options) (*Value, interface{}, error) {
 	bb := d.BitBufRange(d.Pos(), d.BitsLeft())
-	opts = append(opts, DecodeOptions{IsRoot: false, StartOffset: d.Pos()})
+	opts = append(opts, DecodeOptions{ReadBuf: d.readBuf, IsRoot: false, StartOffset: d.Pos()})
 	dv, v, err := decode(name, "", bb, formats, opts)
 	if dv == nil || dv.Errors() != nil {
 		return nil, nil, err
@@ -894,7 +907,7 @@ func (d *D) FieldFormat(name string, formats []*Format, opts ...Options) (*Value
 
 func (d *D) FieldTryFormatLen(name string, nBits int64, formats []*Format, opts ...Options) (*Value, interface{}, error) {
 	bb := d.BitBufRange(d.Pos(), nBits)
-	opts = append(opts, DecodeOptions{IsRoot: false, StartOffset: d.Pos()})
+	opts = append(opts, DecodeOptions{ReadBuf: d.readBuf, IsRoot: false, StartOffset: d.Pos()})
 	dv, v, err := decode(name, "", bb, formats, opts)
 	if dv == nil || dv.Errors() != nil {
 		return nil, nil, err
@@ -919,7 +932,7 @@ func (d *D) FieldFormatLen(name string, nBits int64, formats []*Format, opts ...
 // TODO: return decooder?
 func (d *D) FieldTryFormatRange(name string, firstBit int64, nBits int64, formats []*Format, opts ...Options) (*Value, interface{}, error) {
 	bb := d.BitBufRange(firstBit, nBits)
-	opts = append(opts, DecodeOptions{IsRoot: false, StartOffset: firstBit})
+	opts = append(opts, DecodeOptions{ReadBuf: d.readBuf, IsRoot: false, StartOffset: firstBit})
 	dv, v, err := decode(name, "", bb, formats, opts)
 	if dv == nil || dv.Errors() != nil {
 		return nil, nil, err
@@ -940,7 +953,7 @@ func (d *D) FieldFormatRange(name string, firstBit int64, nBits int64, formats [
 }
 
 func (d *D) FieldTryFormatBitBuf(name string, bb *bitio.Buffer, formats []*Format, opts ...Options) (*Value, interface{}, error) {
-	opts = append(opts, DecodeOptions{IsRoot: true})
+	opts = append(opts, DecodeOptions{ReadBuf: d.readBuf, IsRoot: true})
 	dv, v, err := decode(name, "", bb, formats, opts)
 	if dv == nil || dv.Errors() != nil {
 		return nil, nil, err

diff --git a/pkg/decode/decode_readers.go b/pkg/decode/decode_readers.go
@@ -71,12 +71,7 @@ func (d *D) TryPeekBits(nBits int) (uint64, error) {
 // Bits reads nBits bits from buffer
 func (d *D) bits(nBits int) (uint64, error) {
 	// 64 bits max, 9 byte worse case if not byte aligned
-	buf := d.bitsBuf
-	if buf == nil {
-		d.bitsBuf = make([]byte, 9)
-		buf = d.bitsBuf
-	}
-
+	buf := d.AllocReadBuf(9)
 	_, err := bitio.ReadFull(d.bitBuf, buf, nBits)
 	if err != nil {
 		return 0, err

diff --git a/pkg/decode/io.go b/pkg/decode/io.go
@@ -7,18 +7,24 @@ import (
 	"github.com/wader/fq/pkg/bitio"
 )
 
-func MustCopy(r io.Writer, w io.Reader) int64 {
-	n, err := io.Copy(r, w)
+func Copy(d *D, r io.Writer, w io.Reader) (int64, error) {
+	// TODO: what size?
+	buf := d.AllocReadBuf(64 * 1024)
+	return io.CopyBuffer(r, w, buf)
+}
+
+func MustCopy(d *D, r io.Writer, w io.Reader) int64 {
+	n, err := Copy(d, r, w)
 	if err != nil {
-		panic(IOError{Err: err, Op: "MustCopy"})
+		panic(IOError{Err: err, Op: "MustCopyBuffer"})
 	}
 	return n
 }
 
-func MustNewBitBufFromReader(r io.Reader) *bitio.Buffer {
-	buf := &bytes.Buffer{}
-	MustCopy(buf, r)
-	return bitio.NewBufferFromBytes(buf.Bytes(), -1)
+func MustNewBitBufFromReader(d *D, r io.Reader) *bitio.Buffer {
+	b := &bytes.Buffer{}
+	MustCopy(d, b, r)
+	return bitio.NewBufferFromBytes(b.Bytes(), -1)
 }
 
 // TODO: move?