More extensibility

adapter.go

@@ -100,7 +100,7 @@ func (adapter *Decoder) Buffered() io.Reader {
 func (adapter *Decoder) UseNumber() {
 	cfg := adapter.iter.cfg.configBeforeFrozen
 	cfg.UseNumber = true
-	adapter.iter.cfg = cfg.frozeWithCacheReuse(adapter.iter.cfg.extraExtensions)
+	adapter.iter.cfg = cfg.frozeWithCacheReuse(adapter.iter.cfg)
 }
 
 // DisallowUnknownFields causes the Decoder to return an error when the destination
@@ -109,7 +109,7 @@ func (adapter *Decoder) UseNumber() {
 func (adapter *Decoder) DisallowUnknownFields() {
 	cfg := adapter.iter.cfg.configBeforeFrozen
 	cfg.DisallowUnknownFields = true
-	adapter.iter.cfg = cfg.frozeWithCacheReuse(adapter.iter.cfg.extraExtensions)
+	adapter.iter.cfg = cfg.frozeWithCacheReuse(adapter.iter.cfg)
 }
 
 // NewEncoder same as json.NewEncoder
@@ -134,14 +134,14 @@ func (adapter *Encoder) Encode(val interface{}) error {
 func (adapter *Encoder) SetIndent(prefix, indent string) {
 	config := adapter.stream.cfg.configBeforeFrozen
 	config.IndentionStep = len(indent)
-	adapter.stream.cfg = config.frozeWithCacheReuse(adapter.stream.cfg.extraExtensions)
+	adapter.stream.cfg = config.frozeWithCacheReuse(adapter.stream.cfg)
 }
 
 // SetEscapeHTML escape html by default, set to false to disable
 func (adapter *Encoder) SetEscapeHTML(escapeHTML bool) {
 	config := adapter.stream.cfg.configBeforeFrozen
 	config.EscapeHTML = escapeHTML
-	adapter.stream.cfg = config.frozeWithCacheReuse(adapter.stream.cfg.extraExtensions)
+	adapter.stream.cfg = config.frozeWithCacheReuse(adapter.stream.cfg)
 }
 
 // Valid reports whether data is a valid JSON encoding.

config.go

@@ -113,16 +113,21 @@ func (cfg *frozenConfig) getEncoderFromCache(cacheKey uintptr) ValEncoder {
 
 var cfgCache = concurrent.NewMap()
 
-func getFrozenConfigFromCache(cfg Config) *frozenConfig {
-	obj, found := cfgCache.Load(cfg)
+type cfgKey struct {
+	Config
+	cause *frozenConfig
+}
+
+func getFrozenConfigFromCache(key cfgKey) *frozenConfig {
+	obj, found := cfgCache.Load(key)
 	if found {
 		return obj.(*frozenConfig)
 	}
 	return nil
 }
 
-func addFrozenConfigToCache(cfg Config, frozenConfig *frozenConfig) {
-	cfgCache.Store(cfg, frozenConfig)
+func addFrozenConfigToCache(key cfgKey, frozenConfig *frozenConfig) {
+	cfgCache.Store(key, frozenConfig)
 }
 
 // Froze forge API from config
@@ -166,16 +171,17 @@ func (cfg Config) Froze() API {
 	return api
 }
 
-func (cfg Config) frozeWithCacheReuse(extraExtensions []Extension) *frozenConfig {
-	api := getFrozenConfigFromCache(cfg)
+func (cfg Config) frozeWithCacheReuse(cause *frozenConfig) *frozenConfig {
+	key := cfgKey{cfg, cause}
+	api := getFrozenConfigFromCache(key)
 	if api != nil {
 		return api
 	}
 	api = cfg.Froze().(*frozenConfig)
-	for _, extension := range extraExtensions {
+	for _, extension := range cause.extraExtensions {
 		api.RegisterExtension(extension)
 	}
-	addFrozenConfigToCache(cfg, api)
+	addFrozenConfigToCache(key, api)
 	return api
 }
 
@@ -317,7 +323,7 @@ func (cfg *frozenConfig) MarshalIndent(v interface{}, prefix, indent string) ([]
 	}
 	newCfg := cfg.configBeforeFrozen
 	newCfg.IndentionStep = len(indent)
-	return newCfg.frozeWithCacheReuse(cfg.extraExtensions).Marshal(v)
+	return newCfg.frozeWithCacheReuse(cfg).Marshal(v)
 }
 
 func (cfg *frozenConfig) UnmarshalFromString(str string, v interface{}) error {
@@ -373,3 +379,7 @@ func (cfg *frozenConfig) Valid(data []byte) bool {
 	iter.Skip()
 	return iter.Error == nil
 }
+
+func (cfg *frozenConfig) GetConfig() Config {
+	return cfg.configBeforeFrozen
+}

extension_tests/extension_test.go

@@ -1,13 +1,14 @@
 package test
 
 import (
-	"github.com/json-iterator/go"
-	"github.com/modern-go/reflect2"
-	"github.com/stretchr/testify/require"
 	"reflect"
 	"strconv"
 	"testing"
 	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+	"github.com/stretchr/testify/require"
 )
 
 type TestObject1 struct {
@@ -59,6 +60,19 @@ func Test_customize_map_key_encoder(t *testing.T) {
 	m = map[int]int{}
 	should.NoError(cfg.UnmarshalFromString(output, &m))
 	should.Equal(map[int]int{1: 2}, m)
+
+	b, err := cfg.MarshalIndent(m, "", "  ")
+	should.NoError(err)
+	should.Equal(`{
+  "2": 2
+}`, string(b))
+
+	cfg = jsoniter.Config{}.Froze() // without testMapKeyExtension
+	b, err = cfg.MarshalIndent(m, "", "  ")
+	should.NoError(err)
+	should.Equal(`{
+  "1": 2
+}`, string(b))
 }
 
 type testMapKeyExtension struct {

go.mod

@@ -8,4 +8,5 @@ require (
 	github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421
 	github.com/modern-go/reflect2 v1.0.2
 	github.com/stretchr/testify v1.8.0
+	github.com/valyala/bytebufferpool v1.0.0
 )

go.sum

@@ -14,6 +14,8 @@ github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSS
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
+github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
+github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

iter.go

@@ -122,6 +122,11 @@ func ParseString(cfg API, input string) *Iterator {
 	return ParseBytes(cfg, []byte(input))
 }
 
+// API returns API
+func (iter *Iterator) API() API {
+	return iter.cfg
+}
+
 // Pool returns a pool can provide more iterator with same configuration
 func (iter *Iterator) Pool() IteratorPool {
 	return iter.cfg
@@ -249,6 +254,15 @@ func (iter *Iterator) readByte() (ret byte) {
 	return ret
 }
 
+func (iter *Iterator) NextToken() (ret byte) {
+	return iter.nextToken()
+}
+
+func (iter *Iterator) UnreadByte() error {
+	iter.unreadByte()
+	return nil
+}
+
 func (iter *Iterator) loadMore() bool {
 	if iter.reader == nil {
 		if iter.Error == nil {

iter_str.go

@@ -3,6 +3,8 @@ package jsoniter
 import (
 	"fmt"
 	"unicode/utf16"
+
+	"github.com/valyala/bytebufferpool"
 )
 
 // ReadString read string from iterator
@@ -32,19 +34,24 @@ func (iter *Iterator) ReadString() (ret string) {
 	return
 }
 
+var byteBufPool = bytebufferpool.Pool{}
+
 func (iter *Iterator) readStringSlowPath() (ret string) {
-	var str []byte
+	// reduce runtime.growslice
+	b := byteBufPool.Get()
+	defer byteBufPool.Put(b)
+
 	var c byte
 	for iter.Error == nil {
 		c = iter.readByte()
 		if c == '"' {
-			return string(str)
+			return string(b.B)
 		}
 		if c == '\\' {
 			c = iter.readByte()
-			str = iter.readEscapedChar(c, str)
+			b.B = iter.readEscapedChar(c, b.B)
 		} else {
-			str = append(str, c)
+			b.WriteByte(c)
 		}
 	}
 	iter.ReportError("readStringSlowPath", "unexpected end of input")

reflect_array.go

@@ -2,9 +2,10 @@ package jsoniter
 
 import (
 	"fmt"
-	"github.com/modern-go/reflect2"
 	"io"
 	"unsafe"
+
+	"github.com/modern-go/reflect2"
 )
 
 func decoderOfArray(ctx *ctx, typ reflect2.Type) ValDecoder {
@@ -13,13 +14,43 @@ func decoderOfArray(ctx *ctx, typ reflect2.Type) ValDecoder {
 	return &arrayDecoder{arrayType, decoder}
 }
 
+type ArrayEncoderConstructor struct {
+	ArrayType    *reflect2.UnsafeArrayType
+	ElemEncoder  ValEncoder
+	API          API
+	DecorateFunc func(arrayEncoder ValEncoder) ValEncoder
+}
+
+func updateArrayEncoderConstructor(v *ArrayEncoderConstructor, exts ...Extension) {
+	for _, ext := range exts {
+		if e, ok := ext.(interface {
+			UpdateArrayEncoderConstructor(v *ArrayEncoderConstructor)
+		}); ok {
+			e.UpdateArrayEncoderConstructor(v)
+		}
+	}
+}
+
 func encoderOfArray(ctx *ctx, typ reflect2.Type) ValEncoder {
 	arrayType := typ.(*reflect2.UnsafeArrayType)
 	if arrayType.Len() == 0 {
 		return emptyArrayEncoder{}
 	}
-	encoder := encoderOfType(ctx.append("[arrayElem]"), arrayType.Elem())
-	return &arrayEncoder{arrayType, encoder}
+	elemEncoder := encoderOfType(ctx.append("[arrayElem]"), arrayType.Elem())
+
+	c := &ArrayEncoderConstructor{
+		ArrayType:   arrayType,
+		ElemEncoder: elemEncoder,
+		API:         ctx,
+		DecorateFunc: func(arrayEncoder ValEncoder) ValEncoder {
+			return arrayEncoder
+		},
+	}
+	updateArrayEncoderConstructor(c, extensions...)
+	updateArrayEncoderConstructor(c, ctx.encoderExtension)
+	updateArrayEncoderConstructor(c, ctx.extraExtensions...)
+	enc := &arrayEncoder{arrayType, c.ElemEncoder}
+	return c.DecorateFunc(enc)
 }
 
 type emptyArrayEncoder struct{}