Files
2026-07-13 12:32:25 +08:00

419 lines
11 KiB
Go

// Copyright 2024 Dolthub, Inc.
//
// Licensed 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 utils
import (
"encoding/binary"
"math"
"github.com/dolthub/doltgresql/core/id"
)
// Reader handles type-safe reading from a byte slice, which was created by a Writer. This is not safe for concurrent
// use.
type Reader struct {
buf []byte
offset uint64
}
// NewReader creates a new Reader that will read from the given data.
func NewReader(data []byte) *Reader {
return &Reader{
buf: data,
offset: 0,
}
}
// Bool reads a bool.
func (reader *Reader) Bool() bool {
reader.offset += 1
if reader.buf[reader.offset-1] == 1 {
return true
} else {
return false
}
}
// Int8 reads an int8.
func (reader *Reader) Int8() int8 {
return int8(reader.Uint8() - (1 << 7))
}
// Int16 reads an int16.
func (reader *Reader) Int16() int16 {
return int16(reader.Uint16() - (1 << 15))
}
// Int32 reads an int32.
func (reader *Reader) Int32() int32 {
return int32(reader.Uint32() - (1 << 31))
}
// Int64 reads an int64.
func (reader *Reader) Int64() int64 {
return int64(reader.Uint64() - (1 << 63))
}
// Uint8 reads a uint8.
func (reader *Reader) Uint8() uint8 {
reader.offset += 1
return reader.buf[reader.offset-1]
}
// Uint16 reads a uint16.
func (reader *Reader) Uint16() uint16 {
reader.offset += 2
return binary.BigEndian.Uint16(reader.buf[reader.offset-2:])
}
// Uint32 reads a uint32.
func (reader *Reader) Uint32() uint32 {
reader.offset += 4
return binary.BigEndian.Uint32(reader.buf[reader.offset-4:])
}
// Uint64 reads a uint64.
func (reader *Reader) Uint64() uint64 {
reader.offset += 8
return binary.BigEndian.Uint64(reader.buf[reader.offset-8:])
}
// Byte reads a byte. This is equivalent to Uint8, but is included since it is more common to refer to a byte rather
// than a uint8.
func (reader *Reader) Byte() byte {
return reader.Uint8()
}
// Float32 reads a float32.
func (reader *Reader) Float32() float32 {
// For more details, look at Writer.Float32
uval := reader.Uint32()
uval ^= 0x80000000
uval ^= uint32(int32(uval)>>31) & 0x7FFFFFFF
return math.Float32frombits(uval)
}
// Float64 reads a float64.
func (reader *Reader) Float64() float64 {
// For more details, look at Writer.Float64
uval := reader.Uint64()
uval ^= 0x8000000000000000
uval ^= uint64(int64(uval)>>63) & 0x7FFFFFFFFFFFFFFF
return math.Float64frombits(uval)
}
// VariableInt reads an int64 that was written using variable-length encoding.
func (reader *Reader) VariableInt() int64 {
uval := reader.VariableUint()
// binary.PutVarint performs the inverse of this and writes a uint64, so we're undoing it to get the original int64
val := int64(uval >> 1)
if uval&1 != 0 {
val = ^val
}
return val
}
// VariableUint reads a uint64 that was written using variable-length encoding.
func (reader *Reader) VariableUint() uint64 {
// This has been adapted from one of our blogs:
// https://www.dolthub.com/blog/2021-01-08-optimizing-varint-decoding/
b := uint64(reader.buf[reader.offset])
if b < 0x80 {
reader.offset += 1
return b
}
x := b & 0x7f
b = uint64(reader.buf[reader.offset+1])
if b < 0x80 {
reader.offset += 2
return x | (b << 7)
}
x |= (b & 0x7f) << 7
b = uint64(reader.buf[reader.offset+2])
if b < 0x80 {
reader.offset += 3
return x | (b << 14)
}
x |= (b & 0x7f) << 14
b = uint64(reader.buf[reader.offset+3])
if b < 0x80 {
reader.offset += 4
return x | (b << 21)
}
x |= (b & 0x7f) << 21
b = uint64(reader.buf[reader.offset+4])
if b < 0x80 {
reader.offset += 5
return x | (b << 28)
}
x |= (b & 0x7f) << 28
b = uint64(reader.buf[reader.offset+5])
if b < 0x80 {
reader.offset += 6
return x | (b << 35)
}
x |= (b & 0x7f) << 35
b = uint64(reader.buf[reader.offset+6])
if b < 0x80 {
reader.offset += 7
return x | (b << 42)
}
x |= (b & 0x7f) << 42
b = uint64(reader.buf[reader.offset+7])
if b < 0x80 {
reader.offset += 8
return x | (b << 49)
}
x |= (b & 0x7f) << 49
b = uint64(reader.buf[reader.offset+8])
if b < 0x80 {
reader.offset += 9
return x | (b << 56)
}
x |= (b & 0x7f) << 56
b = uint64(reader.buf[reader.offset+9])
if b < 0x80 {
reader.offset += 10
return x | (b << 63)
}
reader.offset += 10
return 0xffffffffffffffff
}
// String reads a string.
func (reader *Reader) String() string {
length := reader.VariableUint()
reader.offset += length
return string(reader.buf[reader.offset-length : reader.offset])
}
// Id reads an internal ID.
func (reader *Reader) Id() id.Id {
return id.Id(reader.String())
}
// BoolSlice reads a bool slice.
func (reader *Reader) BoolSlice() []bool {
count := reader.VariableUint()
vals := make([]bool, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.Bool()
}
return vals
}
// Int8Slice reads an int8 slice.
func (reader *Reader) Int8Slice() []int8 {
count := reader.VariableUint()
vals := make([]int8, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.Int8()
}
return vals
}
// Int16Slice reads an int16 slice.
func (reader *Reader) Int16Slice() []int16 {
count := reader.VariableUint()
vals := make([]int16, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.Int16()
}
return vals
}
// Int32Slice reads an int32 slice.
func (reader *Reader) Int32Slice() []int32 {
count := reader.VariableUint()
vals := make([]int32, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.Int32()
}
return vals
}
// Int64Slice reads an int64 slice.
func (reader *Reader) Int64Slice() []int64 {
count := reader.VariableUint()
vals := make([]int64, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.Int64()
}
return vals
}
// Uint8Slice reads a uint8 slice.
func (reader *Reader) Uint8Slice() []uint8 {
count := reader.VariableUint()
vals := make([]uint8, count)
copy(vals, reader.buf[reader.offset:reader.offset+count])
reader.offset += count
return vals
}
// Uint16Slice reads a uint16 slice.
func (reader *Reader) Uint16Slice() []uint16 {
count := reader.VariableUint()
vals := make([]uint16, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.Uint16()
}
return vals
}
// Uint32Slice reads a uint32 slice.
func (reader *Reader) Uint32Slice() []uint32 {
count := reader.VariableUint()
vals := make([]uint32, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.Uint32()
}
return vals
}
// Uint64Slice reads a uint64 slice.
func (reader *Reader) Uint64Slice() []uint64 {
count := reader.VariableUint()
vals := make([]uint64, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.Uint64()
}
return vals
}
// ByteSlice reads a byte slice. This is equivalent to Uint8Slice, but is included since it is more common to refer to
// byte slices than uint8 slices.
func (reader *Reader) ByteSlice() []byte {
return reader.Uint8Slice()
}
// Float32Slice reads a float32 slice.
func (reader *Reader) Float32Slice() []float32 {
count := reader.VariableUint()
vals := make([]float32, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.Float32()
}
return vals
}
// Float64Slice reads a float64 slice.
func (reader *Reader) Float64Slice() []float64 {
count := reader.VariableUint()
vals := make([]float64, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.Float64()
}
return vals
}
// VariableIntSlice reads an int64 slice that was written using variable-length encoding.
func (reader *Reader) VariableIntSlice() []int64 {
count := reader.VariableUint()
vals := make([]int64, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.VariableInt()
}
return vals
}
// VariableUintSlice reads a uint64 slice that was written using variable-length encoding.
func (reader *Reader) VariableUintSlice() []uint64 {
count := reader.VariableUint()
vals := make([]uint64, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.VariableUint()
}
return vals
}
// StringSlice reads a string slice.
func (reader *Reader) StringSlice() []string {
count := reader.VariableUint()
vals := make([]string, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.String()
}
return vals
}
// StringMap reads a map of strings, keyed by strings.
func (reader *Reader) StringMap() map[string]string {
count := reader.VariableUint()
vals := make(map[string]string, count)
for i := uint64(0); i < count; i++ {
key := reader.String()
vals[key] = reader.String()
}
return vals
}
// IdSlice reads a slice of internal IDs.
func (reader *Reader) IdSlice() []id.Id {
count := reader.VariableUint()
vals := make([]id.Id, count)
for i := uint64(0); i < count; i++ {
vals[i] = reader.Id()
}
return vals
}
// IdTypeSlice reads a slice of internal type IDs.
func (reader *Reader) IdTypeSlice() []id.Type {
count := reader.VariableUint()
vals := make([]id.Type, count)
for i := uint64(0); i < count; i++ {
vals[i] = id.Type(reader.Id())
}
return vals
}
// IsEmpty returns true when all of the data has been read.
func (reader *Reader) IsEmpty() bool {
return reader.offset >= uint64(len(reader.buf))
}
// RemainingBytes returns the number of bytes that have not yet been read.
func (reader *Reader) RemainingBytes() uint64 {
if reader.IsEmpty() {
return 0
}
return uint64(len(reader.buf)) - reader.offset
}
// BytesRead returns the number of bytes that have been read.
func (reader *Reader) BytesRead() uint64 {
return reader.offset
}
// AdvanceReader reads the next N bytes from the given Reader. This is only available for specific, performance-oriented
// circumstances, and should never be used otherwise. This is a standalone function to discourage its use, as it will
// not show up as a function of the Reader object in most IDEs. This uses a branchless comparison to limit the size of n
// to the end of the reader, and it does not allocate a new byte slice (returns a portion from the original byte slice).
func AdvanceReader(reader *Reader, n uint64) []byte {
// This branchless code makes an assumption that the reader contains less than 9223372036854775808 bytes.
// With this assumption, it is equivalent to the following:
// if reader.offset + n > uint64(len(reader.buf)) || n >= 0x8000000000000000 {
// n = uint64(len(reader.buf)) - reader.offset
// }
maxN := uint64(len(reader.buf)) - reader.offset
delta := int64(n - maxN)
mask := (delta >> 63) ^ (int64(n&0x8000000000000000) >> 63)
n = (n & uint64(mask)) | (maxN & ^uint64(mask))
reader.offset += n
return reader.buf[reader.offset-n : reader.offset]
}