compress/zstd/decoder.go

// Copyright 2019+ Klaus Post. All rights reserved.
// License information can be found in the LICENSE file.
// Based on work by Yann Collet, released under BSD License.

package zstd

import (
	"bytes"
	"errors"
	"io"
	"sync"
)

// Decoder provides decoding of zstandard streams.
// The decoder has been designed to operate without allocations after a warmup.
// This means that you should store the decoder for best performance.
// To re-use a stream decoder, use the Reset(r io.Reader) error to switch to another stream.
// A decoder can safely be re-used even if the previous stream failed.
// To release the resources, you must call the Close() function on a decoder.
type Decoder struct {
	o decoderOptions

	// Unreferenced decoders, ready for use.
	decoders chan *blockDec

	// Streams ready to be decoded.
	stream chan decodeStream

	// Current read position used for Reader functionality.
	current decoderState

	// Custom dictionaries.
	// Always uses copies.
	dicts map[uint32]dict

	// streamWg is the waitgroup for all streams
	streamWg sync.WaitGroup
}

// decoderState is used for maintaining state when the decoder
// is used for streaming.
type decoderState struct {
	// current block being written to stream.
	decodeOutput

	// output in order to be written to stream.
	output chan decodeOutput

	// cancel remaining output.
	cancel chan struct{}

	flushed bool
}

var (
	// Check the interfaces we want to support.
	_ = io.WriterTo(&Decoder{})
	_ = io.Reader(&Decoder{})
)

// NewReader creates a new decoder.
// A nil Reader can be provided in which case Reset can be used to start a decode.
//
// A Decoder can be used in two modes:
//
// 1) As a stream, or
// 2) For stateless decoding using DecodeAll.
//
// Only a single stream can be decoded concurrently, but the same decoder
// can run multiple concurrent stateless decodes. It is even possible to
// use stateless decodes while a stream is being decoded.
//
// The Reset function can be used to initiate a new stream, which is will considerably
// reduce the allocations normally caused by NewReader.
func NewReader(r io.Reader, opts ...DOption) (*Decoder, error) {
	initPredefined()
	var d Decoder
	d.o.setDefault()
	for _, o := range opts {
		err := o(&d.o)
		if err != nil {
			return nil, err
		}
	}
	d.current.output = make(chan decodeOutput, d.o.concurrent)
	d.current.flushed = true

	// Transfer option dicts.
	d.dicts = make(map[uint32]dict, len(d.o.dicts))
	for _, dc := range d.o.dicts {
		d.dicts[dc.id] = dc
	}
	d.o.dicts = nil

	// Create decoders
	d.decoders = make(chan *blockDec, d.o.concurrent)
	for i := 0; i < d.o.concurrent; i++ {
		dec := newBlockDec(d.o.lowMem)
		dec.localFrame = newFrameDec(d.o)
		d.decoders <- dec
	}

	if r == nil {
		return &d, nil
	}
	return &d, d.Reset(r)
}

// Read bytes from the decompressed stream into p.
// Returns the number of bytes written and any error that occurred.
// When the stream is done, io.EOF will be returned.
func (d *Decoder) Read(p []byte) (int, error) {
	if d.stream == nil {
		return 0, errors.New("no input has been initialized")
	}
	var n int
	for {
		if len(d.current.b) > 0 {
			filled := copy(p, d.current.b)
			p = p[filled:]
			d.current.b = d.current.b[filled:]
			n += filled
		}
		if len(p) == 0 {
			break
		}
		if len(d.current.b) == 0 {
			// We have an error and no more data
			if d.current.err != nil {
				break
			}
			if !d.nextBlock(n == 0) {
				return n, nil
			}
		}
	}
	if len(d.current.b) > 0 {
		if debug {
			println("returning", n, "still bytes left:", len(d.current.b))
		}
		// Only return error at end of block
		return n, nil
	}
	if d.current.err != nil {
		d.drainOutput()
	}
	if debug {
		println("returning", n, d.current.err, len(d.decoders))
	}
	return n, d.current.err
}

// Reset will reset the decoder the supplied stream after the current has finished processing.
// Note that this functionality cannot be used after Close has been called.
func (d *Decoder) Reset(r io.Reader) error {
	if d.current.err == ErrDecoderClosed {
		return d.current.err
	}
	if r == nil {
		return errors.New("nil Reader sent as input")
	}

	if d.stream == nil {
		d.stream = make(chan decodeStream, 1)
		d.streamWg.Add(1)
		go d.startStreamDecoder(d.stream)
	}

	d.drainOutput()

	// If bytes buffer and < 1MB, do sync decoding anyway.
	if bb, ok := r.(*bytes.Buffer); ok && bb.Len() < 1<<20 {
		if debug {
			println("*bytes.Buffer detected, doing sync decode, len:", bb.Len())
		}
		b := bb.Bytes()
		var dst []byte
		if cap(d.current.b) > 0 {
			dst = d.current.b
		}

		dst, err := d.DecodeAll(b, dst[:0])
		if err == nil {
			err = io.EOF
		}
		d.current.b = dst
		d.current.err = err
		d.current.flushed = true
		if debug {
			println("sync decode to", len(dst), "bytes, err:", err)
		}
		return nil
	}

	// Remove current block.
	d.current.decodeOutput = decodeOutput{}
	d.current.err = nil
	d.current.cancel = make(chan struct{})
	d.current.flushed = false
	d.current.d = nil

	d.stream <- decodeStream{
		r:      r,
		output: d.current.output,
		cancel: d.current.cancel,
	}
	return nil
}

