Merge pull request #330 from alpineiq/v5

small optimizations.

decode.go

@@ -14,14 +14,16 @@ import (
 )
 
 const (
-	looseInterfaceDecodingFlag uint32 = 1 << iota
-	disallowUnknownFieldsFlag
+	bytesAllocLimit = 1 << 20 // 1mb
+	sliceAllocLimit = 1e6     // 1m elements
+	maxMapSize      = 1e6     // 1m elements
 )
 
 const (
-	bytesAllocLimit = 1e6 // 1mb
-	sliceAllocLimit = 1e4
-	maxMapSize      = 1e6
+	looseInterfaceDecodingFlag uint32 = 1 << iota
+	disallowUnknownFieldsFlag
+	usePreallocateValues
+	disableAllocLimitFlag
 )
 
 type bufReader interface {
@@ -53,7 +55,7 @@ func PutDecoder(dec *Decoder) {
 // in the value pointed to by v.
 func Unmarshal(data []byte, v interface{}) error {
 	dec := GetDecoder()
-
+	dec.UsePreallocateValues(true)
 	dec.Reset(bytes.NewReader(data))
 	err := dec.Decode(v)
 
@@ -64,16 +66,14 @@ func Unmarshal(data []byte, v interface{}) error {
 
 // A Decoder reads and decodes MessagePack values from an input stream.
 type Decoder struct {
-	r   io.Reader
-	s   io.ByteScanner
-	buf []byte
-
-	rec []byte // accumulates read data if not nil
-
+	r          io.Reader
+	s          io.ByteScanner
+	mapDecoder func(*Decoder) (interface{}, error)
+	structTag  string
+	buf        []byte
+	rec        []byte
 	dict       []string
 	flags      uint32
-	structTag  string
-	mapDecoder func(*Decoder) (interface{}, error)
 }
 
 // NewDecoder returns a new decoder that reads from r.
@@ -95,10 +95,9 @@ func (d *Decoder) Reset(r io.Reader) {
 
 // ResetDict is like Reset, but also resets the dict.
 func (d *Decoder) ResetDict(r io.Reader, dict []string) {
-	d.resetReader(r)
+	d.ResetReader(r)
 	d.flags = 0
 	d.structTag = ""
-	d.mapDecoder = nil
 	d.dict = dict
 }
 
@@ -110,10 +109,16 @@ func (d *Decoder) WithDict(dict []string, fn func(*Decoder) error) error {
 	return err
 }
 
-func (d *Decoder) resetReader(r io.Reader) {
+func (d *Decoder) ResetReader(r io.Reader) {
+	d.mapDecoder = nil
+	d.dict = nil
+
 	if br, ok := r.(bufReader); ok {
 		d.r = br
 		d.s = br
+	} else if r == nil {
+		d.r = nil
+		d.s = nil
 	} else {
 		br := bufio.NewReader(r)
 		d.r = br
@@ -161,6 +166,24 @@ func (d *Decoder) UseInternedStrings(on bool) {
 	}
 }
 
+// UsePreallocateValues enables preallocating values in chunks
+func (d *Decoder) UsePreallocateValues(on bool) {
+	if on {
+		d.flags |= usePreallocateValues
+	} else {
+		d.flags &= ^usePreallocateValues
+	}
+}
+
+// DisableAllocLimit enables fully allocating slices/maps when the size is known
+func (d *Decoder) DisableAllocLimit(on bool) {
+	if on {
+		d.flags |= disableAllocLimitFlag
+	} else {
+		d.flags &= ^disableAllocLimitFlag
+	}
+}
+
 // Buffered returns a reader of the data remaining in the Decoder's buffer.
 // The reader is valid until the next call to Decode.
 func (d *Decoder) Buffered() io.Reader {
@@ -603,7 +626,11 @@ func (d *Decoder) readFull(b []byte) error {
 
 func (d *Decoder) readN(n int) ([]byte, error) {
 	var err error
-	d.buf, err = readN(d.r, d.buf, n)
+	if d.flags&disableAllocLimitFlag != 0 {
+		d.buf, err = readN(d.r, d.buf, n)
+	} else {
+		d.buf, err = readNGrow(d.r, d.buf, n)
+	}
 	if err != nil {
 		return nil, err
 	}
@@ -615,6 +642,24 @@ func (d *Decoder) readN(n int) ([]byte, error) {
 }
 
 func readN(r io.Reader, b []byte, n int) ([]byte, error) {
+	if b == nil {
+		if n == 0 {
+			return make([]byte, 0), nil
+		}
+		b = make([]byte, 0, n)
+	}
+
+	if n > cap(b) {
+		b = append(b, make([]byte, n-len(b))...)
+	} else if n <= cap(b) {
+		b = b[:n]
+	}
+
+	_, err := io.ReadFull(r, b)
+	return b, err
+}
+
+func readNGrow(r io.Reader, b []byte, n int) ([]byte, error) {
 	if b == nil {
 		if n == 0 {
 			return make([]byte, 0), nil

decode_map.go

@@ -35,7 +35,11 @@ func decodeMapValue(d *Decoder, v reflect.Value) error {
 	}
 
 	if v.IsNil() {
-		v.Set(reflect.MakeMap(typ))
+		ln := n
+		if d.flags&disableAllocLimitFlag == 0 {
+			ln = min(ln, maxMapSize)
+		}
+		v.Set(reflect.MakeMapWithSize(typ, ln))
 	}
 	if n == 0 {
 		return nil
@@ -106,7 +110,11 @@ func (d *Decoder) decodeMapStringStringPtr(ptr *map[string]string) error {
 
 	m := *ptr
 	if m == nil {
-		*ptr = make(map[string]string, min(size, maxMapSize))
+		ln := size
+		if d.flags&disableAllocLimitFlag == 0 {
+			ln = min(size, maxMapSize)
+		}
+		*ptr = make(map[string]string, ln)
 		m = *ptr
 	}
 
@@ -149,7 +157,7 @@ func (d *Decoder) DecodeMap() (map[string]interface{}, error) {
 		return nil, nil
 	}
 
-	m := make(map[string]interface{}, min(n, maxMapSize))
+	m := make(map[string]interface{}, n)
 
 	for i := 0; i < n; i++ {
 		mk, err := d.DecodeString()
@@ -176,7 +184,7 @@ func (d *Decoder) DecodeUntypedMap() (map[interface{}]interface{}, error) {
 		return nil, nil
 	}
 
-	m := make(map[interface{}]interface{}, min(n, maxMapSize))
+	m := make(map[interface{}]interface{}, n)
 
 	for i := 0; i < n; i++ {
 		mk, err := d.decodeInterfaceCond()
@@ -224,7 +232,13 @@ func (d *Decoder) DecodeTypedMap() (interface{}, error) {
 	}
 
 	mapType := reflect.MapOf(keyType, valueType)
-	mapValue := reflect.MakeMap(mapType)
+
+	ln := n
+	if d.flags&disableAllocLimitFlag == 0 {
+		ln = min(ln, maxMapSize)
+	}
+
+	mapValue := reflect.MakeMapWithSize(mapType, ln)
 	mapValue.SetMapIndex(reflect.ValueOf(key), reflect.ValueOf(value))
 
 	n--
@@ -236,17 +250,18 @@ func (d *Decoder) DecodeTypedMap() (interface{}, error) {
 }
 
 func (d *Decoder) decodeTypedMapValue(v reflect.Value, n int) error {
-	typ := v.Type()
-	keyType := typ.Key()
-	valueType := typ.Elem()
-
+	var (
+		typ       = v.Type()
+		keyType   = typ.Key()
+		valueType = typ.Elem()
+	)
 	for i := 0; i < n; i++ {
-		mk := reflect.New(keyType).Elem()
+		mk := d.newValue(keyType).Elem()
 		if err := d.DecodeValue(mk); err != nil {
 			return err
 		}
 
-		mv := reflect.New(valueType).Elem()
+		mv := d.newValue(valueType).Elem()
 		if err := d.DecodeValue(mv); err != nil {
 			return err
 		}

decode_query.go

@@ -11,9 +11,8 @@ import (
 type queryResult struct {
 	query       string
 	key         string
+	values      []interface{}
 	hasAsterisk bool
-
-	values []interface{}
 }
 
 func (q *queryResult) nextKey() {

decode_slice.go

@@ -49,7 +49,7 @@ func (d *Decoder) decodeStringSlicePtr(ptr *[]string) error {
 		return nil
 	}
 
-	ss := makeStrings(*ptr, n)
+	ss := makeStrings(*ptr, n, d.flags&disableAllocLimitFlag != 0)
 	for i := 0; i < n; i++ {
 		s, err := d.DecodeString()
 		if err != nil {
@@ -62,8 +62,8 @@ func (d *Decoder) decodeStringSlicePtr(ptr *[]string) error {
 	return nil
 }
 
-func makeStrings(s []string, n int) []string {
-	if n > sliceAllocLimit {
+func makeStrings(s []string, n int, noLimit bool) []string {
+	if !noLimit && n > sliceAllocLimit {
 		n = sliceAllocLimit
 	}
 
@@ -101,10 +101,17 @@ func decodeSliceValue(d *Decoder, v reflect.Value) error {
 		v.Set(v.Slice(0, v.Cap()))
 	}
 
+	noLimit := d.flags&disableAllocLimitFlag != 1
+
+	if noLimit && n > v.Len() {
+		v.Set(growSliceValue(v, n, noLimit))
+	}
+
 	for i := 0; i < n; i++ {
-		if i >= v.Len() {
-			v.Set(growSliceValue(v, n))
+		if !noLimit && i >= v.Len() {
+			v.Set(growSliceValue(v, n, noLimit))
 		}
+
 		elem := v.Index(i)
 		if err := d.DecodeValue(elem); err != nil {
 			return err
@@ -114,9 +121,9 @@ func decodeSliceValue(d *Decoder, v reflect.Value) error {
 	return nil
 }
 
-func growSliceValue(v reflect.Value, n int) reflect.Value {
+func growSliceValue(v reflect.Value, n int, noLimit bool) reflect.Value {
 	diff := n - v.Len()
-	if diff > sliceAllocLimit {
+	if !noLimit && diff > sliceAllocLimit {
 		diff = sliceAllocLimit
 	}
 	v = reflect.AppendSlice(v, reflect.MakeSlice(v.Type(), diff, diff))
@@ -163,7 +170,7 @@ func (d *Decoder) decodeSlice(c byte) ([]interface{}, error) {
 		return nil, nil
 	}
 
-	s := make([]interface{}, 0, min(n, sliceAllocLimit))
+	s := make([]interface{}, 0, n)
 	for i := 0; i < n; i++ {
 		v, err := d.decodeInterfaceCond()
 		if err != nil {

decode_typgen.go

@@ -0,0 +1,46 @@
+package msgpack
+
+import (
+	"reflect"
+	"sync"
+)
+
+var cachedValues struct {
+	m map[reflect.Type]chan reflect.Value
+	sync.RWMutex
+}
+
+func cachedValue(t reflect.Type) reflect.Value {
+	cachedValues.RLock()
+	ch := cachedValues.m[t]
+	cachedValues.RUnlock()
+	if ch != nil {
+		return <-ch
+	}
+
+	cachedValues.Lock()
+	defer cachedValues.Unlock()
+	if ch = cachedValues.m[t]; ch != nil {
+		return <-ch
+	}
+
+	ch = make(chan reflect.Value, 256)
+	go func() {
+		for {
+			ch <- reflect.New(t)
+		}
+	}()
+	if cachedValues.m == nil {
+		cachedValues.m = make(map[reflect.Type]chan reflect.Value, 8)
+	}
+	cachedValues.m[t] = ch
+	return <-ch
+}
+
+func (d *Decoder) newValue(t reflect.Type) reflect.Value {
+	if d.flags&usePreallocateValues == 0 {
+		return reflect.New(t)
+	}
+
+	return cachedValue(t)
+}

decode_value.go

@@ -128,12 +128,12 @@ func ptrValueDecoder(typ reflect.Type) decoderFunc {
 	return func(d *Decoder, v reflect.Value) error {
 		if d.hasNilCode() {
 			if !v.IsNil() {
-				v.Set(reflect.Zero(v.Type()))
+				v.Set(d.newValue(typ))
 			}
 			return d.DecodeNil()
 		}
 		if v.IsNil() {
-			v.Set(reflect.New(v.Type().Elem()))
+			v.Set(d.newValue(typ.Elem()))
 		}
 		return decoder(d, v.Elem())
 	}
@@ -155,7 +155,7 @@ func nilAwareDecoder(typ reflect.Type, fn decoderFunc) decoderFunc {
 				return d.decodeNilValue(v)
 			}
 			if v.IsNil() {
-				v.Set(reflect.New(v.Type().Elem()))
+				v.Set(d.newValue(typ.Elem()))
 			}
 			return fn(d, v)
 		}