From 399354d64d4c6eac4923cdac52c34c0787ff4396 Mon Sep 17 00:00:00 2001
From: Masaaki Goshima <goccy54@gmail.com>
Date: Sun, 24 Jan 2021 15:17:39 +0900
Subject: [PATCH] Fix encoding runner interface

---
 encode.go                   | 55 +++++++++++++++++--------------------
 encode_vm.go                |  3 +-
 encode_vm_escaped.go        |  4 ++-
 encode_vm_escaped_indent.go |  3 +-
 encode_vm_indent.go         |  3 +-
 5 files changed, 34 insertions(+), 34 deletions(-)

diff --git a/encode.go b/encode.go
index 50cb56a..cf4cdaa 100644
--- a/encode.go
+++ b/encode.go
@@ -112,6 +112,7 @@ func (e *Encoder) EncodeWithOption(v interface{}, opts ...EncodeOption) error {
 		}
 	}
 	header := (*interfaceHeader)(unsafe.Pointer(&v))
+	e.ptr = header.ptr
 	buf, err := e.encode(header, v == nil)
 	if err != nil {
 		return err
@@ -192,23 +193,31 @@ func (e *Encoder) encode(header *interfaceHeader, isNil bool) ([]byte, error) {
 	typ := header.typ
 
 	typeptr := uintptr(unsafe.Pointer(typ))
+	codeSet, err := e.compileToGetCodeSet(typeptr)
+	if err != nil {
+		return nil, err
+	}
+
+	ctx := e.ctx
+	p := uintptr(header.ptr)
+	ctx.init(p, codeSet.codeLength)
+	if e.enabledIndent {
+		if e.enabledHTMLEscape {
+			return e.runEscapedIndent(ctx, b, codeSet)
+		} else {
+			return e.runIndent(ctx, b, codeSet)
+		}
+	}
+	if e.enabledHTMLEscape {
+		return e.runEscaped(ctx, b, codeSet)
+	}
+	return e.run(ctx, b, codeSet)
+}
+
+func (e *Encoder) compileToGetCodeSet(typeptr uintptr) (*opcodeSet, error) {
 	opcodeMap := loadOpcodeMap()
 	if codeSet, exists := opcodeMap[typeptr]; exists {
-		ctx := e.ctx
-		p := uintptr(header.ptr)
-		ctx.init(p, codeSet.codeLength)
-
-		if e.enabledIndent {
-			if e.enabledHTMLEscape {
-				return e.runEscapedIndent(ctx, b, codeSet.code)
-			} else {
-				return e.runIndent(ctx, b, codeSet.code)
-			}
-		}
-		if e.enabledHTMLEscape {
-			return e.runEscaped(ctx, b, codeSet.code)
-		}
-		return e.run(ctx, b, codeSet.code)
+		return codeSet, nil
 	}
 
 	// noescape trick for header.typ ( reflect.*rtype )
@@ -230,21 +239,7 @@ func (e *Encoder) encode(header *interfaceHeader, isNil bool) ([]byte, error) {
 	}
 
 	storeOpcodeSet(typeptr, codeSet, opcodeMap)
-	p := uintptr(header.ptr)
-	ctx := e.ctx
-	ctx.init(p, codeLength)
-
-	if e.enabledIndent {
-		if e.enabledHTMLEscape {
-			return e.runEscapedIndent(ctx, b, codeSet.code)
-		} else {
-			return e.runIndent(ctx, b, codeSet.code)
-		}
-	}
-	if e.enabledHTMLEscape {
-		return e.runEscaped(ctx, b, codeSet.code)
-	}
-	return e.run(ctx, b, codeSet.code)
+	return codeSet, nil
 }
 
 func encodeFloat32(b []byte, v float32) []byte {
diff --git a/encode_vm.go b/encode_vm.go
index 14a755c..fe05b32 100644
--- a/encode_vm.go
+++ b/encode_vm.go
@@ -49,10 +49,11 @@ func errMarshaler(code *opcode, err error) *MarshalerError {
 	}
 }
 
-func (e *Encoder) run(ctx *encodeRuntimeContext, b []byte, code *opcode) ([]byte, error) {
+func (e *Encoder) run(ctx *encodeRuntimeContext, b []byte, codeSet *opcodeSet) ([]byte, error) {
 	recursiveLevel := 0
 	ptrOffset := uintptr(0)
 	ctxptr := ctx.ptr()
+	code := codeSet.code
 
 	for {
 		switch code.op {
diff --git a/encode_vm_escaped.go b/encode_vm_escaped.go
index 619c21a..3d59fcf 100644
--- a/encode_vm_escaped.go
+++ b/encode_vm_escaped.go
@@ -11,11 +11,13 @@ import (
 	"unsafe"
 )
 
-func (e *Encoder) runEscaped(ctx *encodeRuntimeContext, b []byte, code *opcode) ([]byte, error) {
+func (e *Encoder) runEscaped(ctx *encodeRuntimeContext, b []byte, codeSet *opcodeSet) ([]byte, error) {
 	recursiveLevel := 0
 	ptrOffset := uintptr(0)
 	ctxptr := ctx.ptr()
 
+	code := codeSet.code
+
 	for {
 		switch code.op {
 		default:
diff --git a/encode_vm_escaped_indent.go b/encode_vm_escaped_indent.go
index 319e812..2cfd745 100644
--- a/encode_vm_escaped_indent.go
+++ b/encode_vm_escaped_indent.go
@@ -11,11 +11,12 @@ import (
 	"unsafe"
 )
 
-func (e *Encoder) runEscapedIndent(ctx *encodeRuntimeContext, b []byte, code *opcode) ([]byte, error) {
+func (e *Encoder) runEscapedIndent(ctx *encodeRuntimeContext, b []byte, codeSet *opcodeSet) ([]byte, error) {
 	recursiveLevel := 0
 	var seenPtr map[uintptr]struct{}
 	ptrOffset := uintptr(0)
 	ctxptr := ctx.ptr()
+	code := codeSet.code
 
 	for {
 		switch code.op {
diff --git a/encode_vm_indent.go b/encode_vm_indent.go
index 536e962..78a6379 100644
--- a/encode_vm_indent.go
+++ b/encode_vm_indent.go
@@ -11,11 +11,12 @@ import (
 	"unsafe"
 )
 
-func (e *Encoder) runIndent(ctx *encodeRuntimeContext, b []byte, code *opcode) ([]byte, error) {
+func (e *Encoder) runIndent(ctx *encodeRuntimeContext, b []byte, codeSet *opcodeSet) ([]byte, error) {
 	recursiveLevel := 0
 	var seenPtr map[uintptr]struct{}
 	ptrOffset := uintptr(0)
 	ctxptr := ctx.ptr()
+	code := codeSet.code
 
 	for {
 		switch code.op {