pkg/encoding/bytes.go - skywalking-banyandb - Git at Google

 // Licensed to Apache Software Foundation (ASF) under one or more contributor
 // license agreements. See the NOTICE file distributed with
 // this work for additional information regarding copyright
 // ownership. Apache Software Foundation (ASF) licenses this file to you under
 // the Apache License, Version 2.0 (the "License"); you may
 // not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //     http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing,
 // software distributed under the License is distributed on an
 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.

 package encoding

 import (
 	"fmt"

 	"github.com/apache/skywalking-banyandb/pkg/bytes"
 	"github.com/apache/skywalking-banyandb/pkg/compress/zstd"
 )

 // EncodeBytes encodes a string into dst.
 func EncodeBytes(dst, b []byte) []byte {
 	dst = VarUint64ToBytes(dst, uint64(len(b)))
 	dst = append(dst, b...)
 	return dst
 }

 // DecodeBytes decodes a string from src.
 func DecodeBytes(src []byte) ([]byte, []byte, error) {
 	var n uint64
 	src, n = BytesToVarUint64(src)
 	if uint64(len(src)) < n {
 		return nil, nil, fmt.Errorf("src is too short for reading string with size %d; len(src)=%d", n, len(src))
 	}
 	return src[n:], src[:n], nil
 }

 // EncodeBytesBlock encodes a block of strings into dst.
 func EncodeBytesBlock(dst []byte, a [][]byte) []byte {
 	u64s := GenerateUint64List(len(a))
 	aLens := u64s.L[:0]
 	for _, s := range a {
 		if s == nil {
 			aLens = append(aLens, 0)
 		} else {
 			// Offset by 1: "" (len 0) becomes 1, "a" (len 1) becomes 2, etc.
 			aLens = append(aLens, uint64(len(s))+1)
 		}
 	}
 	u64s.L = aLens
 	dst = EncodeUint64Block(dst, u64s.L)
 	ReleaseUint64List(u64s)

 	bb := bbPool.Generate()
 	b := bb.Buf
 	for _, s := range a {
 		b = append(b, s...)
 	}
 	bb.Buf = b
 	dst = compressBlock(dst, bb.Buf)
 	bbPool.Release(bb)

 	return dst
 }

 // BytesBlockDecoder decodes a block of strings from src.
 // It reuses the underlying buffer for storing the decoded strings.
 type BytesBlockDecoder struct {
 	data []byte
 }

 // Reset resets the decoder.
 func (bbd *BytesBlockDecoder) Reset() {
 	bbd.data = bbd.data[:0]
 }

 // Decode decodes a block of strings from src.
 func (bbd *BytesBlockDecoder) Decode(dst [][]byte, src []byte, itemsCount uint64) ([][]byte, error) {
 	u64List := GenerateUint64List(0)
 	defer ReleaseUint64List(u64List)

 	var tail []byte
 	var err error
 	u64List.L, tail, err = DecodeUint64Block(u64List.L[:0], src, itemsCount)
 	if err != nil {
 		return dst, fmt.Errorf("cannot decode string lengths: %w", err)
 	}
 	aLens := u64List.L
 	src = tail

 	dataLen := len(bbd.data)
 	bbd.data, tail, err = decompressBlock(bbd.data, src)
 	if err != nil {
 		return dst, fmt.Errorf("cannot decode bytes block with strings: %w", err)
 	}
 	if len(tail) > 0 {
 		return dst, fmt.Errorf("unexpected non-empty tail after reading bytes block with strings; len(tail)=%d", len(tail))
 	}

 	data := bbd.data[dataLen:]
 	for _, sLen := range aLens {
 		if sLen == 0 {
 			dst = append(dst, nil)
 			continue
 		}
 		// Reverse the offset
 		actualLen := sLen - 1
 		if uint64(len(data)) < actualLen {
 			return dst, fmt.Errorf("cannot decode a string with the length %d bytes from %d bytes", actualLen, len(data))
 		}
 		if actualLen == 0 {
 			// Explicitly create non-nil empty slice for empty strings
 			dst = append(dst, []byte{})
 		} else {
 			dst = append(dst, data[:actualLen])
 		}
 		data = data[actualLen:]
 	}

 	return dst, nil
 }

 // DecodeWithTail decodes a block of strings from src and returns the remaining tail.
 func (bbd *BytesBlockDecoder) DecodeWithTail(dst [][]byte, src []byte, itemsCount uint64) ([][]byte, []byte, error) {
 	u64List := GenerateUint64List(0)
 	defer ReleaseUint64List(u64List)

 	var tail []byte
 	var err error
 	u64List.L, tail, err = DecodeUint64Block(u64List.L[:0], src, itemsCount)
 	if err != nil {
 		return dst, nil, fmt.Errorf("cannot decode string lengths: %w", err)
 	}
 	aLens := u64List.L
 	src = tail

 	dataLen := len(bbd.data)
 	bbd.data, tail, err = decompressBlock(bbd.data, src)
 	if err != nil {
 		return dst, tail, fmt.Errorf("cannot decode bytes block with strings: %w", err)
 	}

 	data := bbd.data[dataLen:]
 	for _, sLen := range aLens {
 		if sLen == 0 {
 			dst = append(dst, nil)
 			continue
 		}
 		// Reverse the offset
 		actualLen := sLen - 1
 		if uint64(len(data)) < actualLen {
 			return dst, tail, fmt.Errorf("cannot decode a string with the length %d bytes from %d bytes", actualLen, len(data))
 		}
 		if actualLen == 0 {
 			// Explicitly create non-nil empty slice for empty strings
 			dst = append(dst, []byte{})
 		} else {
 			dst = append(dst, data[:actualLen])
 		}
 		data = data[actualLen:]
 	}

 	return dst, tail, nil
 }

 // EncodeUint64Block encodes a block of uint64 values into dst.
 func EncodeUint64Block(dst []byte, a []uint64) []byte {
 	bb := bbPool.Generate()
 	bb.Buf = encodeUint64List(bb.Buf[:0], a)
 	dst = compressBlock(dst, bb.Buf)
 	bbPool.Release(bb)
 	return dst
 }

 // DecodeUint64Block decodes a block of uint64 values from src.
 func DecodeUint64Block(dst []uint64, src []byte, itemsCount uint64) ([]uint64, []byte, error) {
 	bb := bbPool.Generate()
 	defer bbPool.Release(bb)

 	var err error
 	bb.Buf, src, err = decompressBlock(bb.Buf[:0], src)
 	if err != nil {
 		return dst, src, fmt.Errorf("cannot decode bytes block: %w", err)
 	}

 	dst, err = decodeUint64List(dst, bb.Buf, itemsCount)
 	if err != nil {
 		return dst, src, fmt.Errorf("cannot decode %d uint64 items from bytes block of length %d bytes: %w", itemsCount, len(bb.Buf), err)
 	}
 	return dst, src, nil
 }

 const (
 	uintBlockType8  = 0
 	uintBlockType16 = 1
 	uintBlockType32 = 2
 	uintBlockType64 = 3
 )

 func encodeUint64List(dst []byte, a []uint64) []byte {
 	nMax := uint64(0)
 	for _, n := range a {
 		if n > nMax {
 			nMax = n
 		}
 	}
 	switch {
 	case nMax < (1 << 8):
 		dst = append(dst, uintBlockType8)
 		for _, n := range a {
 			dst = append(dst, byte(n))
 		}
 	case nMax < (1 << 16):
 		dst = append(dst, uintBlockType16)
 		for _, n := range a {
 			dst = Uint16ToBytes(dst, uint16(n))
 		}
 	case nMax < (1 << 32):
 		dst = append(dst, uintBlockType32)
 		for _, n := range a {
 			dst = Uint32ToBytes(dst, uint32(n))
 		}
 	default:
 		dst = append(dst, uintBlockType64)
 		for _, n := range a {
 			dst = Uint64ToBytes(dst, n)
 		}
 	}
 	return dst
 }

 func decodeUint64List(dst []uint64, src []byte, itemsCount uint64) ([]uint64, error) {
 	if len(src) < 1 {
 		return dst, fmt.Errorf("cannot decode uint64 block type from empty src")
 	}
 	blockType := src[0]
 	src = src[1:]

 	switch blockType {
 	case uintBlockType8:
 		if uint64(len(src)) != itemsCount {
 			return dst, fmt.Errorf("unexpected block length for %d items; got %d bytes; want %d bytes", itemsCount, len(src), itemsCount)
 		}
 		for _, v := range src {
 			dst = append(dst, uint64(v))
 		}
 	case uintBlockType16:
 		if uint64(len(src)) != 2*itemsCount {
 			return dst, fmt.Errorf("unexpected block length for %d items; got %d bytes; want %d bytes", itemsCount, len(src), 2*itemsCount)
 		}
 		for len(src) > 0 {
 			v := BytesToUint16(src)
 			src = src[2:]
 			dst = append(dst, uint64(v))
 		}
 	case uintBlockType32:
 		if uint64(len(src)) != 4*itemsCount {
 			return dst, fmt.Errorf("unexpected block length for %d items; got %d bytes; want %d bytes", itemsCount, len(src), 4*itemsCount)
 		}
 		for len(src) > 0 {
 			v := BytesToUint32(src)
 			src = src[4:]
 			dst = append(dst, uint64(v))
 		}
 	case uintBlockType64:
 		if uint64(len(src)) != 8*itemsCount {
 			return dst, fmt.Errorf("unexpected block length for %d items; got %d bytes; want %d bytes", itemsCount, len(src), 8*itemsCount)
 		}
 		for len(src) > 0 {
 			v := BytesToUint64(src)
 			src = src[8:]
 			dst = append(dst, v)
 		}
 	default:
 		return dst, fmt.Errorf("unexpected uint64 block type: %d; want 0, 1, 2 or 3", blockType)
 	}
 	return dst, nil
 }

 const (
 	compressTypePlain = 0
 	compressTypeZSTD  = 1
 )

 func compressBlock(dst, src []byte) []byte {
 	if len(src) < 128 {
 		dst = append(dst, compressTypePlain, byte(len(src)))
 		return append(dst, src...)
 	}

 	dst = append(dst, compressTypeZSTD)
 	bb := bbPool.Generate()
 	bb.Buf = zstd.Compress(bb.Buf[:0], src, 1)
 	dst = VarUint64ToBytes(dst, uint64(len(bb.Buf)))
 	dst = append(dst, bb.Buf...)
 	bbPool.Release(bb)
 	return dst
 }

 func decompressBlock(dst, src []byte) ([]byte, []byte, error) {
 	if len(src) < 1 {
 		return dst, src, fmt.Errorf("cannot decode block type from empty src")
 	}
 	blockType := src[0]
 	src = src[1:]
 	switch blockType {
 	case compressTypePlain:
 		if len(src) < 1 {
 			return dst, src, fmt.Errorf("cannot decode plain block size from empty src")
 		}
 		blockLen := int(src[0])
 		src = src[1:]
 		if len(src) < blockLen {
 			return dst, src, fmt.Errorf("cannot read plain block with the size %d bytes from %b bytes", blockLen, len(src))
 		}

 		dst = append(dst, src[:blockLen]...)
 		src = src[blockLen:]
 		return dst, src, nil
 	case compressTypeZSTD:
 		tail, blockLen := BytesToVarUint64(src)
 		src = tail
 		if uint64(len(src)) < blockLen {
 			return dst, src, fmt.Errorf("cannot read compressed block with the size %d bytes from %d bytes", blockLen, len(src))
 		}
 		compressedBlock := src[:blockLen]
 		src = src[blockLen:]

 		// Decompress the block
 		var err error
 		bb := bbPool.Generate()
 		bb.Buf, err = zstd.Decompress(bb.Buf[:0], compressedBlock)
 		if err != nil {
 			return dst, src, fmt.Errorf("cannot decompress block: %w", err)
 		}

 		// Copy the decompressed block to dst.
 		dst = append(dst, bb.Buf...)
 		bbPool.Release(bb)
 		return dst, src, nil
 	default:
 		return dst, src, fmt.Errorf("unexpected block type: %d; supported types: 0, 1", blockType)
 	}
 }

 var bbPool = bytes.NewBufferPool("encoding.bytesBlock")
	// Licensed to Apache Software Foundation (ASF) under one or more contributor
	// license agreements. See the NOTICE file distributed with
	// this work for additional information regarding copyright
	// ownership. Apache Software Foundation (ASF) licenses this file to you under
	// the Apache License, Version 2.0 (the "License"); you may
	// not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing,
	// software distributed under the License is distributed on an
	// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	// KIND, either express or implied. See the License for the
	// specific language governing permissions and limitations
	// under the License.

	package encoding

	import (
	"fmt"

	"github.com/apache/skywalking-banyandb/pkg/bytes"
	"github.com/apache/skywalking-banyandb/pkg/compress/zstd"
	)

	// EncodeBytes encodes a string into dst.
	func EncodeBytes(dst, b []byte) []byte {
	dst = VarUint64ToBytes(dst, uint64(len(b)))
	dst = append(dst, b...)
	return dst
	}

	// DecodeBytes decodes a string from src.
	func DecodeBytes(src []byte) ([]byte, []byte, error) {
	var n uint64
	src, n = BytesToVarUint64(src)
	if uint64(len(src)) < n {
	return nil, nil, fmt.Errorf("src is too short for reading string with size %d; len(src)=%d", n, len(src))
	}
	return src[n:], src[:n], nil
	}

	// EncodeBytesBlock encodes a block of strings into dst.
	func EncodeBytesBlock(dst []byte, a [][]byte) []byte {
	u64s := GenerateUint64List(len(a))
	aLens := u64s.L[:0]
	for _, s := range a {
	if s == nil {
	aLens = append(aLens, 0)
	} else {
	// Offset by 1: "" (len 0) becomes 1, "a" (len 1) becomes 2, etc.
	aLens = append(aLens, uint64(len(s))+1)
	}
	}
	u64s.L = aLens
	dst = EncodeUint64Block(dst, u64s.L)
	ReleaseUint64List(u64s)

	bb := bbPool.Generate()
	b := bb.Buf
	for _, s := range a {
	b = append(b, s...)
	}
	bb.Buf = b
	dst = compressBlock(dst, bb.Buf)
	bbPool.Release(bb)

	return dst
	}

	// BytesBlockDecoder decodes a block of strings from src.
	// It reuses the underlying buffer for storing the decoded strings.
	type BytesBlockDecoder struct {
	data []byte
	}

	// Reset resets the decoder.
	func (bbd *BytesBlockDecoder) Reset() {
	bbd.data = bbd.data[:0]
	}

	// Decode decodes a block of strings from src.
	func (bbd *BytesBlockDecoder) Decode(dst [][]byte, src []byte, itemsCount uint64) ([][]byte, error) {
	u64List := GenerateUint64List(0)
	defer ReleaseUint64List(u64List)

	var tail []byte
	var err error
	u64List.L, tail, err = DecodeUint64Block(u64List.L[:0], src, itemsCount)
	if err != nil {
	return dst, fmt.Errorf("cannot decode string lengths: %w", err)
	}
	aLens := u64List.L
	src = tail

	dataLen := len(bbd.data)
	bbd.data, tail, err = decompressBlock(bbd.data, src)
	if err != nil {
	return dst, fmt.Errorf("cannot decode bytes block with strings: %w", err)
	}
	if len(tail) > 0 {
	return dst, fmt.Errorf("unexpected non-empty tail after reading bytes block with strings; len(tail)=%d", len(tail))
	}

	data := bbd.data[dataLen:]
	for _, sLen := range aLens {
	if sLen == 0 {
	dst = append(dst, nil)
	continue
	}
	// Reverse the offset
	actualLen := sLen - 1
	if uint64(len(data)) < actualLen {
	return dst, fmt.Errorf("cannot decode a string with the length %d bytes from %d bytes", actualLen, len(data))
	}
	if actualLen == 0 {
	// Explicitly create non-nil empty slice for empty strings
	dst = append(dst, []byte{})
	} else {
	dst = append(dst, data[:actualLen])
	}
	data = data[actualLen:]
	}

	return dst, nil
	}

	// DecodeWithTail decodes a block of strings from src and returns the remaining tail.
	func (bbd *BytesBlockDecoder) DecodeWithTail(dst [][]byte, src []byte, itemsCount uint64) ([][]byte, []byte, error) {
	u64List := GenerateUint64List(0)
	defer ReleaseUint64List(u64List)

	var tail []byte
	var err error
	u64List.L, tail, err = DecodeUint64Block(u64List.L[:0], src, itemsCount)
	if err != nil {
	return dst, nil, fmt.Errorf("cannot decode string lengths: %w", err)
	}
	aLens := u64List.L
	src = tail

	dataLen := len(bbd.data)
	bbd.data, tail, err = decompressBlock(bbd.data, src)
	if err != nil {
	return dst, tail, fmt.Errorf("cannot decode bytes block with strings: %w", err)
	}

	data := bbd.data[dataLen:]
	for _, sLen := range aLens {
	if sLen == 0 {
	dst = append(dst, nil)
	continue
	}
	// Reverse the offset
	actualLen := sLen - 1
	if uint64(len(data)) < actualLen {
	return dst, tail, fmt.Errorf("cannot decode a string with the length %d bytes from %d bytes", actualLen, len(data))
	}
	if actualLen == 0 {
	// Explicitly create non-nil empty slice for empty strings
	dst = append(dst, []byte{})
	} else {
	dst = append(dst, data[:actualLen])
	}
	data = data[actualLen:]
	}

	return dst, tail, nil
	}

	// EncodeUint64Block encodes a block of uint64 values into dst.
	func EncodeUint64Block(dst []byte, a []uint64) []byte {
	bb := bbPool.Generate()
	bb.Buf = encodeUint64List(bb.Buf[:0], a)
	dst = compressBlock(dst, bb.Buf)
	bbPool.Release(bb)
	return dst
	}

	// DecodeUint64Block decodes a block of uint64 values from src.
	func DecodeUint64Block(dst []uint64, src []byte, itemsCount uint64) ([]uint64, []byte, error) {
	bb := bbPool.Generate()
	defer bbPool.Release(bb)

	var err error
	bb.Buf, src, err = decompressBlock(bb.Buf[:0], src)
	if err != nil {
	return dst, src, fmt.Errorf("cannot decode bytes block: %w", err)
	}

	dst, err = decodeUint64List(dst, bb.Buf, itemsCount)
	if err != nil {
	return dst, src, fmt.Errorf("cannot decode %d uint64 items from bytes block of length %d bytes: %w", itemsCount, len(bb.Buf), err)
	}
	return dst, src, nil
	}

	const (
	uintBlockType8 = 0
	uintBlockType16 = 1
	uintBlockType32 = 2
	uintBlockType64 = 3
	)

	func encodeUint64List(dst []byte, a []uint64) []byte {
	nMax := uint64(0)
	for _, n := range a {
	if n > nMax {
	nMax = n
	}
	}
	switch {
	case nMax < (1 << 8):
	dst = append(dst, uintBlockType8)
	for _, n := range a {
	dst = append(dst, byte(n))
	}
	case nMax < (1 << 16):
	dst = append(dst, uintBlockType16)
	for _, n := range a {
	dst = Uint16ToBytes(dst, uint16(n))
	}
	case nMax < (1 << 32):
	dst = append(dst, uintBlockType32)
	for _, n := range a {
	dst = Uint32ToBytes(dst, uint32(n))
	}
	default:
	dst = append(dst, uintBlockType64)
	for _, n := range a {
	dst = Uint64ToBytes(dst, n)
	}
	}
	return dst
	}

	func decodeUint64List(dst []uint64, src []byte, itemsCount uint64) ([]uint64, error) {
	if len(src) < 1 {
	return dst, fmt.Errorf("cannot decode uint64 block type from empty src")
	}
	blockType := src[0]
	src = src[1:]

	switch blockType {
	case uintBlockType8:
	if uint64(len(src)) != itemsCount {
	return dst, fmt.Errorf("unexpected block length for %d items; got %d bytes; want %d bytes", itemsCount, len(src), itemsCount)
	}
	for _, v := range src {
	dst = append(dst, uint64(v))
	}
	case uintBlockType16:
	if uint64(len(src)) != 2*itemsCount {
	return dst, fmt.Errorf("unexpected block length for %d items; got %d bytes; want %d bytes", itemsCount, len(src), 2*itemsCount)
	}
	for len(src) > 0 {
	v := BytesToUint16(src)
	src = src[2:]
	dst = append(dst, uint64(v))
	}
	case uintBlockType32:
	if uint64(len(src)) != 4*itemsCount {
	return dst, fmt.Errorf("unexpected block length for %d items; got %d bytes; want %d bytes", itemsCount, len(src), 4*itemsCount)
	}
	for len(src) > 0 {
	v := BytesToUint32(src)
	src = src[4:]
	dst = append(dst, uint64(v))
	}
	case uintBlockType64:
	if uint64(len(src)) != 8*itemsCount {
	return dst, fmt.Errorf("unexpected block length for %d items; got %d bytes; want %d bytes", itemsCount, len(src), 8*itemsCount)
	}
	for len(src) > 0 {
	v := BytesToUint64(src)
	src = src[8:]
	dst = append(dst, v)
	}
	default:
	return dst, fmt.Errorf("unexpected uint64 block type: %d; want 0, 1, 2 or 3", blockType)
	}
	return dst, nil
	}

	const (
	compressTypePlain = 0
	compressTypeZSTD = 1
	)

	func compressBlock(dst, src []byte) []byte {
	if len(src) < 128 {
	dst = append(dst, compressTypePlain, byte(len(src)))
	return append(dst, src...)
	}

	dst = append(dst, compressTypeZSTD)
	bb := bbPool.Generate()
	bb.Buf = zstd.Compress(bb.Buf[:0], src, 1)
	dst = VarUint64ToBytes(dst, uint64(len(bb.Buf)))
	dst = append(dst, bb.Buf...)
	bbPool.Release(bb)
	return dst
	}

	func decompressBlock(dst, src []byte) ([]byte, []byte, error) {
	if len(src) < 1 {
	return dst, src, fmt.Errorf("cannot decode block type from empty src")
	}
	blockType := src[0]
	src = src[1:]
	switch blockType {
	case compressTypePlain:
	if len(src) < 1 {
	return dst, src, fmt.Errorf("cannot decode plain block size from empty src")
	}
	blockLen := int(src[0])
	src = src[1:]
	if len(src) < blockLen {
	return dst, src, fmt.Errorf("cannot read plain block with the size %d bytes from %b bytes", blockLen, len(src))
	}

	dst = append(dst, src[:blockLen]...)
	src = src[blockLen:]
	return dst, src, nil
	case compressTypeZSTD:
	tail, blockLen := BytesToVarUint64(src)
	src = tail
	if uint64(len(src)) < blockLen {
	return dst, src, fmt.Errorf("cannot read compressed block with the size %d bytes from %d bytes", blockLen, len(src))
	}
	compressedBlock := src[:blockLen]
	src = src[blockLen:]

	// Decompress the block
	var err error
	bb := bbPool.Generate()
	bb.Buf, err = zstd.Decompress(bb.Buf[:0], compressedBlock)
	if err != nil {
	return dst, src, fmt.Errorf("cannot decompress block: %w", err)
	}

	// Copy the decompressed block to dst.
	dst = append(dst, bb.Buf...)
	bbPool.Release(bb)
	return dst, src, nil
	default:
	return dst, src, fmt.Errorf("unexpected block type: %d; supported types: 0, 1", blockType)
	}
	}

	var bbPool = bytes.NewBufferPool("encoding.bytesBlock")