json: refactor ref-cycle handling

Also set UnsupportedValueError.Value (better stdlib compat).

Author: extemporalgenome

json/codec.go

@@ -10,6 +10,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"time"
 	"unicode"
@@ -32,13 +33,28 @@ type codec struct {
 
 type encoder struct {
 	flags AppendFlags
-	// ptrDepth tracks the depth of pointer cycles, when it reaches the value
+	// refDepth tracks the depth of pointer cycles, when it reaches the value
 	// of startDetectingCyclesAfter, the ptrSeen map is allocated and the
 	// encoder starts tracking pointers it has seen as an attempt to detect
 	// whether it has entered a pointer cycle and needs to error before the
 	// goroutine runs out of stack space.
-	ptrDepth uint32
-	ptrSeen  map[unsafe.Pointer]struct{}
+	//
+	// This relies on encoder being passed as a value,
+	// and encoder methods calling each other in a traditional stack
+	// (not using trampoline techniques),
+	// since refDepth is never decremented.
+	refDepth uint32
+	refSeen  cycleMap
+}
+
+type cycleKey struct {
+	ptr unsafe.Pointer
+}
+
+type cycleMap map[cycleKey]struct{}
+
+var cycleMapPool = sync.Pool{
+	New: func() any { return make(cycleMap) },
 }
 
 type decoder struct {

json/encode.go

@@ -812,22 +812,23 @@ func (e encoder) encodeEmbeddedStructPointer(b []byte, p unsafe.Pointer, t refle
 }
 
 func (e encoder) encodePointer(b []byte, p unsafe.Pointer, t reflect.Type, encode encodeFunc) ([]byte, error) {
-	if p = *(*unsafe.Pointer)(p); p != nil {
-		if e.ptrDepth++; e.ptrDepth >= startDetectingCyclesAfter {
-			if _, seen := e.ptrSeen[p]; seen {
-				// TODO: reconstruct the reflect.Value from p + t so we can set
-				// the erorr's Value field?
-				return b, &UnsupportedValueError{Str: fmt.Sprintf("encountered a cycle via %s", t)}
-			}
-			if e.ptrSeen == nil {
-				e.ptrSeen = make(map[unsafe.Pointer]struct{})
-			}
-			e.ptrSeen[p] = struct{}{}
-			defer delete(e.ptrSeen, p)
+	// p was a pointer to the actual user data pointer:
+	// dereference it to operate on the user data pointer.
+	p = *(*unsafe.Pointer)(p)
+	if p == nil {
+		return e.encodeNull(b, nil)
+	}
+
+	if shouldCheckForRefCycle(&e) {
+		key := cycleKey{ptr: p}
+		if hasRefCycle(&e, key) {
+			return b, refCycleError(t, p)
 		}
-		return encode(e, b, p)
+
+		defer freeRefCycleInfo(&e, key)
 	}
-	return e.encodeNull(b, nil)
+
+	return encode(e, b, p)
 }
 
 func (e encoder) encodeInterface(b []byte, p unsafe.Pointer) ([]byte, error) {
@@ -986,3 +987,64 @@ func appendCompactEscapeHTML(dst []byte, src []byte) []byte {
 
 	return dst
 }
+
+// shouldCheckForRefCycle determines whether checking for reference cycles
+// is reasonable to do at this time.
+//
+// When true, checkRefCycle should be called and any error handled,
+// and then a deferred call to freeRefCycleInfo should be made.
+//
+// This should only be called from encoder methods that are possible points
+// that could directly contribute to a reference cycle.
+func shouldCheckForRefCycle(e *encoder) bool {
+	// Note: do not combine this with checkRefCycle,
+	// because checkRefCycle is too large to be inlined,
+	// and a non-inlined depth check leads to ~5%+ benchmark degradation.
+	e.refDepth++
+	return e.refDepth >= startDetectingCyclesAfter
+}
+
+// refCycleError constructs an [UnsupportedValueError].
+func refCycleError(t reflect.Type, p unsafe.Pointer) error {
+	v := reflect.NewAt(t, p)
+	return &UnsupportedValueError{
+		Value: v,
+		Str:   fmt.Sprintf("encountered a cycle via %s", t),
+	}
+}
+
+// hasRefCycle returns an error if a reference cycle was detected.
+// The data pointer passed in should be equivalent to one of:
+//
+//   - A normal Go pointer, e.g. `unsafe.Pointer(&T)`
+//   - The pointer to a map header, e.g. `*(*unsafe.Pointer)(&map[K]V)`
+//
+// Many [encoder] methods accept a pointer-to-a-pointer,
+// and so those may need to be derenced in order to safely pass them here.
+func hasRefCycle(e *encoder, key cycleKey) bool {
+	_, seen := e.refSeen[key]
+	if seen {
+		return true
+	}
+
+	if e.refSeen == nil {
+		e.refSeen = cycleMapPool.Get().(cycleMap)
+	}
+
+	e.refSeen[key] = struct{}{}
+
+	return false
+}
+
+// freeRefCycle performs the cleanup operation for [checkRefCycle].
+// p must be the same value passed into a prior call to checkRefCycle.
+func freeRefCycleInfo(e *encoder, key cycleKey) {
+	delete(e.refSeen, key)
+	if len(e.refSeen) == 0 {
+		// There are no remaining elements,
+		// so we can release this map for later reuse.
+		m := e.refSeen
+		e.refSeen = nil
+		cycleMapPool.Put(m)
+	}
+}