Merge branch 'lukeshu/go1.20'

7 files changed, 246 insertions, 42 deletions

diff --git a/ReleaseNotes.md b/ReleaseNotes.md
index f71967c..493252a 100644
--- a/ReleaseNotes.md
+++ b/ReleaseNotes.md
@@ -9,6 +9,15 @@
      poisons future Decodes.  This is something that `encoding/json`
      supports.
 
+   - Encoder: Fixes a bug where if an encode error is encountered, all
+     future Encode calls will fail.  Reusing an Encoder is something
+     that `encoding/json` supports.
+
+   - compat/json.Encoder: Now buffers the output, to avoid partial
+     writes if an encode error is encountered.  This matches the
+     behavior of `encoding/json`.  For memory consumption reasons, the
+     native lowmemjson Encoder still does not buffer.
+
 # v0.3.4 (2023-02-05)
 
   Theme: Fix compilation with Go 1.20
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) {
diff --git a/encode.go b/encode.go
index fa558b9..2d16891 100644
--- a/encode.go
+++ b/encode.go
@@ -40,8 +40,8 @@ type Encodable interface {
 // lowmemjson/compat/json.Encoder offers those .SetEscapeHTML and
 // .SetIndent methods.
 type Encoder struct {
-	w                *ReEncoder
-	closeAfterEncode bool
+	w      *ReEncoder
+	isRoot bool
 }
 
 // NewEncoder returns a new Encoder that writes to w.
@@ -60,8 +60,8 @@ func NewEncoder(w io.Writer) *Encoder {
 		})
 	}
 	return &Encoder{
-		w:                re,
-		closeAfterEncode: re.par.StackIsEmpty(),
+		w:      re,
+		isRoot: re.par.StackIsEmpty(),
 	}
 }
 
@@ -73,12 +73,20 @@ func NewEncoder(w io.Writer) *Encoder {
 // that, with the exception that in addition to the json.Marshaler
 // interface it also checks for the Encodable interface.
 //
+// Unlike encoding/json.Encoder.Encode, lowmemjson.Encoder.Encode does
+// not buffer its output; if a encode-error is encountered, lowmemjson
+// may write partial output, whereas encodin/json would not have
+// written anything.
+//
 // [documentation for encoding/json.Marshal]: https://pkg.go.dev/encoding/json@go1.18#Marshal
 func (enc *Encoder) Encode(obj any) (err error) {
+	if enc.isRoot {
+		enc.w.par.Reset()
+	}
 	if err := encode(enc.w, reflect.ValueOf(obj), enc.w.BackslashEscape, false, 0, map[any]struct{}{}); err != nil {
 		return err
 	}
-	if enc.closeAfterEncode {
+	if enc.isRoot {
 		return enc.w.Close()
 	}
 	return nil