From cb194687a0a04f65d05944a062cf49bfeef5afac Mon Sep 17 00:00:00 2001
From: Masaaki Goshima <goccy54@gmail.com>
Date: Sun, 30 Aug 2020 17:52:59 +0900
Subject: [PATCH] Add encodeRuntimeContext

---
 encode.go         | 17 ++++++++++++++---
 encode_context.go | 16 ++++++++++++++++
 encode_opcode.go  | 18 ++++++++++++++++++
 encode_vm.go      |  4 ++--
 4 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/encode.go b/encode.go
index 05b15d9..ed77942 100644
--- a/encode.go
+++ b/encode.go
@@ -38,6 +38,7 @@ type opcodeMap struct {
 type opcodeSet struct {
 	codeIndent sync.Pool
 	code       sync.Pool
+	ctx        sync.Pool
 }
 
 func (m *opcodeMap) get(k uintptr) *opcodeSet {
@@ -161,9 +162,11 @@ func (e *Encoder) encode(v interface{}) error {
 		} else {
 			code = codeSet.code.Get().(*opcode)
 		}
+		ctx := codeSet.ctx.Get().(encodeRuntimeContext)
 		p := uintptr(header.ptr)
+		ctx.init(p)
 		code.ptr = p
-		if err := e.run(code); err != nil {
+		if err := e.run(ctx, code); err != nil {
 			return err
 		}
 		if e.enabledIndent {
@@ -193,6 +196,7 @@ func (e *Encoder) encode(v interface{}) error {
 	if err != nil {
 		return err
 	}
+	codeLength := code.length()
 	codeSet := &opcodeSet{
 		codeIndent: sync.Pool{
 			New: func() interface{} {
@@ -204,15 +208,22 @@ func (e *Encoder) encode(v interface{}) error {
 				return copyOpcode(code)
 			},
 		},
+		ctx: sync.Pool{
+			New: func() interface{} {
+				return make(encodeRuntimeContext, codeLength)
+			},
+		},
 	}
 	cachedOpcode.set(typeptr, codeSet)
 	p := uintptr(header.ptr)
+	ctx := codeSet.ctx.Get().(encodeRuntimeContext)
+	ctx.init(p)
 	if e.enabledIndent {
 		codeIndent.ptr = p
-		return e.run(codeIndent)
+		return e.run(ctx, codeIndent)
 	}
 	code.ptr = p
-	return e.run(code)
+	return e.run(ctx, code)
 }
 
 func (e *Encoder) encodeInt(v int) {
diff --git a/encode_context.go b/encode_context.go
index 6c6f4a0..458bb4c 100644
--- a/encode_context.go
+++ b/encode_context.go
@@ -1,5 +1,10 @@
 package json
 
+import (
+	"reflect"
+	"unsafe"
+)
+
 type encodeCompileContext struct {
 	typ         *rtype
 	withIndent  bool
@@ -52,3 +57,14 @@ func (c *encodeCompileContext) decOpcodeIndex() {
 		c.parent.decOpcodeIndex()
 	}
 }
+
+type encodeRuntimeContext []uintptr
+
+func (c *encodeRuntimeContext) init(p uintptr) {
+	(*c)[0] = p
+}
+
+func (c *encodeRuntimeContext) ptr() *uintptr {
+	header := (*reflect.SliceHeader)(unsafe.Pointer(c))
+	return (*uintptr)(unsafe.Pointer(&header.Data))
+}
diff --git a/encode_opcode.go b/encode_opcode.go
index 6740f67..04f7c23 100644
--- a/encode_opcode.go
+++ b/encode_opcode.go
@@ -78,6 +78,24 @@ func (c *opcode) beforeLastCode() *opcode {
 	return nil
 }
 
+func (c *opcode) length() int {
+	var idx int
+	for code := c; code.op != opEnd; {
+		idx = code.idx
+		switch code.op.codeType() {
+		case codeArrayElem:
+			code = code.toArrayElemCode().end
+		case codeSliceElem:
+			code = code.toSliceElemCode().end
+		case codeMapKey:
+			code = code.toMapKeyCode().end
+		default:
+			code = code.next
+		}
+	}
+	return idx + 1
+}
+
 func (c *opcode) decOpcodeIndex() {
 	for code := c; code.op != opEnd; {
 		code.idx--
diff --git a/encode_vm.go b/encode_vm.go
index 3f8f012..b8d0dbf 100644
--- a/encode_vm.go
+++ b/encode_vm.go
@@ -11,7 +11,7 @@ import (
 	"unsafe"
 )
 
-func (e *Encoder) run(code *opcode) error {
+func (e *Encoder) run(ctx encodeRuntimeContext, code *opcode) error {
 	seenPtr := map[uintptr]struct{}{}
 	for {
 		switch code.op {
@@ -546,7 +546,7 @@ func (e *Encoder) run(code *opcode) error {
 				}
 			}
 			recursive.seenPtr = recursive.ptr
-			if err := e.run(newRecursiveCode(recursive)); err != nil {
+			if err := e.run(nil, newRecursiveCode(recursive)); err != nil {
 				return err
 			}
 			code = recursive.next