diff options
Diffstat (limited to 'go/cbrotli/writer.go')
-rwxr-xr-x | go/cbrotli/writer.go | 160 |
1 files changed, 160 insertions, 0 deletions
diff --git a/go/cbrotli/writer.go b/go/cbrotli/writer.go new file mode 100755 index 0000000..279a2f2 --- /dev/null +++ b/go/cbrotli/writer.go @@ -0,0 +1,160 @@ +// Copyright 2016 Google Inc. All Rights Reserved. +// +// Distributed under MIT license. +// See file LICENSE for detail or copy at https://opensource.org/licenses/MIT + +package cbrotli + +/* +#include <stdbool.h> +#include <stddef.h> +#include <stdint.h> + +#include <brotli/encode.h> + +struct CompressStreamResult { + size_t bytes_consumed; + const uint8_t* output_data; + size_t output_data_size; + int success; + int has_more; +}; + +static struct CompressStreamResult CompressStream( + BrotliEncoderState* s, BrotliEncoderOperation op, + const uint8_t* data, size_t data_size) { + struct CompressStreamResult result; + size_t available_in = data_size; + const uint8_t* next_in = data; + size_t available_out = 0; + result.success = BrotliEncoderCompressStream(s, op, + &available_in, &next_in, &available_out, 0, 0) ? 1 : 0; + result.bytes_consumed = data_size - available_in; + result.output_data = 0; + result.output_data_size = 0; + if (result.success) { + result.output_data = BrotliEncoderTakeOutput(s, &result.output_data_size); + } + result.has_more = BrotliEncoderHasMoreOutput(s) ? 1 : 0; + return result; +} +*/ +import "C" + +import ( + "bytes" + "errors" + "io" + "unsafe" +) + +// WriterOptions configures Writer. +type WriterOptions struct { + // Quality controls the compression-speed vs compression-density trade-offs. + // The higher the quality, the slower the compression. Range is 0 to 11. + Quality int + // LGWin is the base 2 logarithm of the sliding window size. + // Range is 10 to 24. 0 indicates automatic configuration based on Quality. + LGWin int + // BufferSize is the number of bytes to use to buffer encoded output. + // 0 indicates an implementation-defined default. + BufferSize int +} + +// Writer implements io.WriteCloser by writing Brotli-encoded data to an +// underlying Writer. +type Writer struct { + dst io.Writer + state *C.BrotliEncoderState + buf, encoded []byte +} + +var ( + errEncode = errors.New("cbrotli: encode error") + errWriterClosed = errors.New("cbrotli: Writer is closed") +) + +// NewWriter initializes new Writer instance. +// Close MUST be called to free resources. +func NewWriter(dst io.Writer, options WriterOptions) *Writer { + state := C.BrotliEncoderCreateInstance(nil, nil, nil) + C.BrotliEncoderSetParameter( + state, C.BROTLI_PARAM_QUALITY, (C.uint32_t)(options.Quality)) + C.BrotliEncoderSetParameter( + state, C.BROTLI_PARAM_LGWIN, (C.uint32_t)(options.LGWin)) + return &Writer{ + dst: dst, + state: state, + } +} + +func (w *Writer) writeChunk(p []byte, op C.BrotliEncoderOperation) (n int, err error) { + if w.state == nil { + return 0, errWriterClosed + } + + for { + var data *C.uint8_t + if len(p) != 0 { + data = (*C.uint8_t)(&p[0]) + } + result := C.CompressStream(w.state, op, data, C.size_t(len(p))) + if result.success == 0 { + return n, errEncode + } + p = p[int(result.bytes_consumed):] + n += int(result.bytes_consumed) + + length := int(result.output_data_size) + if length != 0 { + // It is a workaround for non-copying-wrapping of native memory. + // C-encoder never pushes output block longer than ((2 << 25) + 502). + // TODO: use natural wrapper, when it becomes available, see + // https://golang.org/issue/13656. + output := (*[1 << 30]byte)(unsafe.Pointer(result.output_data))[:length:length] + _, err = w.dst.Write(output) + if err != nil { + return n, err + } + } + if len(p) == 0 && result.has_more == 0 { + return n, nil + } + } +} + +// Flush outputs encoded data for all input provided to Write. The resulting +// output can be decoded to match all input before Flush, but the stream is +// not yet complete until after Close. +// Flush has a negative impact on compression. +func (w *Writer) Flush() error { + _, err := w.writeChunk(nil, C.BROTLI_OPERATION_FLUSH) + return err +} + +// Close flushes remaining data to the decorated writer and frees C resources. +func (w *Writer) Close() error { + // If stream is already closed, it is reported by `writeChunk`. + _, err := w.writeChunk(nil, C.BROTLI_OPERATION_FINISH) + // C-Brotli tolerates `nil` pointer here. + C.BrotliEncoderDestroyInstance(w.state) + w.state = nil + return err +} + +// Write implements io.Writer. Flush or Close must be called to ensure that the +// encoded bytes are actually flushed to the underlying Writer. +func (w *Writer) Write(p []byte) (n int, err error) { + return w.writeChunk(p, C.BROTLI_OPERATION_PROCESS) +} + +// Encode returns content encoded with Brotli. +func Encode(content []byte, options WriterOptions) ([]byte, error) { + var buf bytes.Buffer + writer := NewWriter(&buf, options) + _, err := writer.Write(content) + if closeErr := writer.Close(); err == nil { + err = closeErr + } + return buf.Bytes(), err +} |