dataoutx.go

package protocol

import (
	"encoding/binary"
	"io"
	"math"
	"sync"
)

const (
	INT3MinValue  = int32(-8388608)      // 0xff800000 sign-extended
	INT3MaxValue  = int32(8388607)       // 0x007fffff
	LONG5MinValue = int64(-549755813888) // 0xffffff8000000000 sign-extended
	LONG5MaxValue = int64(549755813887)  // 0x0000007fffffffff
)

type DataOutputX struct {
	buf     []byte
	written int
	writer  io.Writer // optional: when set, writes to stream instead of buffer
	scratch [8]byte   // reusable scratch buffer — already on heap with the struct, avoids per-call allocation
}

var dataOutputXPool = sync.Pool{
	New: func() any {
		return &DataOutputX{buf: make([]byte, 0, 256)}
	},
}

// AcquireDataOutputX returns a DataOutputX from the pool, ready for reuse.
func AcquireDataOutputX() *DataOutputX {
	o := dataOutputXPool.Get().(*DataOutputX)
	o.Reset()
	return o
}

// ReleaseDataOutputX returns a DataOutputX to the pool.
// Do not use the DataOutputX after releasing; copy ToByteArray() before calling this.
func ReleaseDataOutputX(o *DataOutputX) {
	if o.writer != nil {
		return
	}
	dataOutputXPool.Put(o)
}

func NewDataOutputX() *DataOutputX {
	return &DataOutputX{buf: make([]byte, 0, 256)}
}

func NewDataOutputXWithSize(size int) *DataOutputX {
	return &DataOutputX{buf: make([]byte, 0, size)}
}

// NewDataOutputXStream creates a stream-based DataOutputX that writes to an io.Writer.
func NewDataOutputXStream(w io.Writer) *DataOutputX {
	return &DataOutputX{writer: w}
}

func (o *DataOutputX) ToByteArray() []byte {
	return o.buf
}

// Reset clears the buffer for reuse (e.g. from sync.Pool) without reallocating.
func (o *DataOutputX) Reset() {
	o.buf = o.buf[:0]
	o.written = 0
}

func (o *DataOutputX) Size() int {
	return o.written
}

// Flush writes any buffered data to the underlying writer (stream mode only).
func (o *DataOutputX) Flush() error {
	if o.writer == nil {
		return nil
	}
	if f, ok := o.writer.(interface{ Flush() error }); ok {
		return f.Flush()
	}
	return nil
}

func (o *DataOutputX) Write(b []byte) *DataOutputX {
	if o.writer != nil {
		o.writer.Write(b)
		o.written += len(b)
		return o
	}
	o.buf = append(o.buf, b...)
	o.written += len(b)
	return o
}

func (o *DataOutputX) WriteByte(v byte) *DataOutputX {
	if o.writer != nil {
		o.scratch[0] = v
		o.writer.Write(o.scratch[:1])
		o.written++
		return o
	}
	o.buf = append(o.buf, v)
	o.written++
	return o
}

func (o *DataOutputX) WriteBoolean(v bool) *DataOutputX {
	if v {
		o.WriteByte(1)
	} else {
		o.WriteByte(0)
	}
	return o
}

func (o *DataOutputX) WriteInt16(v int16) *DataOutputX {
	u := uint16(v)
	if o.writer != nil {
		binary.BigEndian.PutUint16(o.scratch[:2], u)
		o.writer.Write(o.scratch[:2])
		o.written += 2
		return o
	}
	o.buf = append(o.buf, byte(u>>8), byte(u))
	o.written += 2
	return o
}

func (o *DataOutputX) WriteUint16(v uint16) *DataOutputX {
	if o.writer != nil {
		binary.BigEndian.PutUint16(o.scratch[:2], v)
		o.writer.Write(o.scratch[:2])
		o.written += 2
		return o
	}
	o.buf = append(o.buf, byte(v>>8), byte(v))
	o.written += 2
	return o
}

func (o *DataOutputX) WriteInt32(v int32) *DataOutputX {
	u := uint32(v)
	if o.writer != nil {
		binary.BigEndian.PutUint32(o.scratch[:4], u)
		o.writer.Write(o.scratch[:4])
		o.written += 4
		return o
	}
	o.buf = append(o.buf, byte(u>>24), byte(u>>16), byte(u>>8), byte(u))
	o.written += 4
	return o
}

func (o *DataOutputX) WriteInt64(v int64) *DataOutputX {
	u := uint64(v)
	if o.writer != nil {
		binary.BigEndian.PutUint64(o.scratch[:8], u)
		o.writer.Write(o.scratch[:8])
		o.written += 8
		return o
	}
	o.buf = append(o.buf, byte(u>>56), byte(u>>48), byte(u>>40), byte(u>>32),
		byte(u>>24), byte(u>>16), byte(u>>8), byte(u))
	o.written += 8
	return o
}

