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