From 284be3f68b1eaf2ba693e0a8ae03baa80ebc973f Mon Sep 17 00:00:00 2001
From: Luke Shumaker <lukeshu@lukeshu.com>
Date: Sat, 28 Jan 2023 10:43:23 -0700
Subject: reencode: Reuse speculation buffers

---
 encode.go   |  4 ++--
 reencode.go | 34 +++++++++++++++++++++-------------
 2 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/encode.go b/encode.go
index 57f3852..949fd55 100644
--- a/encode.go
+++ b/encode.go
@@ -116,7 +116,7 @@ func encode(w internal.AllWriter, val reflect.Value, escaper BackslashEscaper, q
 			return discardInt(w.WriteString("null"))
 		}
 		// Use a sub-ReEncoder to check that it's a full element.
-		validator := NewReEncoder(w, ReEncoderConfig{BackslashEscape: escaper})
+		validator := &ReEncoder{out: w, ReEncoderConfig: ReEncoderConfig{BackslashEscape: escaper}}
 		if err := obj.EncodeJSON(validator); err != nil {
 			return &EncodeMethodError{
 				Type:       val.Type(),
@@ -152,7 +152,7 @@ func encode(w internal.AllWriter, val reflect.Value, escaper BackslashEscaper, q
 			}
 		}
 		// Use a sub-ReEncoder to check that it's a full element.
-		validator := NewReEncoder(w, ReEncoderConfig{BackslashEscape: escaper})
+		validator := &ReEncoder{out: w, ReEncoderConfig: ReEncoderConfig{BackslashEscape: escaper}}
 		if _, err := validator.Write(dat); err != nil {
 			return &EncodeMethodError{
 				Type:       val.Type(),
diff --git a/reencode.go b/reencode.go
index d588b1b..49d8ddb 100644
--- a/reencode.go
+++ b/reencode.go
@@ -72,6 +72,7 @@ func NewReEncoder(out io.Writer, cfg ReEncoderConfig) *ReEncoder {
 	return &ReEncoder{
 		ReEncoderConfig: cfg,
 		out:             internal.NewAllWriter(out),
+		specu:           new(speculation),
 	}
 }
 
@@ -108,12 +109,21 @@ type ReEncoder struct {
 }
 
 type speculation struct {
+	speculating      bool
 	endWhenStackSize int
 	fmt              ReEncoder
 	compact          bytes.Buffer
 	buf              []inputTuple
 }
 
+func (specu *speculation) Reset() {
+	specu.speculating = false
+	specu.endWhenStackSize = 0
+	specu.fmt = ReEncoder{}
+	specu.compact.Reset()
+	specu.buf = specu.buf[:0]
+}
+
 type inputTuple struct {
 	c         rune
 	t         internal.RuneType
@@ -278,27 +288,25 @@ func (enc *ReEncoder) handleRune(c rune, t internal.RuneType, stackSize int) err
 	}
 
 	// main
-	if enc.specu == nil { // not speculating
+	if !enc.specu.speculating { // not speculating
 		switch t {
 		case internal.RuneTypeObjectBeg, internal.RuneTypeArrayBeg: // start speculating
 			if err, _ := enc.handleRunePre(c, t); err != nil {
 				return err
 			}
-			specu := &speculation{
-				endWhenStackSize: stackSize - 1,
-				fmt: ReEncoder{
-					ReEncoderConfig: enc.ReEncoderConfig,
-				},
+			enc.specu.speculating = true
+			enc.specu.endWhenStackSize = stackSize - 1
+			enc.specu.fmt = ReEncoder{
+				ReEncoderConfig: enc.ReEncoderConfig,
+				out:             &enc.specu.compact,
 			}
-			specu.fmt.Compact = true
-			specu.fmt.out = &specu.compact
-			enc.specu = specu
+			enc.specu.fmt.Compact = true
 			enc.specu.buf = append(enc.specu.buf, inputTuple{
 				c:         c,
 				t:         t,
 				stackSize: stackSize,
 			})
-			if err := specu.fmt.handleRuneMain(c, t); err != nil {
+			if err := enc.specu.fmt.handleRuneMain(c, t); err != nil {
 				return err
 			}
 		default:
@@ -317,8 +325,8 @@ func (enc *ReEncoder) handleRune(c rune, t internal.RuneType, stackSize int) err
 		}
 		switch {
 		case enc.specu.compact.Len() >= enc.CompactIfUnder: // stop speculating; use indent
-			buf := enc.specu.buf
-			enc.specu = nil
+			buf := append([]inputTuple(nil), enc.specu.buf...)
+			enc.specu.Reset()
 			if err := enc.handleRuneMain(buf[0].c, buf[0].t); err != nil {
 				return err
 			}
@@ -331,9 +339,9 @@ func (enc *ReEncoder) handleRune(c rune, t internal.RuneType, stackSize int) err
 			if _, err := enc.specu.compact.WriteTo(enc.out); err != nil {
 				return err
 			}
+			enc.specu.Reset()
 			enc.lastNonSpace = t
 			enc.curIndent--
-			enc.specu = nil
 		}
 	}
 
-- 
cgit v1.2.3