encode.go 3.5 KB

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