wehub-resource-sync
313 lines
9.3 KiB
Go
313 lines
9.3 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 (
|
|
"bytes"
|
|
"encoding/binary"
|
|
"maps"
|
|
"math"
|
|
"slices"
|
|
|
|
"github.com/dolthub/doltgresql/core/id"
|
|
)
|
|
|
|
// Writer handles type-safe writing into a byte buffer, which may later be read from using Reader. The Writer will
|
|
// automatically grow as it is written to. The serialized forms of booleans, ints, uints, and floats are
|
|
// byte-comparable, meaning it is valid to use bytes.Compare() without needing to deserialize them. Variable-length
|
|
// encoded values, strings, and slices are not byte-comparable. This is not safe for concurrent use.
|
|
type Writer struct {
|
|
buf *bytes.Buffer
|
|
numSlice []byte
|
|
}
|
|
|
|
// NewWriter creates a new Writer with the given starting capacity. A larger starting capacity reduces reallocations at
|
|
// the cost of potentially wasted memory.
|
|
func NewWriter(capacity uint64) *Writer {
|
|
// If capacity is zero, then we'll set it to something arbitrary to try and minimize reallocations
|
|
if capacity == 0 {
|
|
capacity = 32
|
|
}
|
|
return &Writer{
|
|
buf: bytes.NewBuffer(make([]byte, 0, capacity)),
|
|
numSlice: make([]byte, 10), // 10 bytes will cover all integers and variable-length integers
|
|
}
|
|
}
|
|
|
|
// Bool writes a bool.
|
|
func (writer *Writer) Bool(val bool) {
|
|
if val {
|
|
writer.buf.WriteByte(1)
|
|
} else {
|
|
writer.buf.WriteByte(0)
|
|
}
|
|
}
|
|
|
|
// Int8 writes an int8.
|
|
func (writer *Writer) Int8(val int8) {
|
|
writer.buf.WriteByte(byte(val) + (1 << 7))
|
|
}
|
|
|
|
// Int16 writes an int16.
|
|
func (writer *Writer) Int16(val int16) {
|
|
writer.Uint16(uint16(val) + (1 << 15))
|
|
}
|
|
|
|
// Int32 writes an int32.
|
|
func (writer *Writer) Int32(val int32) {
|
|
writer.Uint32(uint32(val) + (1 << 31))
|
|
}
|
|
|
|
// Int64 writes an int64.
|
|
func (writer *Writer) Int64(val int64) {
|
|
writer.Uint64(uint64(val) + (1 << 63))
|
|
}
|
|
|
|
// Uint8 writes a uint8.
|
|
func (writer *Writer) Uint8(val uint8) {
|
|
writer.buf.WriteByte(val)
|
|
}
|
|
|
|
// Uint16 writes a uint16.
|
|
func (writer *Writer) Uint16(val uint16) {
|
|
binary.BigEndian.PutUint16(writer.numSlice, val)
|
|
writer.buf.Write(writer.numSlice[:2])
|
|
}
|
|
|
|
// Uint32 writes a uint32.
|
|
func (writer *Writer) Uint32(val uint32) {
|
|
binary.BigEndian.PutUint32(writer.numSlice, val)
|
|
writer.buf.Write(writer.numSlice[:4])
|
|
}
|
|
|
|
// Uint64 writes a uint64.
|
|
func (writer *Writer) Uint64(val uint64) {
|
|
binary.BigEndian.PutUint64(writer.numSlice, val)
|
|
writer.buf.Write(writer.numSlice[:8])
|
|
}
|
|
|
|
// Byte writes 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 (writer *Writer) Byte(val byte) {
|
|
writer.buf.WriteByte(val)
|
|
}
|
|
|
|
// Float32 writes a float32.
|
|
func (writer *Writer) Float32(val float32) {
|
|
// Float encoding produces byte-comparable serialized values when looking at the exponent and mantissa. This means
|
|
// that we just have to flip to exponent and mantissa for negative values, and flip the sign bit so that negatives
|
|
// sort before positives. To do this, we take advantage of arithmetic shifting by casting to a signed integer to
|
|
// create a mask that only exists for negative values.
|
|
uval := math.Float32bits(val)
|
|
uval ^= uint32(int32(uval)>>31) & 0x7FFFFFFF
|
|
uval ^= 0x80000000
|
|
writer.Uint32(uval)
|
|
}
|
|
|
|
// Float64 writes a float64.
|
|
func (writer *Writer) Float64(val float64) {
|
|
// Float encoding produces byte-comparable serialized values when looking at the exponent and mantissa. This means
|
|
// that we just have to flip to exponent and mantissa for negative values, and flip the sign bit so that negatives
|
|
// sort before positives. To do this, we take advantage of arithmetic shifting by casting to a signed integer to
|
|
// create a mask that only exists for negative values.
|
|
uval := math.Float64bits(val)
|
|
uval ^= uint64(int64(uval)>>63) & 0x7FFFFFFFFFFFFFFF
|
|
uval ^= 0x8000000000000000
|
|
writer.Uint64(uval)
|
|
}
|
|
|
|
// VariableInt writes an int64 using variable-length encoding. Smaller values use less space at the cost of larger
|
|
// values using more bytes, but this is generally more space-efficient. This does carry a small computational hit when
|
|
// reading.
|
|
func (writer *Writer) VariableInt(val int64) {
|
|
count := binary.PutVarint(writer.numSlice, val)
|
|
writer.buf.Write(writer.numSlice[:count])
|
|
}
|
|
|
|
// VariableUint writes a uint64 using variable-length encoding. Smaller values use less space at the cost of larger
|
|
// values using more bytes, but this is generally more space-efficient. This does carry a small computational hit when
|
|
// reading.
|
|
func (writer *Writer) VariableUint(val uint64) {
|
|
count := binary.PutUvarint(writer.numSlice, val)
|
|
writer.buf.Write(writer.numSlice[:count])
|
|
}
|
|
|
|
// String writes a string.
|
|
func (writer *Writer) String(val string) {
|
|
writer.VariableUint(uint64(len(val)))
|
|
writer.buf.WriteString(val)
|
|
}
|
|
|
|
// Id writes an internal ID.
|
|
func (writer *Writer) Id(val id.Id) {
|
|
writer.String(string(val))
|
|
}
|
|
|
|
// BoolSlice writes a bool slice.
|
|
func (writer *Writer) BoolSlice(vals []bool) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.Bool(vals[i])
|
|
}
|
|
}
|
|
|
|
// Int8Slice writes an int8 slice.
|
|
func (writer *Writer) Int8Slice(vals []int8) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.Int8(vals[i])
|
|
}
|
|
}
|
|
|
|
// Int16Slice writes an int16 slice.
|
|
func (writer *Writer) Int16Slice(vals []int16) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.Int16(vals[i])
|
|
}
|
|
}
|
|
|
|
// Int32Slice writes an int32 slice.
|
|
func (writer *Writer) Int32Slice(vals []int32) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.Int32(vals[i])
|
|
}
|
|
}
|
|
|
|
// Int64Slice writes an int64 slice.
|
|
func (writer *Writer) Int64Slice(vals []int64) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.Int64(vals[i])
|
|
}
|
|
}
|
|
|
|
// Uint8Slice writes a uint8 slice.
|
|
func (writer *Writer) Uint8Slice(vals []uint8) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
writer.buf.Write(vals)
|
|
}
|
|
|
|
// Uint16Slice writes a uint16 slice.
|
|
func (writer *Writer) Uint16Slice(vals []uint16) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.Uint16(vals[i])
|
|
}
|
|
}
|
|
|
|
// Uint32Slice writes a uint32 slice.
|
|
func (writer *Writer) Uint32Slice(vals []uint32) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.Uint32(vals[i])
|
|
}
|
|
}
|
|
|
|
// Uint64Slice writes a uint64 slice.
|
|
func (writer *Writer) Uint64Slice(vals []uint64) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.Uint64(vals[i])
|
|
}
|
|
}
|
|
|
|
// ByteSlice writes 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 (writer *Writer) ByteSlice(vals []byte) {
|
|
writer.Uint8Slice(vals)
|
|
}
|
|
|
|
// Float32Slice writes a float32 slice.
|
|
func (writer *Writer) Float32Slice(vals []float32) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.Float32(vals[i])
|
|
}
|
|
}
|
|
|
|
// Float64Slice writes a float64 slice.
|
|
func (writer *Writer) Float64Slice(vals []float64) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.Float64(vals[i])
|
|
}
|
|
}
|
|
|
|
// VariableIntSlice writes an int64 slice using variable-length encoding. Smaller values use less space at the cost of
|
|
// larger values using more space, but this is generally more space-efficient. This does carry a computational hit when
|
|
// reading.
|
|
func (writer *Writer) VariableIntSlice(vals []int64) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.VariableInt(vals[i])
|
|
}
|
|
}
|
|
|
|
// VariableUintSlice writes a uint64 slice using variable-length encoding. Smaller values use less space at the cost of
|
|
// larger values using more space, but this is generally more space-efficient. This does carry a computational hit when
|
|
// reading.
|
|
func (writer *Writer) VariableUintSlice(vals []uint64) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.VariableUint(vals[i])
|
|
}
|
|
}
|
|
|
|
// StringSlice writes a string slice.
|
|
func (writer *Writer) StringSlice(vals []string) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.String(vals[i])
|
|
}
|
|
}
|
|
|
|
// StringMap writes a map of strings, keyed by strings.
|
|
func (writer *Writer) StringMap(m map[string]string) {
|
|
writer.VariableUint(uint64(len(m)))
|
|
// We iterate over the sorted set of keys for determinism
|
|
for _, k := range slices.Sorted(maps.Keys(m)) {
|
|
writer.String(k)
|
|
writer.String(m[k])
|
|
}
|
|
}
|
|
|
|
// IdSlice writes a slice of internal IDs.
|
|
func (writer *Writer) IdSlice(vals []id.Id) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.Id(vals[i])
|
|
}
|
|
}
|
|
|
|
// IdTypeSlice writes a slice of internal type IDs.
|
|
func (writer *Writer) IdTypeSlice(vals []id.Type) {
|
|
writer.VariableUint(uint64(len(vals)))
|
|
for i := range vals {
|
|
writer.Id(vals[i].AsId())
|
|
}
|
|
}
|
|
|
|
// Data returns the data written to the Writer.
|
|
func (writer *Writer) Data() []byte {
|
|
return writer.buf.Bytes()
|
|
}
|
|
|
|
// Reset resets the Writer to be empty, but it retains the underlying storage for use by future writes.
|
|
func (writer *Writer) Reset() {
|
|
writer.buf.Reset()
|
|
}
|