// drainOutput will drain the output until errEndOfStream is sent.
func (d *Decoder) drainOutput() {
	if d.current.cancel != nil {
		println("cancelling current")
		close(d.current.cancel)
		d.current.cancel = nil
	}
	if d.current.d != nil {
		if debug {
			printf("re-adding current decoder %p, decoders: %d", d.current.d, len(d.decoders))
		}
		d.decoders <- d.current.d
		d.current.d = nil
		d.current.b = nil
	}
	if d.current.output == nil || d.current.flushed {
		println("current already flushed")
		return
	}
	for {
		select {
		case v := <-d.current.output:
			if v.d != nil {
				if debug {
					printf("re-adding decoder %p", v.d)
				}
				d.decoders <- v.d
			}
			if v.err == errEndOfStream {
				println("current flushed")
				d.current.flushed = true
				return
			}
		}
	}
}

// WriteTo writes data to w until there's no more data to write or when an error occurs.
// The return value n is the number of bytes written.
// Any error encountered during the write is also returned.
func (d *Decoder) WriteTo(w io.Writer) (int64, error) {
	if d.stream == nil {
		return 0, errors.New("no input has been initialized")
	}
	var n int64
	for {
		if len(d.current.b) > 0 {
			n2, err2 := w.Write(d.current.b)
			n += int64(n2)
			if err2 != nil && d.current.err == nil {
				d.current.err = err2
				break
			}
		}
		if d.current.err != nil {
			break
		}
		d.nextBlock(true)
	}
	err := d.current.err
	if err != nil {
		d.drainOutput()
	}
	if err == io.EOF {
		err = nil
	}
	return n, err
}

// DecodeAll allows stateless decoding of a blob of bytes.
// Output will be appended to dst, so if the destination size is known
// you can pre-allocate the destination slice to avoid allocations.
// DecodeAll can be used concurrently.
// The Decoder concurrency limits will be respected.
func (d *Decoder) DecodeAll(input, dst []byte) ([]byte, error) {
	if d.current.err == ErrDecoderClosed {
		return dst, ErrDecoderClosed
	}

	// Grab a block decoder and frame decoder.
	block := <-d.decoders
	frame := block.localFrame
	defer func() {
		if debug {
			printf("re-adding decoder: %p", block)
		}
		frame.rawInput = nil
		frame.bBuf = nil
		d.decoders <- block
	}()
	frame.bBuf = input

	for {
		frame.history.reset()
		err := frame.reset(&frame.bBuf)
		if err == io.EOF {
			if debug {
				println("frame reset return EOF")
			}
			return dst, nil
		}
		if frame.DictionaryID != nil {
			dict, ok := d.dicts[*frame.DictionaryID]
			if !ok {
				return nil, ErrUnknownDictionary
			}
			frame.history.setDict(&dict)
		}
		if err != nil {
			return dst, err
		}
		if frame.FrameContentSize > d.o.maxDecodedSize-uint64(len(dst)) {
			return dst, ErrDecoderSizeExceeded
		}
		if frame.FrameContentSize > 0 && frame.FrameContentSize < 1<<30 {
			// Never preallocate moe than 1 GB up front.
			if cap(dst)-len(dst) < int(frame.FrameContentSize) {
				dst2 := make([]byte, len(dst), len(dst)+int(frame.FrameContentSize))
				copy(dst2, dst)
				dst = dst2
			}
		}
		if cap(dst) == 0 {
			// Allocate len(input) * 2 by default if nothing is provided
			// and we didn't get frame content size.
			size := len(input) * 2
			// Cap to 1 MB.
			if size > 1<<20 {
				size = 1 << 20
			}
			if uint64(size) > d.o.maxDecodedSize {
				size = int(d.o.maxDecodedSize)
			}
			dst = make([]byte, 0, size)
		}

		dst, err = frame.runDecoder(dst, block)
		if err != nil {
			return dst, err
		}
		if len(frame.bBuf) == 0 {
			if debug {
				println("frame dbuf empty")
			}
			break
		}
	}
	return dst, nil
}

