Sync borrowed code from Go 1.20

New tests mean encode.go and compat.go also need some bugfixes.

5 files changed, 224 insertions, 37 deletions

diff --git a/compat/json/borrowed_bench_test.go b/compat/json/borrowed_bench_test.go
index 443a13d..0c13df1 100644
--- a/compat/json/borrowed_bench_test.go
+++ b/compat/json/borrowed_bench_test.go
@@ -19,6 +19,7 @@ import (
 	"io"
 	"os"
 	"reflect"
+	"regexp"
 	"runtime"
 	"strings"
 	"sync"
@@ -100,6 +101,36 @@ func BenchmarkCodeEncoder(b *testing.B) {
 	b.SetBytes(int64(len(codeJSON)))
 }
 
+func BenchmarkCodeEncoderError(b *testing.B) {
+	b.ReportAllocs()
+	if codeJSON == nil {
+		b.StopTimer()
+		codeInit(b) // MODIFIED: use the test logger
+		b.StartTimer()
+	}
+
+	// Trigger an error in Marshal with cyclic data.
+	type Dummy struct {
+		Name string
+		Next *Dummy
+	}
+	dummy := Dummy{Name: "Dummy"}
+	dummy.Next = &dummy
+
+	b.RunParallel(func(pb *testing.PB) {
+		enc := NewEncoder(io.Discard)
+		for pb.Next() {
+			if err := enc.Encode(&codeStruct); err != nil {
+				b.Fatal("Encode:", err)
+			}
+			if _, err := Marshal(dummy); err == nil {
+				b.Fatal("expect an error here")
+			}
+		}
+	})
+	b.SetBytes(int64(len(codeJSON)))
+}
+
 func BenchmarkCodeMarshal(b *testing.B) {
 	b.ReportAllocs()
 	if codeJSON == nil {
@@ -117,6 +148,35 @@ func BenchmarkCodeMarshal(b *testing.B) {
 	b.SetBytes(int64(len(codeJSON)))
 }
 
+func BenchmarkCodeMarshalError(b *testing.B) {
+	b.ReportAllocs()
+	if codeJSON == nil {
+		b.StopTimer()
+		codeInit(b) // MODIFIED: use the test logger
+		b.StartTimer()
+	}
+
+	// Trigger an error in Marshal with cyclic data.
+	type Dummy struct {
+		Name string
+		Next *Dummy
+	}
+	dummy := Dummy{Name: "Dummy"}
+	dummy.Next = &dummy
+
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			if _, err := Marshal(&codeStruct); err != nil {
+				b.Fatal("Marshal:", err)
+			}
+			if _, err := Marshal(dummy); err == nil {
+				b.Fatal("expect an error here")
+			}
+		}
+	})
+	b.SetBytes(int64(len(codeJSON)))
+}
+
 func benchMarshalBytes(n int) func(*testing.B) {
 	sample := []byte("hello world")
 	// Use a struct pointer, to avoid an allocation when passing it as an
@@ -135,6 +195,36 @@ func benchMarshalBytes(n int) func(*testing.B) {
 	}
 }
 
+func benchMarshalBytesError(n int) func(*testing.B) {
+	sample := []byte("hello world")
+	// Use a struct pointer, to avoid an allocation when passing it as an
+	// interface parameter to Marshal.
+	v := &struct {
+		Bytes []byte
+	}{
+		bytes.Repeat(sample, (n/len(sample))+1)[:n],
+	}
+
+	// Trigger an error in Marshal with cyclic data.
+	type Dummy struct {
+		Name string
+		Next *Dummy
+	}
+	dummy := Dummy{Name: "Dummy"}
+	dummy.Next = &dummy
+
+	return func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			if _, err := Marshal(v); err != nil {
+				b.Fatal("Marshal:", err)
+			}
+			if _, err := Marshal(dummy); err == nil {
+				b.Fatal("expect an error here")
+			}
+		}
+	}
+}
+
 func BenchmarkMarshalBytes(b *testing.B) {
 	b.ReportAllocs()
 	// 32 fits within encodeState.scratch.
@@ -146,6 +236,17 @@ func BenchmarkMarshalBytes(b *testing.B) {
 	b.Run("4096", benchMarshalBytes(4096))
 }
 
+func BenchmarkMarshalBytesError(b *testing.B) {
+	b.ReportAllocs()
+	// 32 fits within encodeState.scratch.
+	b.Run("32", benchMarshalBytesError(32))
+	// 256 doesn't fit in encodeState.scratch, but is small enough to
+	// allocate and avoid the slower base64.NewEncoder.
+	b.Run("256", benchMarshalBytesError(256))
+	// 4096 is large enough that we want to avoid allocating for it.
+	b.Run("4096", benchMarshalBytesError(4096))
+}
+
 func BenchmarkCodeDecoder(b *testing.B) {
 	b.ReportAllocs()
 	if codeJSON == nil {
@@ -411,3 +512,33 @@ func BenchmarkEncodeMarshaler(b *testing.B) {
 		}
 	})
 }
