encode.go 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269
  1. package msgpack
  2. import (
  3. "bytes"
  4. "io"
  5. "reflect"
  6. "sync"
  7. "time"
  8. "github.com/vmihailenco/msgpack/v5/msgpcode"
  9. )
  10. const (
  11. sortMapKeysFlag uint32 = 1 << iota
  12. arrayEncodedStructsFlag
  13. useCompactIntsFlag
  14. useCompactFloatsFlag
  15. useInternedStringsFlag
  16. omitEmptyFlag
  17. )
  18. type writer interface {
  19. io.Writer
  20. WriteByte(byte) error
  21. }
  22. type byteWriter struct {
  23. io.Writer
  24. }
  25. func newByteWriter(w io.Writer) byteWriter {
  26. return byteWriter{
  27. Writer: w,
  28. }
  29. }
  30. func (bw byteWriter) WriteByte(c byte) error {
  31. _, err := bw.Write([]byte{c})
  32. return err
  33. }
  34. //------------------------------------------------------------------------------
  35. var encPool = sync.Pool{
  36. New: func() interface{} {
  37. return NewEncoder(nil)
  38. },
  39. }
  40. func GetEncoder() *Encoder {
  41. return encPool.Get().(*Encoder)
  42. }
  43. func PutEncoder(enc *Encoder) {
  44. enc.w = nil
  45. encPool.Put(enc)
  46. }
  47. // Marshal returns the MessagePack encoding of v.
  48. func Marshal(v interface{}) ([]byte, error) {
  49. enc := GetEncoder()
  50. var buf bytes.Buffer
  51. enc.Reset(&buf)
  52. err := enc.Encode(v)
  53. b := buf.Bytes()
  54. PutEncoder(enc)
  55. if err != nil {
  56. return nil, err
  57. }
  58. return b, err
  59. }
  60. type Encoder struct {
  61. w writer
  62. buf []byte
  63. timeBuf []byte
  64. dict map[string]int
  65. flags uint32
  66. structTag string
  67. }
  68. // NewEncoder returns a new encoder that writes to w.
  69. func NewEncoder(w io.Writer) *Encoder {
  70. e := &Encoder{
  71. buf: make([]byte, 9),
  72. }
  73. e.Reset(w)
  74. return e
  75. }
  76. // Writer returns the Encoder's writer.
  77. func (e *Encoder) Writer() io.Writer {
  78. return e.w
  79. }
  80. // Reset discards any buffered data, resets all state, and switches the writer to write to w.
  81. func (e *Encoder) Reset(w io.Writer) {
  82. e.ResetDict(w, nil)
  83. }
  84. // ResetDict is like Reset, but also resets the dict.
  85. func (e *Encoder) ResetDict(w io.Writer, dict map[string]int) {
  86. e.resetWriter(w)
  87. e.flags = 0
  88. e.structTag = ""
  89. e.dict = dict
  90. }
  91. func (e *Encoder) WithDict(dict map[string]int, fn func(*Encoder) error) error {
  92. oldDict := e.dict
  93. e.dict = dict
  94. err := fn(e)
  95. e.dict = oldDict
  96. return err
  97. }
  98. func (e *Encoder) resetWriter(w io.Writer) {
  99. if bw, ok := w.(writer); ok {
  100. e.w = bw
  101. } else {
  102. e.w = newByteWriter(w)
  103. }
  104. }
  105. // SetSortMapKeys causes the Encoder to encode map keys in increasing order.
  106. // Supported map types are:
  107. // - map[string]string
  108. // - map[string]interface{}
  109. func (e *Encoder) SetSortMapKeys(on bool) *Encoder {
  110. if on {
  111. e.flags |= sortMapKeysFlag
  112. } else {
  113. e.flags &= ^sortMapKeysFlag
  114. }
  115. return e
  116. }
  117. // SetCustomStructTag causes the Encoder to use a custom struct tag as
  118. // fallback option if there is no msgpack tag.
  119. func (e *Encoder) SetCustomStructTag(tag string) {
  120. e.structTag = tag
  121. }
  122. // SetOmitEmpty causes the Encoder to omit empty values by default.
  123. func (e *Encoder) SetOmitEmpty(on bool) {
  124. if on {
  125. e.flags |= omitEmptyFlag
  126. } else {
  127. e.flags &= ^omitEmptyFlag
  128. }
  129. }
  130. // UseArrayEncodedStructs causes the Encoder to encode Go structs as msgpack arrays.
  131. func (e *Encoder) UseArrayEncodedStructs(on bool) {
  132. if on {
  133. e.flags |= arrayEncodedStructsFlag
  134. } else {
  135. e.flags &= ^arrayEncodedStructsFlag
  136. }
  137. }
  138. // UseCompactEncoding causes the Encoder to chose the most compact encoding.
  139. // For example, it allows to encode small Go int64 as msgpack int8 saving 7 bytes.
  140. func (e *Encoder) UseCompactInts(on bool) {
  141. if on {
  142. e.flags |= useCompactIntsFlag
  143. } else {
  144. e.flags &= ^useCompactIntsFlag
  145. }
  146. }
  147. // UseCompactFloats causes the Encoder to chose a compact integer encoding
  148. // for floats that can be represented as integers.
  149. func (e *Encoder) UseCompactFloats(on bool) {
  150. if on {
  151. e.flags |= useCompactFloatsFlag
  152. } else {
  153. e.flags &= ^useCompactFloatsFlag
  154. }
  155. }
  156. // UseInternedStrings causes the Encoder to intern strings.
  157. func (e *Encoder) UseInternedStrings(on bool) {
  158. if on {
  159. e.flags |= useInternedStringsFlag
  160. } else {
  161. e.flags &= ^useInternedStringsFlag
  162. }
  163. }
  164. func (e *Encoder) Encode(v interface{}) error {
  165. switch v := v.(type) {
  166. case nil:
  167. return e.EncodeNil()
  168. case string:
  169. return e.EncodeString(v)
  170. case []byte:
  171. return e.EncodeBytes(v)
  172. case int:
  173. return e.EncodeInt(int64(v))
  174. case int64:
  175. return e.encodeInt64Cond(v)
  176. case uint:
  177. return e.EncodeUint(uint64(v))
  178. case uint64:
  179. return e.encodeUint64Cond(v)
  180. case bool:
  181. return e.EncodeBool(v)
  182. case float32:
  183. return e.EncodeFloat32(v)
  184. case float64:
  185. return e.EncodeFloat64(v)
  186. case time.Duration:
  187. return e.encodeInt64Cond(int64(v))
  188. case time.Time:
  189. return e.EncodeTime(v)
  190. }
  191. return e.EncodeValue(reflect.ValueOf(v))
  192. }
  193. func (e *Encoder) EncodeMulti(v ...interface{}) error {
  194. for _, vv := range v {
  195. if err := e.Encode(vv); err != nil {
  196. return err
  197. }
  198. }
  199. return nil
  200. }
  201. func (e *Encoder) EncodeValue(v reflect.Value) error {
  202. fn := getEncoder(v.Type())
  203. return fn(e, v)
  204. }
  205. func (e *Encoder) EncodeNil() error {
  206. return e.writeCode(msgpcode.Nil)
  207. }
  208. func (e *Encoder) EncodeBool(value bool) error {
  209. if value {
  210. return e.writeCode(msgpcode.True)
  211. }
  212. return e.writeCode(msgpcode.False)
  213. }
  214. func (e *Encoder) EncodeDuration(d time.Duration) error {
  215. return e.EncodeInt(int64(d))
  216. }
  217. func (e *Encoder) writeCode(c byte) error {
  218. return e.w.WriteByte(c)
  219. }
  220. func (e *Encoder) write(b []byte) error {
  221. _, err := e.w.Write(b)
  222. return err
  223. }
  224. func (e *Encoder) writeString(s string) error {
  225. _, err := e.w.Write(stringToBytes(s))
  226. return err
  227. }