148 lines
4.1 KiB
Go
148 lines
4.1 KiB
Go
package zstd
|
|
|
|
/*
|
|
#define ZSTD_STATIC_LINKING_ONLY
|
|
#include "zstd.h"
|
|
#include "stdint.h" // for uintptr_t
|
|
|
|
// The following *_wrapper function are used for removing superflouos
|
|
// memory allocations when calling the wrapped functions from Go code.
|
|
// See https://github.com/golang/go/issues/24450 for details.
|
|
|
|
static size_t ZSTD_compress_wrapper(uintptr_t dst, size_t maxDstSize, const uintptr_t src, size_t srcSize, int compressionLevel) {
|
|
return ZSTD_compress((void*)dst, maxDstSize, (const void*)src, srcSize, compressionLevel);
|
|
}
|
|
|
|
static size_t ZSTD_decompress_wrapper(uintptr_t dst, size_t maxDstSize, uintptr_t src, size_t srcSize) {
|
|
return ZSTD_decompress((void*)dst, maxDstSize, (const void *)src, srcSize);
|
|
}
|
|
|
|
*/
|
|
import "C"
|
|
import (
|
|
"bytes"
|
|
"errors"
|
|
"io/ioutil"
|
|
"runtime"
|
|
"unsafe"
|
|
)
|
|
|
|
// Defines best and standard values for zstd cli
|
|
const (
|
|
BestSpeed = 1
|
|
BestCompression = 20
|
|
DefaultCompression = 5
|
|
)
|
|
|
|
var (
|
|
// ErrEmptySlice is returned when there is nothing to compress
|
|
ErrEmptySlice = errors.New("Bytes slice is empty")
|
|
)
|
|
|
|
// CompressBound returns the worst case size needed for a destination buffer,
|
|
// which can be used to preallocate a destination buffer or select a previously
|
|
// allocated buffer from a pool.
|
|
// See zstd.h to mirror implementation of ZSTD_COMPRESSBOUND
|
|
func CompressBound(srcSize int) int {
|
|
lowLimit := 128 << 10 // 128 kB
|
|
var margin int
|
|
if srcSize < lowLimit {
|
|
margin = (lowLimit - srcSize) >> 11
|
|
}
|
|
return srcSize + (srcSize >> 8) + margin
|
|
}
|
|
|
|
// cCompressBound is a cgo call to check the go implementation above against the c code.
|
|
func cCompressBound(srcSize int) int {
|
|
return int(C.ZSTD_compressBound(C.size_t(srcSize)))
|
|
}
|
|
|
|
// Compress src into dst. If you have a buffer to use, you can pass it to
|
|
// prevent allocation. If it is too small, or if nil is passed, a new buffer
|
|
// will be allocated and returned.
|
|
func Compress(dst, src []byte) ([]byte, error) {
|
|
return CompressLevel(dst, src, DefaultCompression)
|
|
}
|
|
|
|
// CompressLevel is the same as Compress but you can pass a compression level
|
|
func CompressLevel(dst, src []byte, level int) ([]byte, error) {
|
|
bound := CompressBound(len(src))
|
|
if cap(dst) >= bound {
|
|
dst = dst[0:bound] // Reuse dst buffer
|
|
} else {
|
|
dst = make([]byte, bound)
|
|
}
|
|
|
|
srcPtr := C.uintptr_t(uintptr(0)) // Do not point anywhere, if src is empty
|
|
if len(src) > 0 {
|
|
srcPtr = C.uintptr_t(uintptr(unsafe.Pointer(&src[0])))
|
|
}
|
|
|
|
cWritten := C.ZSTD_compress_wrapper(
|
|
C.uintptr_t(uintptr(unsafe.Pointer(&dst[0]))),
|
|
C.size_t(len(dst)),
|
|
srcPtr,
|
|
C.size_t(len(src)),
|
|
C.int(level))
|
|
|
|
runtime.KeepAlive(src)
|
|
written := int(cWritten)
|
|
// Check if the return is an Error code
|
|
if err := getError(written); err != nil {
|
|
return nil, err
|
|
}
|
|
return dst[:written], nil
|
|
}
|
|
|
|
// Decompress src into dst. If you have a buffer to use, you can pass it to
|
|
// prevent allocation. If it is too small, or if nil is passed, a new buffer
|
|
// will be allocated and returned.
|
|
func Decompress(dst, src []byte) ([]byte, error) {
|
|
if len(src) == 0 {
|
|
return []byte{}, ErrEmptySlice
|
|
}
|
|
decompress := func(dst, src []byte) ([]byte, error) {
|
|
|
|
cWritten := C.ZSTD_decompress_wrapper(
|
|
C.uintptr_t(uintptr(unsafe.Pointer(&dst[0]))),
|
|
C.size_t(len(dst)),
|
|
C.uintptr_t(uintptr(unsafe.Pointer(&src[0]))),
|
|
C.size_t(len(src)))
|
|
|
|
runtime.KeepAlive(src)
|
|
written := int(cWritten)
|
|
// Check error
|
|
if err := getError(written); err != nil {
|
|
return nil, err
|
|
}
|
|
return dst[:written], nil
|
|
}
|
|
|
|
if len(dst) == 0 {
|
|
// Attempt to use zStd to determine decompressed size (may result in error or 0)
|
|
size := int(C.size_t(C.ZSTD_getDecompressedSize(unsafe.Pointer(&src[0]), C.size_t(len(src)))))
|
|
|
|
if err := getError(size); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if size > 0 {
|
|
dst = make([]byte, size)
|
|
} else {
|
|
dst = make([]byte, len(src)*3) // starting guess
|
|
}
|
|
}
|
|
for i := 0; i < 3; i++ { // 3 tries to allocate a bigger buffer
|
|
result, err := decompress(dst, src)
|
|
if !IsDstSizeTooSmallError(err) {
|
|
return result, err
|
|
}
|
|
dst = make([]byte, len(dst)*2) // Grow buffer by 2
|
|
}
|
|
|
|
// We failed getting a dst buffer of correct size, use stream API
|
|
r := NewReader(bytes.NewReader(src))
|
|
defer r.Close()
|
|
return ioutil.ReadAll(r)
|
|
}
|