Improve performance of encoding map type ( escaped and not indented )

encode_context.go

@@ -1,9 +1,67 @@
 package json
 
 import (
+	"bytes"
+	"sync"
 	"unsafe"
 )
 
+type mapItem struct {
+	key   []byte
+	value []byte
+}
+
+type mapslice struct {
+	items []mapItem
+}
+
+func (m *mapslice) Len() int {
+	return len(m.items)
+}
+
+func (m *mapslice) Less(i, j int) bool {
+	return bytes.Compare(m.items[i].key, m.items[j].key) < 0
+}
+
+func (m *mapslice) Swap(i, j int) {
+	m.items[i], m.items[j] = m.items[j], m.items[i]
+}
+
+type encodeMapContext struct {
+	iter  unsafe.Pointer
+	pos   []int
+	slice *mapslice
+	buf   []byte
+}
+
+var mapContextPool = sync.Pool{
+	New: func() interface{} {
+		return &encodeMapContext{}
+	},
+}
+
+func newMapContext(mapLen int) *encodeMapContext {
+	ctx := mapContextPool.Get().(*encodeMapContext)
+	if ctx.slice == nil {
+		ctx.slice = &mapslice{
+			items: make([]mapItem, 0, mapLen),
+		}
+	}
+	if cap(ctx.pos) < (mapLen*2 + 1) {
+		ctx.pos = make([]int, 0, mapLen*2+1)
+		ctx.slice.items = make([]mapItem, 0, mapLen)
+	} else {
+		ctx.pos = ctx.pos[:0]
+		ctx.slice.items = ctx.slice.items[:0]
+	}
+	ctx.buf = ctx.buf[:0]
+	return ctx
+}
+
+func releaseMapContext(c *encodeMapContext) {
+	mapContextPool.Put(c)
+}
+
 type encodeCompileContext struct {
 	typ                      *rtype
 	root                     bool

encode_vm_escaped.go

@@ -341,11 +341,10 @@ func (e *Encoder) runEscaped(ctx *encodeRuntimeContext, b []byte, codeSet *opcod
 					store(ctxptr, code.length, uintptr(mlen))
 					store(ctxptr, code.mapIter, uintptr(iter))
 					if !e.unorderedMap {
-						pos := make([]int, 0, mlen)
-						pos = append(pos, len(b))
-						posPtr := unsafe.Pointer(&pos)
-						ctx.keepRefs = append(ctx.keepRefs, posPtr)
-						store(ctxptr, code.end.mapPos, uintptr(posPtr))
+						mapCtx := newMapContext(mlen)
+						mapCtx.pos = append(mapCtx.pos, len(b))
+						ctx.keepRefs = append(ctx.keepRefs, unsafe.Pointer(mapCtx))
+						store(ctxptr, code.end.mapPos, uintptr(unsafe.Pointer(mapCtx)))
 					}
 					key := mapiterkey(iter)
 					store(ctxptr, code.next.idx, uintptr(key))
@@ -382,11 +381,9 @@ func (e *Encoder) runEscaped(ctx *encodeRuntimeContext, b []byte, codeSet *opcod
 					key := mapiterkey(iter)
 					store(ctxptr, code.next.idx, uintptr(key))
 					if !e.unorderedMap {
-						pos := make([]int, 0, mlen)
-						pos = append(pos, len(b))
-						posPtr := unsafe.Pointer(&pos)
-						ctx.keepRefs = append(ctx.keepRefs, posPtr)
-						store(ctxptr, code.end.mapPos, uintptr(posPtr))
+						mapCtx := newMapContext(mlen)
+						mapCtx.pos = append(mapCtx.pos, len(b))
+						store(ctxptr, code.end.mapPos, uintptr(unsafe.Pointer(mapCtx)))
 					}
 					code = code.next
 				} else {
@@ -414,8 +411,8 @@ func (e *Encoder) runEscaped(ctx *encodeRuntimeContext, b []byte, codeSet *opcod
 				}
 			} else {
 				ptr := load(ctxptr, code.end.mapPos)
-				posPtr := (*[]int)(*(*unsafe.Pointer)(unsafe.Pointer(&ptr)))
-				*posPtr = append(*posPtr, len(b))
+				mapCtx := (*encodeMapContext)(e.ptrToUnsafePtr(ptr))
+				mapCtx.pos = append(mapCtx.pos, len(b))
 				if idx < length {
 					ptr := load(ctxptr, code.mapIter)
 					iter := e.ptrToUnsafePtr(ptr)
@@ -433,8 +430,8 @@ func (e *Encoder) runEscaped(ctx *encodeRuntimeContext, b []byte, codeSet *opcod
 				b[last] = ':'
 			} else {
 				ptr := load(ctxptr, code.end.mapPos)
-				posPtr := (*[]int)(*(*unsafe.Pointer)(unsafe.Pointer(&ptr)))
-				*posPtr = append(*posPtr, len(b))
+				mapCtx := (*encodeMapContext)(e.ptrToUnsafePtr(ptr))
+				mapCtx.pos = append(mapCtx.pos, len(b))
 			}
 			ptr := load(ctxptr, code.mapIter)
 			iter := e.ptrToUnsafePtr(ptr)
@@ -445,14 +442,9 @@ func (e *Encoder) runEscaped(ctx *encodeRuntimeContext, b []byte, codeSet *opcod
 		case opMapEnd:
 			// this operation only used by sorted map.
 			length := int(load(ctxptr, code.length))
-			type mapKV struct {
-				key   string
-				value string
-			}
-			kvs := make([]mapKV, 0, length)
 			ptr := load(ctxptr, code.mapPos)
-			posPtr := e.ptrToUnsafePtr(ptr)
-			pos := *(*[]int)(posPtr)
+			mapCtx := (*encodeMapContext)(e.ptrToUnsafePtr(ptr))
+			pos := mapCtx.pos
 			for i := 0; i < length; i++ {
 				startKey := pos[i*2]
 				startValue := pos[i*2+1]
@@ -462,25 +454,22 @@ func (e *Encoder) runEscaped(ctx *encodeRuntimeContext, b []byte, codeSet *opcod
 				} else {
 					endValue = len(b)
 				}
-				kvs = append(kvs, mapKV{
-					key:   string(b[startKey:startValue]),
-					value: string(b[startValue:endValue]),
+				mapCtx.slice.items = append(mapCtx.slice.items, mapItem{
+					key:   b[startKey:startValue],
+					value: b[startValue:endValue],
 				})
 			}
-			sort.Slice(kvs, func(i, j int) bool {
-				return kvs[i].key < kvs[j].key
-			})
-			buf := b[pos[0]:]
-			buf = buf[:0]
-			for _, kv := range kvs {
-				buf = append(buf, []byte(kv.key)...)
-				buf[len(buf)-1] = ':'
-				buf = append(buf, []byte(kv.value)...)
+			sort.Sort(mapCtx.slice)
+			for _, item := range mapCtx.slice.items {
+				mapCtx.buf = append(mapCtx.buf, item.key...)
+				mapCtx.buf[len(mapCtx.buf)-1] = ':'
+				mapCtx.buf = append(mapCtx.buf, item.value...)
 			}
-			buf[len(buf)-1] = '}'
-			buf = append(buf, ',')
+			mapCtx.buf[len(mapCtx.buf)-1] = '}'
+			mapCtx.buf = append(mapCtx.buf, ',')
 			b = b[:pos[0]]
-			b = append(b, buf...)
+			b = append(b, mapCtx.buf...)
+			releaseMapContext(mapCtx)
 			code = code.next
 		case opStructFieldPtrAnonymousHeadRecursive:
 			store(ctxptr, code.idx, e.ptrToPtr(load(ctxptr, code.idx)))