242 lines
4.5 KiB
Go
242 lines
4.5 KiB
Go
package msgpack
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
"reflect"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/vmihailenco/msgpack/v4/codes"
|
|
)
|
|
|
|
const (
|
|
sortMapKeysFlag uint32 = 1 << iota
|
|
structAsArrayFlag
|
|
encodeUsingJSONFlag
|
|
useCompactIntsFlag
|
|
useCompactFloatsFlag
|
|
)
|
|
|
|
type writer interface {
|
|
io.Writer
|
|
WriteByte(byte) error
|
|
}
|
|
|
|
type byteWriter struct {
|
|
io.Writer
|
|
|
|
buf [1]byte
|
|
}
|
|
|
|
func newByteWriter(w io.Writer) *byteWriter {
|
|
bw := new(byteWriter)
|
|
bw.Reset(w)
|
|
return bw
|
|
}
|
|
|
|
func (bw *byteWriter) Reset(w io.Writer) {
|
|
bw.Writer = w
|
|
}
|
|
|
|
func (bw *byteWriter) WriteByte(c byte) error {
|
|
bw.buf[0] = c
|
|
_, err := bw.Write(bw.buf[:])
|
|
return err
|
|
}
|
|
|
|
//------------------------------------------------------------------------------
|
|
|
|
var encPool = sync.Pool{
|
|
New: func() interface{} {
|
|
return NewEncoder(nil)
|
|
},
|
|
}
|
|
|
|
// Marshal returns the MessagePack encoding of v.
|
|
func Marshal(v interface{}) ([]byte, error) {
|
|
enc := encPool.Get().(*Encoder)
|
|
|
|
var buf bytes.Buffer
|
|
enc.Reset(&buf)
|
|
|
|
err := enc.Encode(v)
|
|
b := buf.Bytes()
|
|
|
|
encPool.Put(enc)
|
|
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return b, err
|
|
}
|
|
|
|
type Encoder struct {
|
|
w writer
|
|
|
|
buf []byte
|
|
timeBuf []byte
|
|
bootstrap [9 + 12]byte
|
|
|
|
intern map[string]int
|
|
|
|
flags uint32
|
|
}
|
|
|
|
// NewEncoder returns a new encoder that writes to w.
|
|
func NewEncoder(w io.Writer) *Encoder {
|
|
e := new(Encoder)
|
|
e.buf = e.bootstrap[:9]
|
|
e.timeBuf = e.bootstrap[9 : 9+12]
|
|
e.Reset(w)
|
|
return e
|
|
}
|
|
|
|
func (e *Encoder) Reset(w io.Writer) {
|
|
if bw, ok := w.(writer); ok {
|
|
e.w = bw
|
|
} else if bw, ok := e.w.(*byteWriter); ok {
|
|
bw.Reset(w)
|
|
} else {
|
|
e.w = newByteWriter(w)
|
|
}
|
|
|
|
for k := range e.intern {
|
|
delete(e.intern, k)
|
|
}
|
|
|
|
//TODO:
|
|
//e.sortMapKeys = false
|
|
//e.structAsArray = false
|
|
//e.useJSONTag = false
|
|
//e.useCompact = false
|
|
}
|
|
|
|
// SortMapKeys causes the Encoder to encode map keys in increasing order.
|
|
// Supported map types are:
|
|
// - map[string]string
|
|
// - map[string]interface{}
|
|
func (e *Encoder) SortMapKeys(on bool) *Encoder {
|
|
if on {
|
|
e.flags |= sortMapKeysFlag
|
|
} else {
|
|
e.flags &= ^sortMapKeysFlag
|
|
}
|
|
return e
|
|
}
|
|
|
|
// StructAsArray causes the Encoder to encode Go structs as msgpack arrays.
|
|
func (e *Encoder) StructAsArray(on bool) *Encoder {
|
|
if on {
|
|
e.flags |= structAsArrayFlag
|
|
} else {
|
|
e.flags &= ^structAsArrayFlag
|
|
}
|
|
return e
|
|
}
|
|
|
|
// UseJSONTag causes the Encoder to use json struct tag as fallback option
|
|
// if there is no msgpack tag.
|
|
func (e *Encoder) UseJSONTag(on bool) *Encoder {
|
|
if on {
|
|
e.flags |= encodeUsingJSONFlag
|
|
} else {
|
|
e.flags &= ^encodeUsingJSONFlag
|
|
}
|
|
return e
|
|
}
|
|
|
|
// UseCompactEncoding causes the Encoder to chose the most compact encoding.
|
|
// For example, it allows to encode small Go int64 as msgpack int8 saving 7 bytes.
|
|
func (e *Encoder) UseCompactEncoding(on bool) *Encoder {
|
|
if on {
|
|
e.flags |= useCompactIntsFlag
|
|
} else {
|
|
e.flags &= ^useCompactIntsFlag
|
|
}
|
|
return e
|
|
}
|
|
|
|
// UseCompactFloats causes the Encoder to chose a compact integer encoding
|
|
// for floats that can be represented as integers.
|
|
func (e *Encoder) UseCompactFloats(on bool) {
|
|
if on {
|
|
e.flags |= useCompactFloatsFlag
|
|
} else {
|
|
e.flags &= ^useCompactFloatsFlag
|
|
}
|
|
}
|
|
|
|
func (e *Encoder) Encode(v interface{}) error {
|
|
switch v := v.(type) {
|
|
case nil:
|
|
return e.EncodeNil()
|
|
case string:
|
|
return e.EncodeString(v)
|
|
case []byte:
|
|
return e.EncodeBytes(v)
|
|
case int:
|
|
return e.encodeInt64Cond(int64(v))
|
|
case int64:
|
|
return e.encodeInt64Cond(v)
|
|
case uint:
|
|
return e.encodeUint64Cond(uint64(v))
|
|
case uint64:
|
|
return e.encodeUint64Cond(v)
|
|
case bool:
|
|
return e.EncodeBool(v)
|
|
case float32:
|
|
return e.EncodeFloat32(v)
|
|
case float64:
|
|
return e.EncodeFloat64(v)
|
|
case time.Duration:
|
|
return e.encodeInt64Cond(int64(v))
|
|
case time.Time:
|
|
return e.EncodeTime(v)
|
|
}
|
|
return e.EncodeValue(reflect.ValueOf(v))
|
|
}
|
|
|
|
func (e *Encoder) EncodeMulti(v ...interface{}) error {
|
|
for _, vv := range v {
|
|
if err := e.Encode(vv); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (e *Encoder) EncodeValue(v reflect.Value) error {
|
|
fn := getEncoder(v.Type())
|
|
return fn(e, v)
|
|
}
|
|
|
|
func (e *Encoder) EncodeNil() error {
|
|
return e.writeCode(codes.Nil)
|
|
}
|
|
|
|
func (e *Encoder) EncodeBool(value bool) error {
|
|
if value {
|
|
return e.writeCode(codes.True)
|
|
}
|
|
return e.writeCode(codes.False)
|
|
}
|
|
|
|
func (e *Encoder) EncodeDuration(d time.Duration) error {
|
|
return e.EncodeInt(int64(d))
|
|
}
|
|
|
|
func (e *Encoder) writeCode(c codes.Code) error {
|
|
return e.w.WriteByte(byte(c))
|
|
}
|
|
|
|
func (e *Encoder) write(b []byte) error {
|
|
_, err := e.w.Write(b)
|
|
return err
|
|
}
|
|
|
|
func (e *Encoder) writeString(s string) error {
|
|
_, err := e.w.Write(stringToBytes(s))
|
|
return err
|
|
}
|