// nextBlock returns the next block.
// If an error occurs d.err will be set.
// Optionally the function can block for new output.
// If non-blocking mode is used the returned boolean will be false
// if no data was available without blocking.
func (d *Decoder) nextBlock(blocking bool) (ok bool) {
	if d.current.d != nil {
		if debug {
			printf("re-adding current decoder %p", d.current.d)
		}
		d.decoders <- d.current.d
		d.current.d = nil
	}
	if d.current.err != nil {
		// Keep error state.
		return blocking
	}

	if blocking {
		d.current.decodeOutput = <-d.current.output
	} else {
		select {
		case d.current.decodeOutput = <-d.current.output:
		default:
			return false
		}
	}
	if debug {
		println("got", len(d.current.b), "bytes, error:", d.current.err)
	}
	return true
}

// Close will release all resources.
// It is NOT possible to reuse the decoder after this.
func (d *Decoder) Close() {
	if d.current.err == ErrDecoderClosed {
		return
	}
	d.drainOutput()
	if d.stream != nil {
		close(d.stream)
		d.streamWg.Wait()
		d.stream = nil
	}
	if d.decoders != nil {
		close(d.decoders)
		for dec := range d.decoders {
			dec.Close()
		}
		d.decoders = nil
	}
	if d.current.d != nil {
		d.current.d.Close()
		d.current.d = nil
	}
	d.current.err = ErrDecoderClosed
}

// IOReadCloser returns the decoder as an io.ReadCloser for convenience.
// Any changes to the decoder will be reflected, so the returned ReadCloser
// can be reused along with the decoder.
// io.WriterTo is also supported by the returned ReadCloser.
func (d *Decoder) IOReadCloser() io.ReadCloser {
	return closeWrapper{d: d}
}

// closeWrapper wraps a function call as a closer.
type closeWrapper struct {
	d *Decoder
}

// WriteTo forwards WriteTo calls to the decoder.
func (c closeWrapper) WriteTo(w io.Writer) (n int64, err error) {
	return c.d.WriteTo(w)
}

// Read forwards read calls to the decoder.
func (c closeWrapper) Read(p []byte) (n int, err error) {
	return c.d.Read(p)
}

// Close closes the decoder.
func (c closeWrapper) Close() error {
	c.d.Close()
	return nil
}

type decodeOutput struct {
	d   *blockDec
	b   []byte
	err error
}

type decodeStream struct {
	r io.Reader

	// Blocks ready to be written to output.
	output chan decodeOutput

	// cancel reading from the input
	cancel chan struct{}
}

// errEndOfStream indicates that everything from the stream was read.
var errEndOfStream = errors.New("end-of-stream")

// Create Decoder:
// Spawn n block decoders. These accept tasks to decode a block.
// Create goroutine that handles stream processing, this will send history to decoders as they are available.
// Decoders update the history as they decode.
// When a block is returned:
// 		a) history is sent to the next decoder,
// 		b) content written to CRC.
// 		c) return data to WRITER.
// 		d) wait for next block to return data.
// Once WRITTEN, the decoders reused by the writer frame decoder for re-use.
func (d *Decoder) startStreamDecoder(inStream chan decodeStream) {
	defer d.streamWg.Done()
	frame := newFrameDec(d.o)
	for stream := range inStream {
		if debug {
			println("got new stream")
		}
		br := readerWrapper{r: stream.r}
	decodeStream:
		for {
			frame.history.reset()
			err := frame.reset(&br)
			if debug && err != nil {
				println("Frame decoder returned", err)
			}
			if err == nil && frame.DictionaryID != nil {
				dict, ok := d.dicts[*frame.DictionaryID]
				if !ok {
					err = ErrUnknownDictionary
				} else {
					frame.history.setDict(&dict)
				}
			}
			if err != nil {
				stream.output <- decodeOutput{
					err: err,
				}
				break
			}
			if debug {
				println("starting frame decoder")
			}

			// This goroutine will forward history between frames.
			frame.frameDone.Add(1)
			frame.initAsync()

			go frame.startDecoder(stream.output)
		decodeFrame:
			// Go through all blocks of the frame.
			for {
				dec := <-d.decoders
				select {
				case <-stream.cancel:
					if !frame.sendErr(dec, io.EOF) {
						// To not let the decoder dangle, send it back.
						stream.output <- decodeOutput{d: dec}
					}
					break decodeStream
				default:
				}
				err := frame.next(dec)
				switch err {
				case io.EOF:
					// End of current frame, no error
					println("EOF on next block")
					break decodeFrame
				case nil:
					continue
				default:
					println("block decoder returned", err)
					break decodeStream
				}
			}
			// All blocks have started decoding, check if there are more frames.
			println("waiting for done")
			frame.frameDone.Wait()
			println("done waiting...")
		}
		frame.frameDone.Wait()
		println("Sending EOS")
		stream.output <- decodeOutput{err: errEndOfStream}
	}
}