Supported sorted map

cmd/generator/main.go

@@ -204,6 +204,7 @@ func (t opType) fieldToStringTagField() opType {
 		"MapHead",
 		"MapKey",
 		"MapValue",
+		"MapEnd",
 		"StructFieldRecursive",
 		"StructField",
 	}
@@ -243,7 +244,7 @@ func (t opType) fieldToStringTagField() opType {
 		{"SortedMapKey", "SortedMapKeyIndent", "MapKey"},
 		{"SortedRootMapKey", "SortedRootMapKeyIndent", "MapKey"},
 		{"SortedMapValue", "SortedMapValueIndent", "MapValue"},
-		{"SortedMapEnd", "SortedMapEndIndent", "Op"},
+		{"SortedMapEnd", "SortedMapEndIndent", "MapEnd"},
 		{"StructFieldHead", "StructFieldHeadIndent", "StructField"},
 		{"StructFieldHeadOmitEmpty", "StructFieldHeadOmitEmptyIndent", "StructField"},
 		{"StructFieldHeadStringTag", "StructFieldHeadStringTagIndent", "StructField"},

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
 }
@@ -462,13 +452,13 @@ func (e *Encoder) typeToHeaderType(op opType) opType {
 		return opStructFieldHeadString
 	case opBool:
 		return opStructFieldHeadBool
-	case opMapHead:
+	case opMapHead, opSortedMapHead:
 		return opStructFieldHeadMap
-	case opMapHeadLoad:
+	case opMapHeadLoad, opSortedMapHeadLoad:
 		return opStructFieldHeadMapLoad
-	case opMapHeadIndent:
+	case opMapHeadIndent, opSortedMapHeadIndent:
 		return opStructFieldHeadMapIndent
-	case opMapHeadLoadIndent:
+	case opMapHeadLoadIndent, opSortedMapHeadLoadIndent:
 		return opStructFieldHeadMapLoadIndent
 	case opArrayHead:
 		return opStructFieldHeadArray
@@ -520,13 +510,13 @@ func (e *Encoder) typeToFieldType(op opType) opType {
 		return opStructFieldString
 	case opBool:
 		return opStructFieldBool
-	case opMapHead:
+	case opMapHead, opSortedMapHead:
 		return opStructFieldMap
-	case opMapHeadLoad:
+	case opMapHeadLoad, opSortedMapHeadLoad:
 		return opStructFieldMapLoad
-	case opMapHeadIndent:
+	case opMapHeadIndent, opSortedMapHeadIndent:
 		return opStructFieldMapIndent
-	case opMapHeadLoadIndent:
+	case opMapHeadLoadIndent, opSortedMapHeadLoadIndent:
 		return opStructFieldMapLoadIndent
 	case opArrayHead:
 		return opStructFieldArray

encode_opcode.go

@@ -24,6 +24,8 @@ 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
+	headPos uintptr // offset to access head position of map
+	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 +89,8 @@ func (c *opcode) copy(codeMap map[uintptr]*opcode) *opcode {
 		elemIdx:      c.elemIdx,
 		length:       c.length,
 		mapIter:      c.mapIter,
+		headPos:      c.headPos,
+		mapPos:       c.mapPos,
 		offset:       c.offset,
 		size:         c.size,
 	}
@@ -194,6 +198,19 @@ func (c *opcode) dumpMapHead(code *opcode) string {
 	)
 }
 
+func (c *opcode) dumpMapEnd(code *opcode) string {
+	return fmt.Sprintf(
+		`[%d]%s%s ([idx:%d][headPos:%d][mapPos:%d][length:%d])`,
+		code.displayIdx,
+		strings.Repeat("-", code.indent),
+		code.op,
+		code.idx/uintptrSize,
+		code.headPos/uintptrSize,
+		code.mapPos/uintptrSize,
+		code.length/uintptrSize,
+	)
+}
+
 func (c *opcode) dumpElem(code *opcode) string {
 	var length uintptr
 	if code.op.codeType() == codeArrayElem {
@@ -270,6 +287,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
@@ -369,9 +389,9 @@ func newArrayElemCode(ctx *encodeCompileContext, head *opcode, length int, size
 func newMapHeaderCode(ctx *encodeCompileContext, withLoad bool) *opcode {
 	var op opType
 	if withLoad {
-		op = opMapHeadLoad
+		op = opSortedMapHeadLoad
 	} else {
-		op = opMapHead
+		op = opSortedMapHead
 	}
 	idx := opcodeOffset(ctx.ptrIndex)
 	ctx.incPtrIndex()
@@ -394,7 +414,7 @@ func newMapHeaderCode(ctx *encodeCompileContext, withLoad bool) *opcode {
 
 func newMapKeyCode(ctx *encodeCompileContext, head *opcode) *opcode {
 	return &opcode{
-		op:         opMapKey,
+		op:         opSortedMapKey,
 		displayIdx: ctx.opcodeIndex,
 		idx:        opcodeOffset(ctx.ptrIndex),
 		elemIdx:    head.elemIdx,
@@ -406,7 +426,7 @@ func newMapKeyCode(ctx *encodeCompileContext, head *opcode) *opcode {
 
 func newMapValueCode(ctx *encodeCompileContext, head *opcode) *opcode {
 	return &opcode{
-		op:         opMapValue,
+		op:         opSortedMapValue,
 		displayIdx: ctx.opcodeIndex,
 		idx:        opcodeOffset(ctx.ptrIndex),
 		elemIdx:    head.elemIdx,
@@ -416,6 +436,24 @@ func newMapValueCode(ctx *encodeCompileContext, head *opcode) *opcode {
 	}
 }
 
+func newMapEndCode(ctx *encodeCompileContext, head *opcode) *opcode {
+	headPos := opcodeOffset(ctx.ptrIndex)
+	ctx.incPtrIndex()
+	mapPos := opcodeOffset(ctx.ptrIndex)
+	ctx.incPtrIndex()
+	idx := opcodeOffset(ctx.ptrIndex)
+	return &opcode{
+		op:         opSortedMapEnd,
+		displayIdx: ctx.opcodeIndex,
+		idx:        idx,
+		length:     head.length,
+		headPos:    headPos,
+		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
@@ -2462,7 +2463,7 @@ func (t opType) codeType() codeType {
   case opSortedMapValue:
     return codeMapValue
   case opSortedMapEnd:
-    return codeOp
+    return codeMapEnd
   case opStructFieldHead:
     return codeStructField
   case opStructFieldHeadOmitEmpty:
@@ -3256,7 +3257,7 @@ func (t opType) codeType() codeType {
   case opSortedMapValueIndent:
     return codeMapValue
   case opSortedMapEndIndent:
-    return codeOp
+    return codeMapEnd
   case opStructFieldHeadIndent:
     return codeStructField
   case opStructFieldHeadOmitEmptyIndent:

encode_test.go

@@ -384,7 +384,7 @@ func Test_Marshal(t *testing.T) {
 				"d": 4,
 			})
 			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))
 		})
 		t.Run("map[string]interface{}", func(t *testing.T) {
 			type T struct {
@@ -400,7 +400,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 +474,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 +493,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 +898,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 +914,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"
 )
@@ -425,6 +426,125 @@ func (e *Encoder) run(ctx *encodeRuntimeContext, code *opcode) error {
 				e.encodeBytes([]byte{']', '\n'})
 				code = code.end.next
 			}
+		case opSortedMapHead:
+			ptr := load(ctxptr, code.idx)
+			if ptr == 0 {
+				e.encodeNull()
+				code = code.end.next
+			} else {
+				e.encodeByte('{')
+				mlen := maplen(unsafe.Pointer(ptr))
+				if mlen > 0 {
+					iter := mapiterinit(code.typ, unsafe.Pointer(ptr))
+					ctx.keepRefs = append(ctx.keepRefs, iter)
+					store(ctxptr, code.elemIdx, 0)
+					store(ctxptr, code.length, uintptr(mlen))
+					store(ctxptr, code.mapIter, uintptr(iter))
+					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
+				} else {
+					e.encodeByte('}')
+					code = code.end.next
+				}
+			}
+		case opSortedMapHeadLoad:
+			ptr := load(ctxptr, code.idx)
+			if ptr == 0 {
+				e.encodeNull()
+				code = code.end.next
+			} else {
+				// load pointer
+				ptr = uintptr(*(*unsafe.Pointer)(unsafe.Pointer(ptr)))
+				e.encodeByte('{')
+				mlen := maplen(unsafe.Pointer(ptr))
+				if mlen > 0 {
+					iter := mapiterinit(code.typ, unsafe.Pointer(ptr))
+					ctx.keepRefs = append(ctx.keepRefs, iter)
+					store(ctxptr, code.elemIdx, 0)
+					store(ctxptr, code.length, uintptr(mlen))
+					store(ctxptr, code.mapIter, uintptr(iter))
+					key := mapiterkey(iter)
+					store(ctxptr, code.next.idx, uintptr(key))
+					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('}')
+					code = code.end.next
+				}
+			}
+		case opSortedMapKey:
+			idx := load(ctxptr, code.elemIdx)
+			length := load(ctxptr, code.length)
+			posPtr := (*[]int)(unsafe.Pointer(load(ctxptr, code.end.mapPos)))
+			*posPtr = append(*posPtr, len(e.buf))
+			idx++
+			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 opSortedMapValue:
+			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 opSortedMapEnd:
+			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 opMapHead:
 			ptr := load(ctxptr, code.idx)
 			if ptr == 0 {
@@ -493,6 +613,136 @@ func (e *Encoder) run(ctx *encodeRuntimeContext, code *opcode) error {
 			store(ctxptr, code.next.idx, uintptr(value))
 			mapiternext(iter)
 			code = code.next
+		case opSortedMapHeadIndent:
+			ptr := load(ctxptr, code.idx)
+			if ptr == 0 {
+				e.encodeNull()
+				code = code.end.next
+			} else {
+				mlen := maplen(unsafe.Pointer(ptr))
+				if mlen > 0 {
+					e.encodeBytes([]byte{'{', '\n'})
+
+					iter := mapiterinit(code.typ, unsafe.Pointer(ptr))
+					ctx.keepRefs = append(ctx.keepRefs, iter)
+					store(ctxptr, code.elemIdx, 0)
+					store(ctxptr, code.length, uintptr(mlen))
+					store(ctxptr, code.mapIter, uintptr(iter))
+					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
+				} else {
+					e.encodeIndent(code.indent)
+					e.encodeBytes([]byte{'{', '}'})
+					code = code.end.next
+				}
+			}
+		case opSortedMapHeadLoadIndent:
+			ptr := load(ctxptr, code.idx)
+			if ptr == 0 {
+				e.encodeNull()
+				code = code.end.next
+			} else {
+				// load pointer
+				ptr = uintptr(*(*unsafe.Pointer)(unsafe.Pointer(ptr)))
+				mlen := maplen(unsafe.Pointer(ptr))
+				if mlen > 0 {
+					e.encodeBytes([]byte{'{', '\n'})
+
+					iter := mapiterinit(code.typ, unsafe.Pointer(ptr))
+					ctx.keepRefs = append(ctx.keepRefs, iter)
+					store(ctxptr, code.elemIdx, 0)
+					store(ctxptr, code.length, uintptr(mlen))
+					store(ctxptr, code.mapIter, uintptr(iter))
+					key := mapiterkey(iter)
+					store(ctxptr, code.next.idx, uintptr(key))
+					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.encodeIndent(code.indent)
+					e.encodeBytes([]byte{'{', '}'})
+					code = code.end.next
+				}
+			}
+		case opSortedMapKeyIndent:
+			idx := load(ctxptr, code.elemIdx)
+			length := load(ctxptr, code.length)
+			posPtr := (*[]int)(unsafe.Pointer(load(ctxptr, code.end.mapPos)))
+			*posPtr = append(*posPtr, len(e.buf))
+			idx++
+			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 opSortedMapValueIndent:
+			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 opSortedMapEndIndent:
+			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]),
+				})
+			}
+			fmt.Println("kvs = ", kvs)
+			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 opMapHeadIndent:
 			ptr := load(ctxptr, code.idx)
 			if ptr == 0 {