From 7007d6ee410970d9cc8e84dbe795b593212637e8 Mon Sep 17 00:00:00 2001 From: Masaaki Goshima Date: Fri, 2 Apr 2021 14:03:00 +0900 Subject: [PATCH] Fix encoding of recursive slice/map --- cover_map_test.go | 37 ++++++++++++++++++++++++ cover_slice_test.go | 32 ++++++++++++++++++++ internal/encoder/compiler.go | 23 +++++++++++---- internal/encoder/vm_escaped_indent/vm.go | 4 +++ internal/encoder/vm_indent/vm.go | 4 +++ 5 files changed, 95 insertions(+), 5 deletions(-) diff --git a/cover_map_test.go b/cover_map_test.go index cc4f0e7..2cf08d2 100644 --- a/cover_map_test.go +++ b/cover_map_test.go @@ -7,6 +7,11 @@ import ( "github.com/goccy/go-json" ) +type recursiveMap struct { + A int + B map[string]*recursiveMap +} + func TestCoverMap(t *testing.T) { type structMap struct { A map[string]int `json:"a"` @@ -54,6 +59,38 @@ func TestCoverMap(t *testing.T) { name: "NestedMap", data: map[string]map[string]int{"a": {"b": 1}}, }, + { + name: "RecursiveMap", + data: map[string]*recursiveMap{ + "keyA": { + A: 1, + B: map[string]*recursiveMap{ + "keyB": { + A: 2, + B: map[string]*recursiveMap{ + "keyC": { + A: 3, + }, + }, + }, + }, + }, + "keyD": { + A: 4, + B: map[string]*recursiveMap{ + "keyE": { + A: 5, + B: map[string]*recursiveMap{ + "keyF": { + A: 6, + }, + }, + }, + }, + }, + }, + }, + // HeadMapZero { name: "HeadMapZero", diff --git a/cover_slice_test.go b/cover_slice_test.go index e2dc015..eb60d94 100644 --- a/cover_slice_test.go +++ b/cover_slice_test.go @@ -23,6 +23,11 @@ func (coverSliceMarshalText) MarshalText() ([]byte, error) { return []byte(`"hello"`), nil } +type recursiveSlice struct { + A int + B []*recursiveSlice +} + func TestCoverSlice(t *testing.T) { type structSlice struct { A []int `json:"a"` @@ -226,6 +231,33 @@ func TestCoverSlice(t *testing.T) { name: "SliceStructPtr", data: []*struct{ A int }{&struct{ A int }{A: 1}, &struct{ A int }{A: 2}}, }, + { + name: "RecursiveSlice", + data: []*recursiveSlice{ + { + A: 1, B: []*recursiveSlice{ + { + A: 2, B: []*recursiveSlice{ + { + A: 3, + }, + }, + }, + }, + }, + { + A: 4, B: []*recursiveSlice{ + { + A: 5, B: []*recursiveSlice{ + { + A: 6, + }, + }, + }, + }, + }, + }, + }, // HeadSliceZero { diff --git a/internal/encoder/compiler.go b/internal/encoder/compiler.go index 76751a9..11b6490 100644 --- a/internal/encoder/compiler.go +++ b/internal/encoder/compiler.go @@ -107,12 +107,24 @@ func compileHead(ctx *compileContext) (*Opcode, error) { return compileBytes(ctx) } } - return compileSlice(ctx) + code, err := compileSlice(ctx) + if err != nil { + return nil, err + } + optimizeStructEnd(code) + linkRecursiveCode(code) + return code, nil case reflect.Map: if isPtr { return compilePtr(ctx.withType(runtime.PtrTo(typ))) } - return compileMap(ctx.withType(typ)) + code, err := compileMap(ctx.withType(typ)) + if err != nil { + return nil, err + } + optimizeStructEnd(code) + linkRecursiveCode(code) + return code, nil case reflect.Struct: code, err := compileStruct(ctx.withType(typ), isPtr) if err != nil { @@ -243,10 +255,11 @@ func linkRecursiveCode(c *Opcode) { lastCode.Idx = beforeLastCode.Idx + uintptrSize lastCode.ElemIdx = lastCode.Idx + uintptrSize + lastCode.Length = lastCode.Idx + 2*uintptrSize - // extend length to alloc slot for elemIdx - totalLength := uintptr(code.TotalLength() + 1) - nextTotalLength := uintptr(c.TotalLength() + 1) + // extend length to alloc slot for elemIdx + length + totalLength := uintptr(code.TotalLength() + 2) + nextTotalLength := uintptr(c.TotalLength() + 2) c.End.Next.Op = OpRecursiveEnd diff --git a/internal/encoder/vm_escaped_indent/vm.go b/internal/encoder/vm_escaped_indent/vm.go index 51e95ff..b221382 100644 --- a/internal/encoder/vm_escaped_indent/vm.go +++ b/internal/encoder/vm_escaped_indent/vm.go @@ -563,6 +563,8 @@ func Run(ctx *encoder.RuntimeContext, b []byte, codeSet *encoder.OpcodeSet, opt offsetNum := ptrOffset / uintptrSize oldOffset := ptrOffset ptrOffset += code.Jmp.CurLen * uintptrSize + oldBaseIndent := ctx.BaseIndent + ctx.BaseIndent += code.Indent - 1 newLen := offsetNum + code.Jmp.CurLen + code.Jmp.NextLen if curlen < newLen { @@ -573,12 +575,14 @@ func Run(ctx *encoder.RuntimeContext, b []byte, codeSet *encoder.OpcodeSet, opt store(ctxptr, c.Idx, ptr) store(ctxptr, c.End.Next.Idx, oldOffset) store(ctxptr, c.End.Next.ElemIdx, uintptr(unsafe.Pointer(code.Next))) + store(ctxptr, c.End.Next.Length, uintptr(oldBaseIndent)) code = c recursiveLevel++ case encoder.OpRecursiveEnd: recursiveLevel-- // restore ctxptr + ctx.BaseIndent = int(load(ctxptr, code.Length)) offset := load(ctxptr, code.Idx) ctx.SeenPtr = ctx.SeenPtr[:len(ctx.SeenPtr)-1] diff --git a/internal/encoder/vm_indent/vm.go b/internal/encoder/vm_indent/vm.go index 4a08160..fdad09e 100644 --- a/internal/encoder/vm_indent/vm.go +++ b/internal/encoder/vm_indent/vm.go @@ -563,6 +563,8 @@ func Run(ctx *encoder.RuntimeContext, b []byte, codeSet *encoder.OpcodeSet, opt offsetNum := ptrOffset / uintptrSize oldOffset := ptrOffset ptrOffset += code.Jmp.CurLen * uintptrSize + oldBaseIndent := ctx.BaseIndent + ctx.BaseIndent += code.Indent - 1 newLen := offsetNum + code.Jmp.CurLen + code.Jmp.NextLen if curlen < newLen { @@ -573,12 +575,14 @@ func Run(ctx *encoder.RuntimeContext, b []byte, codeSet *encoder.OpcodeSet, opt store(ctxptr, c.Idx, ptr) store(ctxptr, c.End.Next.Idx, oldOffset) store(ctxptr, c.End.Next.ElemIdx, uintptr(unsafe.Pointer(code.Next))) + store(ctxptr, c.End.Next.Length, uintptr(oldBaseIndent)) code = c recursiveLevel++ case encoder.OpRecursiveEnd: recursiveLevel-- // restore ctxptr + ctx.BaseIndent = int(load(ctxptr, code.Length)) offset := load(ctxptr, code.Idx) ctx.SeenPtr = ctx.SeenPtr[:len(ctx.SeenPtr)-1]