Merge pull request #47 from goccy/feature/support-sorted-map

Support sorted map and change default behavior encoding for map type

cmd/generator/main.go

@@ -204,6 +204,7 @@ func (t opType) fieldToStringTagField() opType {
 		"MapHead",
 		"MapKey",
 		"MapValue",
+		"MapEnd",
 		"StructFieldRecursive",
 		"StructField",
 	}

encode.go

@@ -18,6 +18,7 @@ type Encoder struct {
 	buf                            []byte
 	enabledIndent                  bool
 	enabledHTMLEscape              bool
+	unorderedMap                   bool
 	prefix                         []byte
 	indentStr                      []byte
 	structTypeToCompiledCode       map[uintptr]*compiledCode
@@ -88,6 +89,16 @@ func NewEncoder(w io.Writer) *Encoder {
 //
 // See the documentation for Marshal for details about the conversion of Go values to JSON.
 func (e *Encoder) Encode(v interface{}) error {
+	return e.EncodeWithOption(v)
+}
+
+// EncodeWithOption call Encode with EncodeOption.
+func (e *Encoder) EncodeWithOption(v interface{}, opts ...EncodeOption) error {
+	for _, opt := range opts {
+		if err := opt(e); err != nil {
+			return err
+		}
+	}
 	if err := e.encode(v); err != nil {
 		return err
 	}
@@ -126,6 +137,7 @@ func (e *Encoder) reset() {
 	e.buf = e.buf[:0]
 	e.enabledHTMLEscape = true
 	e.enabledIndent = false
+	e.unorderedMap = false
 }
 
 func (e *Encoder) encodeForMarshal(v interface{}) ([]byte, error) {

encode_compile.go

@@ -398,26 +398,15 @@ func (e *Encoder) compileMap(ctx *encodeCompileContext, withLoad bool) (*opcode,
 
 	header.mapKey = key
 	header.mapValue = value
-	end := newOpCode(ctx, opMapEnd)
+
+	end := newMapEndCode(ctx, header)
 	ctx.incIndex()
 
 	if ctx.withIndent {
-		if header.op == opMapHead {
+		header.op = header.op.toIndent()
-			if ctx.root {
+		key.op = key.op.toIndent()
-				header.op = opRootMapHeadIndent
+		value.op = value.op.toIndent()
-			} else {
+		end.op = end.op.toIndent()
-				header.op = opMapHeadIndent
-			}
-		} else {
-			header.op = opMapHeadLoadIndent
-		}
-		if ctx.root {
-			key.op = opRootMapKeyIndent
-		} else {
-			key.op = opMapKeyIndent
-		}
-		value.op = opMapValueIndent
-		end.op = opMapEndIndent
 	}
 
 	header.next = keyCode
@@ -428,6 +417,7 @@ func (e *Encoder) compileMap(ctx *encodeCompileContext, withLoad bool) (*opcode,
 
 	header.end = end
 	key.end = end
+	value.end = end
 
 	return (*opcode)(unsafe.Pointer(header)), nil
 }

encode_opcode.go

@@ -24,6 +24,7 @@ type opcode struct {
 	elemIdx uintptr // offset to access array/slice/map elem
 	length  uintptr // offset to access slice/map length or array length
 	mapIter uintptr // offset to access map iterator
+	mapPos  uintptr // offset to access position list for sorted map
 	offset  uintptr // offset size from struct header
 	size    uintptr // array/slice elem size
 
@@ -87,6 +88,7 @@ func (c *opcode) copy(codeMap map[uintptr]*opcode) *opcode {
 		elemIdx:      c.elemIdx,
 		length:       c.length,
 		mapIter:      c.mapIter,
+		mapPos:       c.mapPos,
 		offset:       c.offset,
 		size:         c.size,
 	}
@@ -194,6 +196,18 @@ func (c *opcode) dumpMapHead(code *opcode) string {
 	)
 }
 
+func (c *opcode) dumpMapEnd(code *opcode) string {
+	return fmt.Sprintf(
+		`[%d]%s%s ([idx:%d][mapPos:%d][length:%d])`,
+		code.displayIdx,
+		strings.Repeat("-", code.indent),
+		code.op,
+		code.idx/uintptrSize,
+		code.mapPos/uintptrSize,
+		code.length/uintptrSize,
+	)
+}
+
 func (c *opcode) dumpElem(code *opcode) string {
 	var length uintptr
 	if code.op.codeType() == codeArrayElem {
@@ -270,6 +284,9 @@ func (c *opcode) dump() string {
 		case codeMapValue:
 			codes = append(codes, c.dumpValue(code))
 			code = code.next
+		case codeMapEnd:
+			codes = append(codes, c.dumpMapEnd(code))
+			code = code.next
 		case codeStructField:
 			codes = append(codes, c.dumpField(code))
 			code = code.next
@@ -416,6 +433,21 @@ func newMapValueCode(ctx *encodeCompileContext, head *opcode) *opcode {
 	}
 }
 
+func newMapEndCode(ctx *encodeCompileContext, head *opcode) *opcode {
+	mapPos := opcodeOffset(ctx.ptrIndex)
+	ctx.incPtrIndex()
+	idx := opcodeOffset(ctx.ptrIndex)
+	return &opcode{
+		op:         opMapEnd,
+		displayIdx: ctx.opcodeIndex,
+		idx:        idx,
+		length:     head.length,
+		mapPos:     mapPos,
+		indent:     ctx.indent,
+		next:       newEndOp(ctx),
+	}
+}
+
 func newInterfaceCode(ctx *encodeCompileContext) *opcode {
 	return &opcode{
 		op:         opInterface,

encode_optype.go

@@ -12,8 +12,9 @@ const (
   codeMapHead codeType = 5
   codeMapKey codeType = 6
   codeMapValue codeType = 7
-  codeStructFieldRecursive codeType = 8
+  codeMapEnd codeType = 8
-  codeStructField codeType = 9
+  codeStructFieldRecursive codeType = 9
+  codeStructField codeType = 10
 )
 
 type opType int

encode_test.go

@@ -377,14 +377,18 @@ func Test_Marshal(t *testing.T) {
 	})
 	t.Run("map", func(t *testing.T) {
 		t.Run("map[string]int", func(t *testing.T) {
-			bytes, err := json.Marshal(map[string]int{
+			v := map[string]int{
 				"a": 1,
 				"b": 2,
 				"c": 3,
 				"d": 4,
-			})
+			}
+			bytes, err := json.Marshal(v)
 			assertErr(t, err)
-			assertEq(t, "map", len(`{"a":1,"b":2,"c":3,"d":4}`), len(string(bytes)))
+			assertEq(t, "map", `{"a":1,"b":2,"c":3,"d":4}`, string(bytes))
+			b, err := json.MarshalWithOption(v, json.UnorderedMap())
+			assertErr(t, err)
+			assertEq(t, "unordered map", len(`{"a":1,"b":2,"c":3,"d":4}`), len(string(b)))
 		})
 		t.Run("map[string]interface{}", func(t *testing.T) {
 			type T struct {
@@ -400,7 +404,7 @@ func Test_Marshal(t *testing.T) {
 			}
 			bytes, err := json.Marshal(v)
 			assertErr(t, err)
-			assertEq(t, "map[string]interface{}", len(`{"a":1,"b":2.1,"c":{"A":10},"d":4}`), len(string(bytes)))
+			assertEq(t, "map[string]interface{}", `{"a":1,"b":2.1,"c":{"A":10},"d":4}`, string(bytes))
 		})
 	})
 }
@@ -474,7 +478,7 @@ func Test_MarshalIndent(t *testing.T) {
 			}, prefix, indent)
 			assertErr(t, err)
 			result := "{\n-\t\"a\": 1,\n-\t\"b\": 2,\n-\t\"c\": 3,\n-\t\"d\": 4\n-}"
-			assertEq(t, "map", len(result), len(string(bytes)))
+			assertEq(t, "map", result, string(bytes))
 		})
 		t.Run("map[string]interface{}", func(t *testing.T) {
 			type T struct {
@@ -493,7 +497,7 @@ func Test_MarshalIndent(t *testing.T) {
 			bytes, err := json.MarshalIndent(v, prefix, indent)
 			assertErr(t, err)
 			result := "{\n-\t\"a\": 1,\n-\t\"b\": 2.1,\n-\t\"c\": {\n-\t\t\"E\": 10,\n-\t\t\"F\": 11\n-\t},\n-\t\"d\": 4\n-}"
-			assertEq(t, "map[string]interface{}", len(result), len(string(bytes)))
+			assertEq(t, "map[string]interface{}", result, string(bytes))
 		})
 	})
 }
@@ -898,7 +902,7 @@ func TestTextMarshalerMapKeysAreSorted(t *testing.T) {
 		t.Fatalf("Failed to Marshal text.Marshaler: %v", err)
 	}
 	const want = `{"a:z":3,"x:y":1,"y:x":2,"z:a":4}`
-	if len(string(b)) != len(want) {
+	if string(b) != want {
 		t.Errorf("Marshal map with text.Marshaler keys: got %#q, want %#q", b, want)
 	}
 }
@@ -914,7 +918,7 @@ func TestNilMarshalerTextMapKey(t *testing.T) {
 		t.Fatalf("Failed to Marshal *text.Marshaler: %v", err)
 	}
 	const want = `{"":1,"A:B":2}`
-	if len(string(b)) != len(want) {
+	if string(b) != want {
 		t.Errorf("Marshal map with *text.Marshaler keys: got %#q, want %#q", b, want)
 	}
 }

encode_vm.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"math"
 	"reflect"
+	"sort"
 	"strconv"
 	"unsafe"
 )
@@ -439,6 +440,13 @@ func (e *Encoder) run(ctx *encodeRuntimeContext, code *opcode) error {
 					store(ctxptr, code.elemIdx, 0)
 					store(ctxptr, code.length, uintptr(mlen))
 					store(ctxptr, code.mapIter, uintptr(iter))
+					if !e.unorderedMap {
+						pos := make([]int, 0, mlen)
+						pos = append(pos, len(e.buf))
+						posPtr := unsafe.Pointer(&pos)
+						ctx.keepRefs = append(ctx.keepRefs, posPtr)
+						store(ctxptr, code.end.mapPos, uintptr(posPtr))
+					}
 					key := mapiterkey(iter)
 					store(ctxptr, code.next.idx, uintptr(key))
 					code = code.next
@@ -465,6 +473,13 @@ func (e *Encoder) run(ctx *encodeRuntimeContext, code *opcode) error {
 					store(ctxptr, code.mapIter, uintptr(iter))
 					key := mapiterkey(iter)
 					store(ctxptr, code.next.idx, uintptr(key))
+					if !e.unorderedMap {
+						pos := make([]int, 0, mlen)
+						pos = append(pos, len(e.buf))
+						posPtr := unsafe.Pointer(&pos)
+						ctx.keepRefs = append(ctx.keepRefs, posPtr)
+						store(ctxptr, code.end.mapPos, uintptr(posPtr))
+					}
 					code = code.next
 				} else {
 					e.encodeByte('}')
@@ -475,24 +490,84 @@ func (e *Encoder) run(ctx *encodeRuntimeContext, code *opcode) error {
 			idx := load(ctxptr, code.elemIdx)
 			length := load(ctxptr, code.length)
 			idx++
-			if idx < length {
+			if e.unorderedMap {
-				e.encodeByte(',')
+				if idx < length {
-				iter := unsafe.Pointer(load(ctxptr, code.mapIter))
+					e.encodeByte(',')
-				store(ctxptr, code.elemIdx, idx)
+					iter := unsafe.Pointer(load(ctxptr, code.mapIter))
-				key := mapiterkey(iter)
+					store(ctxptr, code.elemIdx, idx)
-				store(ctxptr, code.next.idx, uintptr(key))
+					key := mapiterkey(iter)
-				code = code.next
+					store(ctxptr, code.next.idx, uintptr(key))
+					code = code.next
+				} else {
+					e.encodeByte('}')
+					code = code.end.next
+				}
 			} else {
-				e.encodeByte('}')
+				posPtr := (*[]int)(unsafe.Pointer(load(ctxptr, code.end.mapPos)))
-				code = code.end.next
+				*posPtr = append(*posPtr, len(e.buf))
+				if idx < length {
+					iter := unsafe.Pointer(load(ctxptr, code.mapIter))
+					store(ctxptr, code.elemIdx, idx)
+					key := mapiterkey(iter)
+					store(ctxptr, code.next.idx, uintptr(key))
+					code = code.next
+				} else {
+					code = code.end
+				}
 			}
 		case opMapValue:
-			e.encodeByte(':')
+			if e.unorderedMap {
+				e.encodeByte(':')
+			} else {
+				posPtr := (*[]int)(unsafe.Pointer(load(ctxptr, code.end.mapPos)))
+				*posPtr = append(*posPtr, len(e.buf))
+			}
 			iter := unsafe.Pointer(load(ctxptr, code.mapIter))
 			value := mapitervalue(iter)
 			store(ctxptr, code.next.idx, uintptr(value))
 			mapiternext(iter)
 			code = code.next
+		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)
+			posPtr := unsafe.Pointer(load(ctxptr, code.mapPos))
+			pos := *(*[]int)(posPtr)
+			for i := 0; i < length; i++ {
+				startKey := pos[i*2]
+				startValue := pos[i*2+1]
+				var endValue int
+				if i+1 < length {
+					endValue = pos[i*2+2]
+				} else {
+					endValue = len(e.buf)
+				}
+				kvs = append(kvs, mapKV{
+					key:   string(e.buf[startKey:startValue]),
+					value: string(e.buf[startValue:endValue]),
+				})
+			}
+			sort.Slice(kvs, func(i, j int) bool {
+				return kvs[i].key < kvs[j].key
+			})
+			buf := e.buf[pos[0]:]
+			buf = buf[:0]
+			for idx, kv := range kvs {
+				if idx != 0 {
+					buf = append(buf, ',')
+				}
+				buf = append(buf, []byte(kv.key)...)
+				buf = append(buf, ':')
+				buf = append(buf, []byte(kv.value)...)
+			}
+			buf = append(buf, '}')
+			e.buf = e.buf[:pos[0]]
+			e.buf = append(e.buf, buf...)
+			code = code.next
 		case opMapHeadIndent:
 			ptr := load(ctxptr, code.idx)
 			if ptr == 0 {
@@ -508,10 +583,20 @@ func (e *Encoder) run(ctx *encodeRuntimeContext, code *opcode) error {
 					store(ctxptr, code.elemIdx, 0)
 					store(ctxptr, code.length, uintptr(mlen))
 					store(ctxptr, code.mapIter, uintptr(iter))
+
+					if !e.unorderedMap {
+						pos := make([]int, 0, mlen)
+						pos = append(pos, len(e.buf))
+						posPtr := unsafe.Pointer(&pos)
+						ctx.keepRefs = append(ctx.keepRefs, posPtr)
+						store(ctxptr, code.end.mapPos, uintptr(posPtr))
+					} else {
+						e.encodeIndent(code.next.indent)
+					}
+
 					key := mapiterkey(iter)
 					store(ctxptr, code.next.idx, uintptr(key))
 					code = code.next
-					e.encodeIndent(code.indent)
 				} else {
 					e.encodeIndent(code.indent)
 					e.encodeBytes([]byte{'{', '}'})
@@ -537,8 +622,18 @@ func (e *Encoder) run(ctx *encodeRuntimeContext, code *opcode) error {
 					store(ctxptr, code.mapIter, uintptr(iter))
 					key := mapiterkey(iter)
 					store(ctxptr, code.next.idx, uintptr(key))
+
+					if !e.unorderedMap {
+						pos := make([]int, 0, mlen)
+						pos = append(pos, len(e.buf))
+						posPtr := unsafe.Pointer(&pos)
+						ctx.keepRefs = append(ctx.keepRefs, posPtr)
+						store(ctxptr, code.end.mapPos, uintptr(posPtr))
+					} else {
+						e.encodeIndent(code.next.indent)
+					}
+
 					code = code.next
-					e.encodeIndent(code.indent)
 				} else {
 					e.encodeIndent(code.indent)
 					e.encodeBytes([]byte{'{', '}'})
@@ -574,19 +669,33 @@ func (e *Encoder) run(ctx *encodeRuntimeContext, code *opcode) error {
 			idx := load(ctxptr, code.elemIdx)
 			length := load(ctxptr, code.length)
 			idx++
-			if idx < length {
+			if e.unorderedMap {
-				e.encodeBytes([]byte{',', '\n'})
+				if idx < length {
-				e.encodeIndent(code.indent)
+					e.encodeBytes([]byte{',', '\n'})
-				store(ctxptr, code.elemIdx, idx)
+					e.encodeIndent(code.indent)
-				iter := unsafe.Pointer(load(ctxptr, code.mapIter))
+					store(ctxptr, code.elemIdx, idx)
-				key := mapiterkey(iter)
+					iter := unsafe.Pointer(load(ctxptr, code.mapIter))
-				store(ctxptr, code.next.idx, uintptr(key))
+					key := mapiterkey(iter)
-				code = code.next
+					store(ctxptr, code.next.idx, uintptr(key))
+					code = code.next
+				} else {
+					e.encodeByte('\n')
+					e.encodeIndent(code.indent - 1)
+					e.encodeByte('}')
+					code = code.end.next
+				}
 			} else {
-				e.encodeByte('\n')
+				posPtr := (*[]int)(unsafe.Pointer(load(ctxptr, code.end.mapPos)))
-				e.encodeIndent(code.indent - 1)
+				*posPtr = append(*posPtr, len(e.buf))
-				e.encodeByte('}')
+				if idx < length {
-				code = code.end.next
+					iter := unsafe.Pointer(load(ctxptr, code.mapIter))
+					store(ctxptr, code.elemIdx, idx)
+					key := mapiterkey(iter)
+					store(ctxptr, code.next.idx, uintptr(key))
+					code = code.next
+				} else {
+					code = code.end
+				}
 			}
 		case opRootMapKeyIndent:
 			idx := load(ctxptr, code.elemIdx)
@@ -607,12 +716,64 @@ func (e *Encoder) run(ctx *encodeRuntimeContext, code *opcode) error {
 				code = code.end.next
 			}
 		case opMapValueIndent:
-			e.encodeBytes([]byte{':', ' '})
+			if e.unorderedMap {
+				e.encodeBytes([]byte{':', ' '})
+			} else {
+				posPtr := (*[]int)(unsafe.Pointer(load(ctxptr, code.end.mapPos)))
+				*posPtr = append(*posPtr, len(e.buf))
+			}
 			iter := unsafe.Pointer(load(ctxptr, code.mapIter))
 			value := mapitervalue(iter)
 			store(ctxptr, code.next.idx, uintptr(value))
 			mapiternext(iter)
 			code = code.next
+		case opMapEndIndent:
+			// 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)
+			pos := *(*[]int)(unsafe.Pointer(load(ctxptr, code.mapPos)))
+			for i := 0; i < length; i++ {
+				startKey := pos[i*2]
+				startValue := pos[i*2+1]
+				var endValue int
+				if i+1 < length {
+					endValue = pos[i*2+2]
+				} else {
+					endValue = len(e.buf)
+				}
+				kvs = append(kvs, mapKV{
+					key:   string(e.buf[startKey:startValue]),
+					value: string(e.buf[startValue:endValue]),
+				})
+			}
+			sort.Slice(kvs, func(i, j int) bool {
+				return kvs[i].key < kvs[j].key
+			})
+			buf := e.buf[pos[0]:]
+			buf = buf[:0]
+			for idx, kv := range kvs {
+				if idx != 0 {
+					buf = append(buf, []byte{',', '\n'}...)
+				}
+
+				buf = append(buf, e.prefix...)
+				buf = append(buf, bytes.Repeat(e.indentStr, code.indent+1)...)
+
+				buf = append(buf, []byte(kv.key)...)
+				buf = append(buf, []byte{':', ' '}...)
+				buf = append(buf, []byte(kv.value)...)
+			}
+			buf = append(buf, '\n')
+			buf = append(buf, e.prefix...)
+			buf = append(buf, bytes.Repeat(e.indentStr, code.indent)...)
+			buf = append(buf, '}')
+			e.buf = e.buf[:pos[0]]
+			e.buf = append(e.buf, buf...)
+			code = code.next
 		case opStructFieldRecursive:
 			ptr := load(ctxptr, code.idx)
 			if ptr != 0 {

json.go

@@ -154,8 +154,18 @@ type Unmarshaler interface {
 // an infinite recursion.
 //
 func Marshal(v interface{}) ([]byte, error) {
+	return MarshalWithOption(v)
+}
+
+// MarshalWithOption returns the JSON encoding of v with EncodeOption.
+func MarshalWithOption(v interface{}, opts ...EncodeOption) ([]byte, error) {
 	var b *bytes.Buffer
 	enc := NewEncoder(b)
+	for _, opt := range opts {
+		if err := opt(enc); err != nil {
+			return nil, err
+		}
+	}
 	bytes, err := enc.encodeForMarshal(v)
 	if err != nil {
 		enc.release()
@@ -169,8 +179,18 @@ func Marshal(v interface{}) ([]byte, error) {
 // Each JSON element in the output will begin on a new line beginning with prefix
 // followed by one or more copies of indent according to the indentation nesting.
 func MarshalIndent(v interface{}, prefix, indent string) ([]byte, error) {
+	return MarshalIndentWithOption(v, prefix, indent)
+}
+
+// MarshalIndentWithOption is like Marshal but applies Indent to format the output with EncodeOption.
+func MarshalIndentWithOption(v interface{}, prefix, indent string, opts ...EncodeOption) ([]byte, error) {
 	var b *bytes.Buffer
 	enc := NewEncoder(b)
+	for _, opt := range opts {
+		if err := opt(enc); err != nil {
+			return nil, err
+		}
+	}
 	enc.SetIndent(prefix, indent)
 	bytes, err := enc.encodeForMarshal(v)
 	if err != nil {

option.go

@@ -0,0 +1,10 @@
+package json
+
+type EncodeOption func(*Encoder) error
+
+func UnorderedMap() EncodeOption {
+	return func(e *Encoder) error {
+		e.unorderedMap = true
+		return nil
+	}
+}