+
+func BenchmarkEncoderEncode(b *testing.B) {
+	b.ReportAllocs()
+	type T struct {
+		X, Y string
+	}
+	v := &T{"foo", "bar"}
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			if err := NewEncoder(io.Discard).Encode(v); err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+}
+
+func BenchmarkNumberIsValid(b *testing.B) {
+	s := "-61657.61667E+61673"
+	for i := 0; i < b.N; i++ {
+		isValidNumber(s)
+	}
+}
+
+func BenchmarkNumberIsValidRegexp(b *testing.B) {
+	var jsonNumberRegexp = regexp.MustCompile(`^-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?$`)
+	s := "-61657.61667E+61673"
+	for i := 0; i < b.N; i++ {
+		jsonNumberRegexp.MatchString(s)
+	}
+}
diff --git a/compat/json/borrowed_encode_test.go b/compat/json/borrowed_encode_test.go
index 3d5d675..999c694 100644
--- a/compat/json/borrowed_encode_test.go
+++ b/compat/json/borrowed_encode_test.go
@@ -14,6 +14,7 @@ import (
 	"math"
 	"reflect"
 	"regexp"
+	"runtime/debug"
 	"strconv"
 	"testing"
 	"unicode"
@@ -781,6 +782,42 @@ func TestIssue10281(t *testing.T) {
 	}
 }
 
+//nolint:paralleltest // MODIFIED: added; can't be parallel because it fusses with the global GC.
+func TestMarshalErrorAndReuseEncodeState(t *testing.T) {
+	// Disable the GC temporarily to prevent encodeState's in Pool being cleaned away during the test.
+	percent := debug.SetGCPercent(-1)
+	defer debug.SetGCPercent(percent)
+
+	// Trigger an error in Marshal with cyclic data.
+	type Dummy struct {
+		Name string
+		Next *Dummy
+	}
+	dummy := Dummy{Name: "Dummy"}
+	dummy.Next = &dummy
+	if b, err := Marshal(dummy); err == nil {
+		t.Errorf("Marshal(dummy) = %#q; want error", b)
+	}
+
+	type Data struct {
+		A string
+		I int
+	}
+	data := Data{A: "a", I: 1}
+	b, err := Marshal(data)
+	if err != nil {
+		t.Errorf("Marshal(%v) = %v", data, err)
+	}
+
+	var data2 Data
+	if err := Unmarshal(b, &data2); err != nil {
+		t.Errorf("Unmarshal(%v) = %v", data2, err)
+	}
+	if data2 != data {
+		t.Errorf("expect: %v, but get: %v", data, data2)
+	}
+}
+
 func TestHTMLEscape(t *testing.T) {
 	t.Parallel() // MODIFIED: added
 	var b, want bytes.Buffer
diff --git a/compat/json/borrowed_number_test.go b/compat/json/borrowed_number_test.go
index e7819c6..9709fb4 100644
--- a/compat/json/borrowed_number_test.go
+++ b/compat/json/borrowed_number_test.go
@@ -119,18 +119,3 @@ func TestNumberIsValid(t *testing.T) {
 		}
 	}
 }
-
-func BenchmarkNumberIsValid(b *testing.B) {
-	s := "-61657.61667E+61673"
-	for i := 0; i < b.N; i++ {
-		isValidNumber(s)
-	}
-}
-
-func BenchmarkNumberIsValidRegexp(b *testing.B) {
-	var jsonNumberRegexp = regexp.MustCompile(`^-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?$`)
-	s := "-61657.61667E+61673"
-	for i := 0; i < b.N; i++ {
-		jsonNumberRegexp.MatchString(s)
-	}
-}
diff --git a/compat/json/borrowed_stream_test.go b/compat/json/borrowed_stream_test.go
index 6c3a403..d90898b 100644
--- a/compat/json/borrowed_stream_test.go
+++ b/compat/json/borrowed_stream_test.go
@@ -13,6 +13,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"reflect"
+	"runtime/debug"
 	"strings"
 	"testing"
 )
