meta/msgpack_writer.go

package meta

import (
	"io"

	"github.com/pkg/errors"

	"git.bit5.ru/backend/msgpack"
)

type msgpWriter struct {
	stack   []*msgpWriteContainer
	current *msgpWriteContainer
}

type msgpWriteContainer struct {
	arr []interface{}
}

func NewMsgpackWriter() Writer {
	wr := &msgpWriter{}
	wr.stack = make([]*msgpWriteContainer, 0, 1)
	wr.current = newMsgpWriteContainer()
	wr.stack = append(wr.stack, wr.current)
	return wr
}

func newMsgpWriteContainer() *msgpWriteContainer {
	return &msgpWriteContainer{make([]interface{}, 0, 1)}
}

func (state *msgpWriteContainer) add(v interface{}) error {
	state.arr = append(state.arr, v)
	return nil
}

func (wr *msgpWriter) WriteI8(v int8, field string) error {
	return wr.current.add(v)
}

func (wr *msgpWriter) WriteU8(v uint8, field string) error {
	return wr.current.add(v)
}
func (wr *msgpWriter) WriteI16(v int16, field string) error {
	return wr.current.add(v)
}
func (wr *msgpWriter) WriteU16(v uint16, field string) error {
	return wr.current.add(v)
}

func (wr *msgpWriter) WriteI32(v int32, field string) error {
	return wr.current.add(v)
}
func (wr *msgpWriter) WriteU32(v uint32, field string) error {
	return wr.current.add(v)
}

func (wr *msgpWriter) WriteI64(v int64, field string) error {
	return wr.current.add(v)
}
func (wr *msgpWriter) WriteU64(v uint64, field string) error {
	return wr.current.add(v)
}

func (wr *msgpWriter) WriteBool(v bool, field string) error {
	return wr.current.add(v)
}

func (wr *msgpWriter) WriteFloat(v float32, field string) error {
	return wr.current.add(v)
}

func (wr *msgpWriter) WriteDouble(v float64, field string) error {
	return wr.current.add(v)
}

func (wr *msgpWriter) WriteString(v string, field string) error {
	return wr.current.add(v)
}

func (wr *msgpWriter) WriteBlob(v []byte, field string) error {
	return wr.current.add(v)
}

func (wr *msgpWriter) BeginContainer(field string) {
	wr.current = newMsgpWriteContainer()
	wr.stack = append(wr.stack, wr.current)
}

func (wr *msgpWriter) EndContainer() error {
	if len(wr.stack) <= 1 {
		return errors.New("No open container")
	}
	last := wr.current
	wr.stack[len(wr.stack)-1] = nil
	wr.stack = wr.stack[0 : len(wr.stack)-1]
	wr.current = wr.stack[len(wr.stack)-1]

	bytes, err := msgpack.Marshal(last.arr)
	if err != nil {
		return err
	}

	//NOTE: using custom msgpack encoder
	wr.current.add(&msgpCustomBytes{bytes})

	return nil
}

func (wr *msgpWriter) GetData() ([]byte, error) {
	if len(wr.stack) != 1 {
		return nil, errors.New("Stack isn't empty")
	}

	if len(wr.current.arr) != 1 {
		return nil, errors.New("Arr size isn't valid")
	}

	return msgpack.Marshal(wr.current.arr[0])
}

type msgpCustomBytes struct {
	v []byte
}

func (msgp *msgpCustomBytes) EncodeMsgpack(writer io.Writer) error {
	_, err := writer.Write(msgp.v)
	return errors.WithStack(err)
}