cloudflared-mirror/vendor/zombiezen.com/go/capnproto2/internal/packed/packed.go

338 lines
6.3 KiB
Go
Raw Permalink Normal View History

// Package packed provides functions to read and write the "packed"
// compression scheme described at https://capnproto.org/encoding.html#packing.
package packed
import (
"bufio"
"errors"
"io"
)
const wordSize = 8
// Special case tags.
const (
zeroTag byte = 0x00
unpackedTag byte = 0xff
)
// Pack appends the packed version of src to dst and returns the
// resulting slice. len(src) must be a multiple of 8 or Pack panics.
func Pack(dst, src []byte) []byte {
if len(src)%wordSize != 0 {
panic("packed.Pack len(src) must be a multiple of 8")
}
var buf [wordSize]byte
for len(src) > 0 {
var hdr byte
n := 0
for i := uint(0); i < wordSize; i++ {
if src[i] != 0 {
hdr |= 1 << i
buf[n] = src[i]
n++
}
}
dst = append(dst, hdr)
dst = append(dst, buf[:n]...)
src = src[wordSize:]
switch hdr {
case zeroTag:
z := min(numZeroWords(src), 0xff)
dst = append(dst, byte(z))
src = src[z*wordSize:]
case unpackedTag:
i := 0
end := min(len(src), 0xff*wordSize)
for i < end {
zeros := 0
for _, b := range src[i : i+wordSize] {
if b == 0 {
zeros++
}
}
if zeros > 1 {
break
}
i += wordSize
}
rawWords := byte(i / wordSize)
dst = append(dst, rawWords)
dst = append(dst, src[:i]...)
src = src[i:]
}
}
return dst
}
// numZeroWords returns the number of leading zero words in b.
func numZeroWords(b []byte) int {
for i, bb := range b {
if bb != 0 {
return i / wordSize
}
}
return len(b) / wordSize
}
// Unpack appends the unpacked version of src to dst and returns the
// resulting slice.
func Unpack(dst, src []byte) ([]byte, error) {
for len(src) > 0 {
tag := src[0]
src = src[1:]
pstart := len(dst)
dst = allocWords(dst, 1)
p := dst[pstart : pstart+wordSize]
if len(src) >= wordSize {
i := 0
nz := tag & 1
p[0] = src[i] & -nz
i += int(nz)
nz = tag >> 1 & 1
p[1] = src[i] & -nz
i += int(nz)
nz = tag >> 2 & 1
p[2] = src[i] & -nz
i += int(nz)
nz = tag >> 3 & 1
p[3] = src[i] & -nz
i += int(nz)
nz = tag >> 4 & 1
p[4] = src[i] & -nz
i += int(nz)
nz = tag >> 5 & 1
p[5] = src[i] & -nz
i += int(nz)
nz = tag >> 6 & 1
p[6] = src[i] & -nz
i += int(nz)
nz = tag >> 7 & 1
p[7] = src[i] & -nz
i += int(nz)
src = src[i:]
} else {
for i := uint(0); i < wordSize; i++ {
if tag&(1<<i) == 0 {
continue
}
if len(src) == 0 {
return dst, io.ErrUnexpectedEOF
}
p[i] = src[0]
src = src[1:]
}
}
switch tag {
case zeroTag:
if len(src) == 0 {
return dst, io.ErrUnexpectedEOF
}
dst = allocWords(dst, int(src[0]))
src = src[1:]
case unpackedTag:
if len(src) == 0 {
return dst, io.ErrUnexpectedEOF
}
start := len(dst)
dst = allocWords(dst, int(src[0]))
src = src[1:]
n := copy(dst[start:], src)
if n < len(dst)-start {
return dst, io.ErrUnexpectedEOF
}
src = src[n:]
}
}
return dst, nil
}
func allocWords(p []byte, n int) []byte {
target := len(p) + n*wordSize
if cap(p) >= target {
pp := p[len(p):target]
for i := range pp {
pp[i] = 0
}
return p[:target]
}
newcap := cap(p)
doublecap := newcap + newcap
if target > doublecap {
newcap = target
} else {
if len(p) < 1024 {
newcap = doublecap
} else {
for newcap < target {
newcap += newcap / 4
}
}
}
pp := make([]byte, target, newcap)
copy(pp, p)
return pp
}
// A Reader decompresses a packed byte stream.
type Reader struct {
// ReadWord state
rd *bufio.Reader
err error
zeroes int
literal int
// Read state
word [wordSize]byte
wordIdx int
}
// NewReader returns a reader that decompresses a packed stream from r.
func NewReader(r *bufio.Reader) *Reader {
return &Reader{rd: r, wordIdx: wordSize}
}
func min(a, b int) int {
if b < a {
return b
}
return a
}
// ReadWord decompresses the next word from the underlying stream.
func (r *Reader) ReadWord(p []byte) error {
if len(p) < wordSize {
return errors.New("packed: read word buffer too small")
}
r.wordIdx = wordSize // if the caller tries to call ReadWord and Read, don't give them partial words.
if r.err != nil {
err := r.err
r.err = nil
return err
}
p = p[:wordSize]
switch {
case r.zeroes > 0:
r.zeroes--
for i := range p {
p[i] = 0
}
return nil
case r.literal > 0:
r.literal--
_, err := io.ReadFull(r.rd, p)
return err
}
var tag byte
if r.rd.Buffered() < wordSize+1 {
var err error
tag, err = r.rd.ReadByte()
if err != nil {
return err
}
for i := range p {
p[i] = 0
}
for i := uint(0); i < wordSize; i++ {
if tag&(1<<i) != 0 {
p[i], err = r.rd.ReadByte()
if err != nil {
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
return err
}
} else {
p[i] = 0
}
}
} else {
b, _ := r.rd.Peek(wordSize + 1)
tag = b[0]
i := 1
nz := tag & 1
p[0] = b[i] & -nz
i += int(nz)
nz = tag >> 1 & 1
p[1] = b[i] & -nz
i += int(nz)
nz = tag >> 2 & 1
p[2] = b[i] & -nz
i += int(nz)
nz = tag >> 3 & 1
p[3] = b[i] & -nz
i += int(nz)
nz = tag >> 4 & 1
p[4] = b[i] & -nz
i += int(nz)
nz = tag >> 5 & 1
p[5] = b[i] & -nz
i += int(nz)
nz = tag >> 6 & 1
p[6] = b[i] & -nz
i += int(nz)
nz = tag >> 7 & 1
p[7] = b[i] & -nz
i += int(nz)
discard(r.rd, i)
}
switch tag {
case zeroTag:
z, err := r.rd.ReadByte()
if err != nil {
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
r.err = err
return err
}
r.zeroes = int(z)
case unpackedTag:
l, err := r.rd.ReadByte()
if err != nil {
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
r.err = err
return err
}
r.literal = int(l)
}
return nil
}
// Read reads up to len(p) bytes into p. This will decompress whole
// words at a time, so mixing calls to Read and ReadWord may lead to
// bytes missing.
func (r *Reader) Read(p []byte) (n int, err error) {
if r.wordIdx < wordSize {
n = copy(p, r.word[r.wordIdx:])
r.wordIdx += n
}
for n < len(p) {
if r.rd.Buffered() < wordSize+1 && n > 0 {
return n, nil
}
if len(p)-n >= wordSize {
err := r.ReadWord(p[n:])
if err != nil {
return n, err
}
n += wordSize
} else {
err := r.ReadWord(r.word[:])
if err != nil {
return n, err
}
r.wordIdx = copy(p[n:], r.word[:])
n += r.wordIdx
}
}
return n, nil
}