write_bytes.go 9.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411
  1. package msgp
  2. import (
  3. "math"
  4. "reflect"
  5. "time"
  6. )
  7. // ensure 'sz' extra bytes in 'b' btw len(b) and cap(b)
  8. func ensure(b []byte, sz int) ([]byte, int) {
  9. l := len(b)
  10. c := cap(b)
  11. if c-l < sz {
  12. o := make([]byte, (2*c)+sz) // exponential growth
  13. n := copy(o, b)
  14. return o[:n+sz], n
  15. }
  16. return b[:l+sz], l
  17. }
  18. // AppendMapHeader appends a map header with the
  19. // given size to the slice
  20. func AppendMapHeader(b []byte, sz uint32) []byte {
  21. switch {
  22. case sz <= 15:
  23. return append(b, wfixmap(uint8(sz)))
  24. case sz <= math.MaxUint16:
  25. o, n := ensure(b, 3)
  26. prefixu16(o[n:], mmap16, uint16(sz))
  27. return o
  28. default:
  29. o, n := ensure(b, 5)
  30. prefixu32(o[n:], mmap32, sz)
  31. return o
  32. }
  33. }
  34. // AppendArrayHeader appends an array header with
  35. // the given size to the slice
  36. func AppendArrayHeader(b []byte, sz uint32) []byte {
  37. switch {
  38. case sz <= 15:
  39. return append(b, wfixarray(uint8(sz)))
  40. case sz <= math.MaxUint16:
  41. o, n := ensure(b, 3)
  42. prefixu16(o[n:], marray16, uint16(sz))
  43. return o
  44. default:
  45. o, n := ensure(b, 5)
  46. prefixu32(o[n:], marray32, sz)
  47. return o
  48. }
  49. }
  50. // AppendNil appends a 'nil' byte to the slice
  51. func AppendNil(b []byte) []byte { return append(b, mnil) }
  52. // AppendFloat64 appends a float64 to the slice
  53. func AppendFloat64(b []byte, f float64) []byte {
  54. o, n := ensure(b, Float64Size)
  55. prefixu64(o[n:], mfloat64, math.Float64bits(f))
  56. return o
  57. }
  58. // AppendFloat32 appends a float32 to the slice
  59. func AppendFloat32(b []byte, f float32) []byte {
  60. o, n := ensure(b, Float32Size)
  61. prefixu32(o[n:], mfloat32, math.Float32bits(f))
  62. return o
  63. }
  64. // AppendInt64 appends an int64 to the slice
  65. func AppendInt64(b []byte, i int64) []byte {
  66. if i >= 0 {
  67. switch {
  68. case i <= math.MaxInt8:
  69. return append(b, wfixint(uint8(i)))
  70. case i <= math.MaxInt16:
  71. o, n := ensure(b, 3)
  72. putMint16(o[n:], int16(i))
  73. return o
  74. case i <= math.MaxInt32:
  75. o, n := ensure(b, 5)
  76. putMint32(o[n:], int32(i))
  77. return o
  78. default:
  79. o, n := ensure(b, 9)
  80. putMint64(o[n:], i)
  81. return o
  82. }
  83. }
  84. switch {
  85. case i >= -32:
  86. return append(b, wnfixint(int8(i)))
  87. case i >= math.MinInt8:
  88. o, n := ensure(b, 2)
  89. putMint8(o[n:], int8(i))
  90. return o
  91. case i >= math.MinInt16:
  92. o, n := ensure(b, 3)
  93. putMint16(o[n:], int16(i))
  94. return o
  95. case i >= math.MinInt32:
  96. o, n := ensure(b, 5)
  97. putMint32(o[n:], int32(i))
  98. return o
  99. default:
  100. o, n := ensure(b, 9)
  101. putMint64(o[n:], i)
  102. return o
  103. }
  104. }
  105. // AppendInt appends an int to the slice
  106. func AppendInt(b []byte, i int) []byte { return AppendInt64(b, int64(i)) }
  107. // AppendInt8 appends an int8 to the slice
  108. func AppendInt8(b []byte, i int8) []byte { return AppendInt64(b, int64(i)) }
  109. // AppendInt16 appends an int16 to the slice
  110. func AppendInt16(b []byte, i int16) []byte { return AppendInt64(b, int64(i)) }
  111. // AppendInt32 appends an int32 to the slice
  112. func AppendInt32(b []byte, i int32) []byte { return AppendInt64(b, int64(i)) }
  113. // AppendUint64 appends a uint64 to the slice
  114. func AppendUint64(b []byte, u uint64) []byte {
  115. switch {
  116. case u <= (1<<7)-1:
  117. return append(b, wfixint(uint8(u)))
  118. case u <= math.MaxUint8:
  119. o, n := ensure(b, 2)
  120. putMuint8(o[n:], uint8(u))
  121. return o
  122. case u <= math.MaxUint16:
  123. o, n := ensure(b, 3)
  124. putMuint16(o[n:], uint16(u))
  125. return o
  126. case u <= math.MaxUint32:
  127. o, n := ensure(b, 5)
  128. putMuint32(o[n:], uint32(u))
  129. return o
  130. default:
  131. o, n := ensure(b, 9)
  132. putMuint64(o[n:], u)
  133. return o
  134. }
  135. }
  136. // AppendUint appends a uint to the slice
  137. func AppendUint(b []byte, u uint) []byte { return AppendUint64(b, uint64(u)) }
  138. // AppendUint8 appends a uint8 to the slice
  139. func AppendUint8(b []byte, u uint8) []byte { return AppendUint64(b, uint64(u)) }
  140. // AppendByte is analogous to AppendUint8
  141. func AppendByte(b []byte, u byte) []byte { return AppendUint8(b, uint8(u)) }
  142. // AppendUint16 appends a uint16 to the slice
  143. func AppendUint16(b []byte, u uint16) []byte { return AppendUint64(b, uint64(u)) }
  144. // AppendUint32 appends a uint32 to the slice
  145. func AppendUint32(b []byte, u uint32) []byte { return AppendUint64(b, uint64(u)) }
  146. // AppendBytes appends bytes to the slice as MessagePack 'bin' data
  147. func AppendBytes(b []byte, bts []byte) []byte {
  148. sz := len(bts)
  149. var o []byte
  150. var n int
  151. switch {
  152. case sz <= math.MaxUint8:
  153. o, n = ensure(b, 2+sz)
  154. prefixu8(o[n:], mbin8, uint8(sz))
  155. n += 2
  156. case sz <= math.MaxUint16:
  157. o, n = ensure(b, 3+sz)
  158. prefixu16(o[n:], mbin16, uint16(sz))
  159. n += 3
  160. default:
  161. o, n = ensure(b, 5+sz)
  162. prefixu32(o[n:], mbin32, uint32(sz))
  163. n += 5
  164. }
  165. return o[:n+copy(o[n:], bts)]
  166. }
  167. // AppendBool appends a bool to the slice
  168. func AppendBool(b []byte, t bool) []byte {
  169. if t {
  170. return append(b, mtrue)
  171. }
  172. return append(b, mfalse)
  173. }
  174. // AppendString appends a string as a MessagePack 'str' to the slice
  175. func AppendString(b []byte, s string) []byte {
  176. sz := len(s)
  177. var n int
  178. var o []byte
  179. switch {
  180. case sz <= 31:
  181. o, n = ensure(b, 1+sz)
  182. o[n] = wfixstr(uint8(sz))
  183. n++
  184. case sz <= math.MaxUint8:
  185. o, n = ensure(b, 2+sz)
  186. prefixu8(o[n:], mstr8, uint8(sz))
  187. n += 2
  188. case sz <= math.MaxUint16:
  189. o, n = ensure(b, 3+sz)
  190. prefixu16(o[n:], mstr16, uint16(sz))
  191. n += 3
  192. default:
  193. o, n = ensure(b, 5+sz)
  194. prefixu32(o[n:], mstr32, uint32(sz))
  195. n += 5
  196. }
  197. return o[:n+copy(o[n:], s)]
  198. }
  199. // AppendStringFromBytes appends a []byte
  200. // as a MessagePack 'str' to the slice 'b.'
  201. func AppendStringFromBytes(b []byte, str []byte) []byte {
  202. sz := len(str)
  203. var n int
  204. var o []byte
  205. switch {
  206. case sz <= 31:
  207. o, n = ensure(b, 1+sz)
  208. o[n] = wfixstr(uint8(sz))
  209. n++
  210. case sz <= math.MaxUint8:
  211. o, n = ensure(b, 2+sz)
  212. prefixu8(o[n:], mstr8, uint8(sz))
  213. n += 2
  214. case sz <= math.MaxUint16:
  215. o, n = ensure(b, 3+sz)
  216. prefixu16(o[n:], mstr16, uint16(sz))
  217. n += 3
  218. default:
  219. o, n = ensure(b, 5+sz)
  220. prefixu32(o[n:], mstr32, uint32(sz))
  221. n += 5
  222. }
  223. return o[:n+copy(o[n:], str)]
  224. }
  225. // AppendComplex64 appends a complex64 to the slice as a MessagePack extension
  226. func AppendComplex64(b []byte, c complex64) []byte {
  227. o, n := ensure(b, Complex64Size)
  228. o[n] = mfixext8
  229. o[n+1] = Complex64Extension
  230. big.PutUint32(o[n+2:], math.Float32bits(real(c)))
  231. big.PutUint32(o[n+6:], math.Float32bits(imag(c)))
  232. return o
  233. }
  234. // AppendComplex128 appends a complex128 to the slice as a MessagePack extension
  235. func AppendComplex128(b []byte, c complex128) []byte {
  236. o, n := ensure(b, Complex128Size)
  237. o[n] = mfixext16
  238. o[n+1] = Complex128Extension
  239. big.PutUint64(o[n+2:], math.Float64bits(real(c)))
  240. big.PutUint64(o[n+10:], math.Float64bits(imag(c)))
  241. return o
  242. }
  243. // AppendTime appends a time.Time to the slice as a MessagePack extension
  244. func AppendTime(b []byte, t time.Time) []byte {
  245. o, n := ensure(b, TimeSize)
  246. t = t.UTC()
  247. o[n] = mext8
  248. o[n+1] = 12
  249. o[n+2] = TimeExtension
  250. putUnix(o[n+3:], t.Unix(), int32(t.Nanosecond()))
  251. return o
  252. }
  253. // AppendMapStrStr appends a map[string]string to the slice
  254. // as a MessagePack map with 'str'-type keys and values
  255. func AppendMapStrStr(b []byte, m map[string]string) []byte {
  256. sz := uint32(len(m))
  257. b = AppendMapHeader(b, sz)
  258. for key, val := range m {
  259. b = AppendString(b, key)
  260. b = AppendString(b, val)
  261. }
  262. return b
  263. }
  264. // AppendMapStrIntf appends a map[string]interface{} to the slice
  265. // as a MessagePack map with 'str'-type keys.
  266. func AppendMapStrIntf(b []byte, m map[string]interface{}) ([]byte, error) {
  267. sz := uint32(len(m))
  268. b = AppendMapHeader(b, sz)
  269. var err error
  270. for key, val := range m {
  271. b = AppendString(b, key)
  272. b, err = AppendIntf(b, val)
  273. if err != nil {
  274. return b, err
  275. }
  276. }
  277. return b, nil
  278. }
  279. // AppendIntf appends the concrete type of 'i' to the
  280. // provided []byte. 'i' must be one of the following:
  281. // - 'nil'
  282. // - A bool, float, string, []byte, int, uint, or complex
  283. // - A map[string]interface{} or map[string]string
  284. // - A []T, where T is another supported type
  285. // - A *T, where T is another supported type
  286. // - A type that satisfieds the msgp.Marshaler interface
  287. // - A type that satisfies the msgp.Extension interface
  288. func AppendIntf(b []byte, i interface{}) ([]byte, error) {
  289. if i == nil {
  290. return AppendNil(b), nil
  291. }
  292. // all the concrete types
  293. // for which we have methods
  294. switch i := i.(type) {
  295. case Marshaler:
  296. return i.MarshalMsg(b)
  297. case Extension:
  298. return AppendExtension(b, i)
  299. case bool:
  300. return AppendBool(b, i), nil
  301. case float32:
  302. return AppendFloat32(b, i), nil
  303. case float64:
  304. return AppendFloat64(b, i), nil
  305. case complex64:
  306. return AppendComplex64(b, i), nil
  307. case complex128:
  308. return AppendComplex128(b, i), nil
  309. case string:
  310. return AppendString(b, i), nil
  311. case []byte:
  312. return AppendBytes(b, i), nil
  313. case int8:
  314. return AppendInt8(b, i), nil
  315. case int16:
  316. return AppendInt16(b, i), nil
  317. case int32:
  318. return AppendInt32(b, i), nil
  319. case int64:
  320. return AppendInt64(b, i), nil
  321. case int:
  322. return AppendInt64(b, int64(i)), nil
  323. case uint:
  324. return AppendUint64(b, uint64(i)), nil
  325. case uint8:
  326. return AppendUint8(b, i), nil
  327. case uint16:
  328. return AppendUint16(b, i), nil
  329. case uint32:
  330. return AppendUint32(b, i), nil
  331. case uint64:
  332. return AppendUint64(b, i), nil
  333. case time.Time:
  334. return AppendTime(b, i), nil
  335. case map[string]interface{}:
  336. return AppendMapStrIntf(b, i)
  337. case map[string]string:
  338. return AppendMapStrStr(b, i), nil
  339. case []interface{}:
  340. b = AppendArrayHeader(b, uint32(len(i)))
  341. var err error
  342. for _, k := range i {
  343. b, err = AppendIntf(b, k)
  344. if err != nil {
  345. return b, err
  346. }
  347. }
  348. return b, nil
  349. }
  350. var err error
  351. v := reflect.ValueOf(i)
  352. switch v.Kind() {
  353. case reflect.Array, reflect.Slice:
  354. l := v.Len()
  355. b = AppendArrayHeader(b, uint32(l))
  356. for i := 0; i < l; i++ {
  357. b, err = AppendIntf(b, v.Index(i).Interface())
  358. if err != nil {
  359. return b, err
  360. }
  361. }
  362. return b, nil
  363. case reflect.Ptr:
  364. if v.IsNil() {
  365. return AppendNil(b), err
  366. }
  367. b, err = AppendIntf(b, v.Elem().Interface())
  368. return b, err
  369. default:
  370. return b, &ErrUnsupportedType{T: v.Type()}
  371. }
  372. }