@@ -61,6 +62,44 @@ func TestEncoder(t *testing.T) {
 	}
 }
 
+//nolint:paralleltest // MODIFIED: added; can't be parallel because it fusses with the global GC.
+func TestEncoderErrorAndReuseEncodeState(t *testing.T) {
+	// Disable the GC temporarily to prevent encodeState's in Pool being cleaned away during the test.
+	percent := debug.SetGCPercent(-1)
+	defer debug.SetGCPercent(percent)
+
+	// Trigger an error in Marshal with cyclic data.
+	type Dummy struct {
+		Name string
+		Next *Dummy
+	}
+	dummy := Dummy{Name: "Dummy"}
+	dummy.Next = &dummy
+
+	var buf bytes.Buffer
+	enc := NewEncoder(&buf)
+	if err := enc.Encode(dummy); err == nil {
+		t.Errorf("Encode(dummy) == nil; want error")
+	}
+
+	type Data struct {
+		A string
+		I int
+	}
+	data := Data{A: "a", I: 1}
+	if err := enc.Encode(data); err != nil {
+		t.Errorf("Marshal(%v) = %v", data, err)
+	}
+
+	var data2 Data
+	if err := Unmarshal(buf.Bytes(), &data2); err != nil {
+		t.Errorf("Unmarshal(%v) = %v", data2, err)
+	}
+	if data2 != data {
+		t.Errorf("expect: %v, but get: %v", data, data2)
+	}
+}
+
 var streamEncodedIndent = `0.1
 "hello"
 null
@@ -324,21 +363,6 @@ func TestBlocking(t *testing.T) {
 	}
 }
 
-func BenchmarkEncoderEncode(b *testing.B) {
-	b.ReportAllocs()
-	type T struct {
-		X, Y string
-	}
-	v := &T{"foo", "bar"}
-	b.RunParallel(func(pb *testing.PB) {
-		for pb.Next() {
-			if err := NewEncoder(io.Discard).Encode(v); err != nil {
-				b.Fatal(err)
-			}
-		}
-	})
-}
-
 //nolint:dupword // False positive, this is commented-out code, not a real comment. // MODIFIED: added nolint declaration
 /* // MODIFIED: we don't have tokens
 type tokenStreamCase struct {
diff --git a/compat/json/compat.go b/compat/json/compat.go
index 688b35c..3678135 100644
--- a/compat/json/compat.go
+++ b/compat/json/compat.go
@@ -72,25 +72,35 @@ func Marshal(v any) ([]byte, error) {
 }
 
 type Encoder struct {
+	out       io.Writer
+	buf       bytes.Buffer
 	encoder   *lowmemjson.Encoder
 	formatter *lowmemjson.ReEncoder
 }
 
 func NewEncoder(w io.Writer) *Encoder {
 	ret := &Encoder{
-		formatter: lowmemjson.NewReEncoder(w, lowmemjson.ReEncoderConfig{
-			AllowMultipleValues: true,
-
-			Compact:               true,
-			ForceTrailingNewlines: true,
-		}),
+		out: w,
 	}
+	ret.formatter = lowmemjson.NewReEncoder(&ret.buf, lowmemjson.ReEncoderConfig{
+		AllowMultipleValues: true,
+
+		Compact:               true,
+		ForceTrailingNewlines: true,
+	})
 	ret.encoder = lowmemjson.NewEncoder(ret.formatter)
 	return ret
 }
 
 func (enc *Encoder) Encode(v any) error {
-	return convertEncodeError(enc.encoder.Encode(v))
+	if err := convertEncodeError(enc.encoder.Encode(v)); err != nil {
+		enc.buf.Reset()
+		return err
+	}
+	if _, err := enc.buf.WriteTo(enc.out); err != nil {
+		return err
+	}
+	return nil
 }
 
 func (enc *Encoder) SetEscapeHTML(on bool) {