func (o *DataOutputX) WriteFloat32(v float32) *DataOutputX {
	return o.WriteInt32(int32(math.Float32bits(v)))
}

func (o *DataOutputX) WriteFloat64(v float64) *DataOutputX {
	return o.WriteInt64(int64(math.Float64bits(v)))
}

func (o *DataOutputX) WriteInt3(v int32) *DataOutputX {
	if o.writer != nil {
		BigEndian.PutInt3(o.scratch[:3], v)
		o.writer.Write(o.scratch[:3])
		o.written += 3
		return o
	}
	o.buf = append(o.buf, byte(v>>16), byte(v>>8), byte(v))
	o.written += 3
	return o
}

func (o *DataOutputX) WriteLong5(v int64) *DataOutputX {
	if o.writer != nil {
		BigEndian.PutInt5(o.scratch[:5], v)
		o.writer.Write(o.scratch[:5])
		o.written += 5
		return o
	}
	o.buf = append(o.buf, byte(v>>32), byte(v>>24), byte(v>>16), byte(v>>8), byte(v))
	o.written += 5
	return o
}

func (o *DataOutputX) WriteDecimal(v int64) *DataOutputX {
	if v == 0 {
		o.WriteByte(0)
	} else if v >= -128 && v <= 127 {
		o.WriteByte(1)
		o.WriteByte(byte(int8(v)))
	} else if v >= -32768 && v <= 32767 {
		o.WriteByte(2)
		o.WriteInt16(int16(v))
	} else if int64(INT3MinValue) <= v && v <= int64(INT3MaxValue) {
		o.WriteByte(3)
		o.WriteInt3(int32(v))
	} else if v >= -2147483648 && v <= 2147483647 {
		o.WriteByte(4)
		o.WriteInt32(int32(v))
	} else if LONG5MinValue <= v && v <= LONG5MaxValue {
		o.WriteByte(5)
		o.WriteLong5(v)
	} else {
		o.WriteByte(8)
		o.WriteInt64(v)
	}
	return o
}

func (o *DataOutputX) WriteBlob(value []byte) *DataOutputX {
	if len(value) == 0 {
		o.WriteByte(0)
	} else {
		length := len(value)
		if length <= 253 {
			o.WriteByte(byte(length))
			o.Write(value)
		} else if length <= 65535 {
			o.WriteByte(255)
			o.WriteUint16(uint16(length))
			o.Write(value)
		} else {
			o.WriteByte(254)
			o.WriteInt32(int32(length))
			o.Write(value)
		}
	}
	return o
}

func (o *DataOutputX) WriteText(s string) *DataOutputX {
	if s == "" {
		o.WriteByte(0)
	} else {
		o.WriteBlob([]byte(s))
	}
	return o
}

func (o *DataOutputX) WriteIntBytes(b []byte) *DataOutputX {
	o.WriteInt32(int32(len(b)))
	o.Write(b)
	return o
}

func (o *DataOutputX) WriteShortBytes(b []byte) *DataOutputX {
	o.WriteInt16(int16(len(b)))
	o.Write(b)
	return o
}

func (o *DataOutputX) WriteArrayInt(v []int32) *DataOutputX {
	if v == nil {
		o.WriteInt16(0)
	} else {
		o.WriteInt16(int16(len(v)))
		for _, val := range v {
			o.WriteInt32(val)
		}
	}
	return o
}

func (o *DataOutputX) WriteArrayLong(v []int64) *DataOutputX {
	if v == nil {
		o.WriteInt16(0)
	} else {
		o.WriteInt16(int16(len(v)))
		for _, val := range v {
			o.WriteInt64(val)
		}
	}
	return o
}

func (o *DataOutputX) WriteArrayFloat(v []float32) *DataOutputX {
	if v == nil {
		o.WriteInt16(0)
	} else {
		o.WriteInt16(int16(len(v)))
		for _, val := range v {
			o.WriteFloat32(val)
		}
	}
	return o
}

func (o *DataOutputX) WriteDecimalArray(v []int64) *DataOutputX {
	if v == nil {
		o.WriteDecimal(0)
	} else {
		o.WriteDecimal(int64(len(v)))
		for _, val := range v {
			o.WriteDecimal(val)
		}
	}
	return o
}

func (o *DataOutputX) WriteDecimalIntArray(v []int32) *DataOutputX {
	if v == nil {
		o.WriteDecimal(0)
	} else {
		o.WriteDecimal(int64(len(v)))
		for _, val := range v {
			o.WriteDecimal(int64(val))
		}
	}
	return o
}