encode.go

   1 package msgpack
   2
   3 import (
   4         "bytes"
   5         "io"
   6         "reflect"
   7         "time"
   8
   9         "github.com/vmihailenco/msgpack/codes"
  10 )
  11
  12 type writer interface {
  13         io.Writer
  14         WriteByte(byte) error
  15         WriteString(string) (int, error)
  16 }
  17
  18 type byteWriter struct {
  19         io.Writer
  20
  21         buf       []byte
  22         bootstrap [64]byte
  23 }
  24
  25 func newByteWriter(w io.Writer) *byteWriter {
  26         bw := &byteWriter{
  27                 Writer: w,
  28         }
  29         bw.buf = bw.bootstrap[:]
  30         return bw
  31 }
  32
  33 func (w *byteWriter) WriteByte(c byte) error {
  34         w.buf = w.buf[:1]
  35         w.buf[0] = c
  36         _, err := w.Write(w.buf)
  37         return err
  38 }
  39
  40 func (w *byteWriter) WriteString(s string) (int, error) {
  41         w.buf = append(w.buf[:0], s...)
  42         return w.Write(w.buf)
  43 }
  44
  45 // Marshal returns the MessagePack encoding of v.
  46 func Marshal(v interface{}) ([]byte, error) {
  47         var buf bytes.Buffer
  48         err := NewEncoder(&buf).Encode(v)
  49         return buf.Bytes(), err
  50 }
  51
  52 type Encoder struct {
  53         w   writer
  54         buf []byte
  55
  56         sortMapKeys   bool
  57         structAsArray bool
  58         useJSONTag    bool
  59         useCompact    bool
  60 }
  61
  62 // NewEncoder returns a new encoder that writes to w.
  63 func NewEncoder(w io.Writer) *Encoder {
  64         bw, ok := w.(writer)
  65         if !ok {
  66                 bw = newByteWriter(w)
  67         }
  68         return &Encoder{
  69                 w:   bw,
  70                 buf: make([]byte, 9),
  71         }
  72 }
  73
  74 // SortMapKeys causes the Encoder to encode map keys in increasing order.
  75 // Supported map types are:
  76 //   - map[string]string
  77 //   - map[string]interface{}
  78 func (e *Encoder) SortMapKeys(flag bool) *Encoder {
  79         e.sortMapKeys = flag
  80         return e
  81 }
  82
  83 // StructAsArray causes the Encoder to encode Go structs as MessagePack arrays.
  84 func (e *Encoder) StructAsArray(flag bool) *Encoder {
  85         e.structAsArray = flag
  86         return e
  87 }
  88
  89 // UseJSONTag causes the Encoder to use json struct tag as fallback option
  90 // if there is no msgpack tag.
  91 func (e *Encoder) UseJSONTag(flag bool) *Encoder {
  92         e.useJSONTag = flag
  93         return e
  94 }
  95
  96 // UseCompactEncoding causes the Encoder to chose the most compact encoding.
  97 // For example, it allows to encode Go int64 as msgpack int8 saving 7 bytes.
  98 func (e *Encoder) UseCompactEncoding(flag bool) *Encoder {
  99         e.useCompact = flag
 100         return e
 101 }
 102
 103 func (e *Encoder) Encode(v interface{}) error {
 104         switch v := v.(type) {
 105         case nil:
 106                 return e.EncodeNil()
 107         case string:
 108                 return e.EncodeString(v)
 109         case []byte:
 110                 return e.EncodeBytes(v)
 111         case int:
 112                 return e.encodeInt64Cond(int64(v))
 113         case int64:
 114                 return e.encodeInt64Cond(v)
 115         case uint:
 116                 return e.encodeUint64Cond(uint64(v))
 117         case uint64:
 118                 return e.encodeUint64Cond(v)
 119         case bool:
 120                 return e.EncodeBool(v)
 121         case float32:
 122                 return e.EncodeFloat32(v)
 123         case float64:
 124                 return e.EncodeFloat64(v)
 125         case time.Duration:
 126                 return e.encodeInt64Cond(int64(v))
 127         case time.Time:
 128                 return e.EncodeTime(v)
 129         }
 130         return e.EncodeValue(reflect.ValueOf(v))
 131 }
 132
 133 func (e *Encoder) EncodeMulti(v ...interface{}) error {
 134         for _, vv := range v {
 135                 if err := e.Encode(vv); err != nil {
 136                         return err
 137                 }
 138         }
 139         return nil
 140 }
 141
 142 func (e *Encoder) EncodeValue(v reflect.Value) error {
 143         fn := getEncoder(v.Type())
 144         return fn(e, v)
 145 }
 146
 147 func (e *Encoder) EncodeNil() error {
 148         return e.writeCode(codes.Nil)
 149 }
 150
 151 func (e *Encoder) EncodeBool(value bool) error {
 152         if value {
 153                 return e.writeCode(codes.True)
 154         }
 155         return e.writeCode(codes.False)
 156 }
 157
 158 func (e *Encoder) writeCode(c codes.Code) error {
 159         return e.w.WriteByte(byte(c))
 160 }
 161
 162 func (e *Encoder) write(b []byte) error {
 163         _, err := e.w.Write(b)
 164         return err
 165 }
 166
 167 func (e *Encoder) writeString(s string) error {
 168         _, err := e.w.WriteString(s)
 169         return err
 170 }