Edit

kc3-lang/brotli/go/cbrotli/writer.go

Branch :

  • Show log

    Commit

  • Author : Evgenii Kliuchnikov
    Date : 2024-11-21 05:38:27
    Hash : 5c667240
    Message : Add dictionary API to cgo wrapper PiperOrigin-RevId: 698745795

  • go/cbrotli/writer.go
  • // 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"
    	"runtime"
    	"unsafe"
    )
    
    // PreparedDictionary is a handle to native object.
    type PreparedDictionary struct {
    	opaque *C.BrotliEncoderPreparedDictionary
    	pinner *runtime.Pinner
    }
    
    // DictionaryType is type for shared dictionary
    type DictionaryType int
    
    const (
    	// DtRaw denotes LZ77 prefix dictionary
    	DtRaw DictionaryType = 0
    	// DtSerialized denotes serialized format
    	DtSerialized DictionaryType = 1
    )
    
    // NewPreparedDictionary prepares dictionary data for encoder.
    // Same instance can be used for multiple encoding sessions.
    // Close MUST be called to free resources.
    func NewPreparedDictionary(data []byte, dictionaryType DictionaryType, quality int) *PreparedDictionary {
    	var ptr *C.uint8_t
    	if len(data) != 0 {
    		ptr = (*C.uint8_t)(&data[0])
    	}
    	p := new(runtime.Pinner)
    	p.Pin(&data[0])
    	d := C.BrotliEncoderPrepareDictionary(C.BrotliSharedDictionaryType(dictionaryType), C.size_t(len(data)), ptr, C.int(quality), nil, nil, nil)
    	return &PreparedDictionary{
    		opaque: d,
    		pinner: p,
    	}
    }
    
    // Close frees C resources.
    // IMPORTANT: calling Close until all encoders that use that dictionary are closed as well will
    // cause crash.
    func (p *PreparedDictionary) Close() error {
    	// C-Brotli tolerates `nil` pointer here.
    	C.BrotliEncoderDestroyPreparedDictionary(p.opaque)
    	p.opaque = nil
    	p.pinner.Unpin()
    	return nil
    }
    
    // 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
    	// Prepared shared dictionary
    	Dictionary *PreparedDictionary
    }
    
    // Writer implements io.WriteCloser by writing Brotli-encoded data to an
    // underlying Writer.
    type Writer struct {
    	healthy      bool
    	dst          io.Writer
    	state        *C.BrotliEncoderState
    	buf, encoded []byte
    }
    
    var (
    	errEncode          = errors.New("cbrotli: encode error")
    	errWriterClosed    = errors.New("cbrotli: Writer is closed")
    	errWriterUnhealthy = errors.New("cbrotli: Writer is unhealthy")
    )
    
    // 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)
    	healthy := state != nil
    	if C.BrotliEncoderSetParameter(
    		state, C.BROTLI_PARAM_QUALITY, (C.uint32_t)(options.Quality)) == 0 {
    		healthy = false
    	}
    	if options.LGWin > 0 {
    		if C.BrotliEncoderSetParameter(
    			state, C.BROTLI_PARAM_LGWIN, (C.uint32_t)(options.LGWin)) == 0 {
    			healthy = false
    		}
    	}
    	if options.Dictionary != nil {
    		if C.BrotliEncoderAttachPreparedDictionary(state, options.Dictionary.opaque) == 0 {
    			healthy = false
    		}
    	}
    	return &Writer{
    		healthy: healthy,
    		dst:     dst,
    		state:   state,
    	}
    }
    
    func (w *Writer) writeChunk(p []byte, op C.BrotliEncoderOperation) (n int, err error) {
    	if !w.healthy {
    		return 0, errWriterUnhealthy
    	}
    	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(eustas): 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
    }