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"
@@ -37,8 +38,19 @@ type encoder struct {
 	// 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.
+	//
+	// This relies on encoder being passed as a value,
+	// and encoder methods calling each other in a traditional stack
+	// (not using trampoline techniques),
+	// since ptrDepth is never decremented.
 	ptrDepth uint32
-	ptrSeen  map[unsafe.Pointer]struct{}
+	ptrSeen  cycleMap
+}
+
+type cycleMap map[unsafe.Pointer]struct{}
+
+var cycleMapPool = sync.Pool{
+	New: func() any { return make(cycleMap) },
 }
 
 type decoder struct {

json/encode.go

@@ -794,22 +794,21 @@ 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)
-		}
-		return encode(e, b, 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)
+	}
+
+	err := checkRefCycle(&e, t, p)
+	if err != nil {
+		return b, err
 	}
-	return e.encodeNull(b, nil)
+
+	defer freeRefCycleInfo(&e, p)
+
+	return encode(e, b, p)
 }
 
 func (e encoder) encodeInterface(b []byte, p unsafe.Pointer) ([]byte, error) {
@@ -968,3 +967,55 @@ func appendCompactEscapeHTML(dst []byte, src []byte) []byte {
 
 	return dst
 }
+
+// checkRefCycle 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 checkRefCycle(e *encoder, t reflect.Type, p unsafe.Pointer) error {
+	e.ptrDepth++
+	if e.ptrDepth < startDetectingCyclesAfter {
+		return nil
+	}
+
+	_, seen := e.ptrSeen[p]
+	if seen {
+		v := reflect.NewAt(t, p)
+		return &UnsupportedValueError{
+			Value: v,
+			Str:   fmt.Sprintf("encountered a cycle via %s", t),
+		}
+	}
+
+	if e.ptrSeen == nil {
+		e.ptrSeen = cycleMapPool.Get().(cycleMap)
+	}
+
+	e.ptrSeen[p] = struct{}{}
+
+	return nil
+}
+
+// freeRefCycle performs the cleanup operation for [checkRefCycle].
+// p must be the same value passed into a prior call to checkRefCycle.
+func freeRefCycleInfo(e *encoder, p unsafe.Pointer) {
+	if e.ptrSeen == nil {
+		// The map hasn't yet been allocated (not enough recursion depth),
+		// so there's not any work to do in this function.
+		return
+	}
+
+	delete(e.ptrSeen, p)
+	if len(e.ptrSeen) != 0 {
+		// There are other keys in the map, so we can't release it into the pool.
+		return
+	}
+
+	m := e.ptrSeen
+	e.ptrSeen = nil
+	cycleMapPool.Put(m)
+}