Merge commit 'b3f4186f2b8e992f56f898784b1cd28bfd7550ca' into lukeshu/fix

1 files changed, 11 insertions, 15 deletions

diff --git a/encode.go b/encode.go
index 5fb4fbf..ca4e060 100644
--- a/encode.go
+++ b/encode.go
@@ -9,17 +9,13 @@ import (
 	"encoding"
 	"encoding/base64"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"io"
-	iofs "io/fs"
 	"reflect"
 	"sort"
 	"strconv"
 	"strings"
 	"unsafe"
-
-	"git.lukeshu.com/go/lowmemjson/internal"
 )
 
 // Encodable is the interface implemented by types that can encode
@@ -98,7 +94,7 @@ var (
 
 const startDetectingCyclesAfter = 1000
 
-func encode(w internal.AllWriter, val reflect.Value, escaper BackslashEscaper, quote bool, cycleDepth uint, cycleSeen map[any]struct{}) error {
+func encode(w *ReEncoder, val reflect.Value, escaper BackslashEscaper, quote bool, cycleDepth uint, cycleSeen map[any]struct{}) error {
 	if !val.IsValid() {
 		return discardInt(w.WriteString("null"))
 	}
@@ -115,22 +111,22 @@ func encode(w internal.AllWriter, val reflect.Value, escaper BackslashEscaper, q
 		if !ok {
 			return discardInt(w.WriteString("null"))
 		}
-		// Use a sub-ReEncoder to check that it's a full element.
-		validator := &ReEncoder{out: w, ReEncoderConfig: ReEncoderConfig{BackslashEscape: EscapePreserve}}
-		if err := obj.EncodeJSON(validator); err != nil {
+		w.pushWriteBarrier()
+		if err := obj.EncodeJSON(w); err != nil {
 			return &EncodeMethodError{
 				Type:       val.Type(),
 				SourceFunc: "EncodeJSON",
 				Err:        err,
 			}
 		}
-		if err := validator.Close(); err != nil && !errors.Is(err, iofs.ErrClosed) {
+		if err := w.Close(); err != nil {
 			return &EncodeMethodError{
 				Type:       val.Type(),
 				SourceFunc: "EncodeJSON",
 				Err:        err,
 			}
 		}
+		w.popWriteBarrier()
 
 	case val.Kind() != reflect.Pointer && val.CanAddr() && reflect.PointerTo(val.Type()).Implements(jsonMarshalerType):
 		val = val.Addr()
@@ -151,22 +147,22 @@ func encode(w internal.AllWriter, val reflect.Value, escaper BackslashEscaper, q
 				Err:        err,
 			}
 		}
-		// Use a sub-ReEncoder to check that it's a full element.
-		validator := &ReEncoder{out: w, ReEncoderConfig: ReEncoderConfig{BackslashEscape: EscapePreserve}}
-		if _, err := validator.Write(dat); err != nil {
+		w.pushWriteBarrier()
+		if _, err := w.Write(dat); err != nil {
 			return &EncodeMethodError{
 				Type:       val.Type(),
 				SourceFunc: "MarshalJSON",
 				Err:        err,
 			}
 		}
-		if err := validator.Close(); err != nil {
+		if err := w.Close(); err != nil {
 			return &EncodeMethodError{
 				Type:       val.Type(),
 				SourceFunc: "MarshalJSON",
 				Err:        err,
 			}
 		}
+		w.popWriteBarrier()
 
 	case val.Kind() != reflect.Pointer && val.CanAddr() && reflect.PointerTo(val.Type()).Implements(textMarshalerType):
 		val = val.Addr()
@@ -361,7 +357,7 @@ func encode(w internal.AllWriter, val reflect.Value, escaper BackslashEscaper, q
 			for i := 0; iter.Next(); i++ {
 				// TODO: Avoid buffering the map key
 				var k strings.Builder
-				if err := encode(&k, iter.Key(), escaper, false, cycleDepth, cycleSeen); err != nil {
+				if err := encode(NewReEncoder(&k, ReEncoderConfig{BackslashEscape: escaper}), iter.Key(), escaper, false, cycleDepth, cycleSeen); err != nil {
 					return err
 				}
 				kStr := k.String()
@@ -496,7 +492,7 @@ func encode(w internal.AllWriter, val reflect.Value, escaper BackslashEscaper, q
 	return nil
 }
 
-func encodeArray(w internal.AllWriter, val reflect.Value, escaper BackslashEscaper, cycleDepth uint, cycleSeen map[any]struct{}) error {
+func encodeArray(w *ReEncoder, val reflect.Value, escaper BackslashEscaper, cycleDepth uint, cycleSeen map[any]struct{}) error {
 	if err := w.WriteByte('['); err != nil {
 		return err
 	}