reencode: Rethink CompactIfUnder to have linear memory

2 files changed, 49 insertions, 34 deletions

diff --git a/internal/parse.go b/internal/parse.go
index b11aae6..9db57fb 100644
--- a/internal/parse.go
+++ b/internal/parse.go
@@ -349,6 +349,10 @@ func (par *Parser) StackIsEmpty() bool {
 	return len(par.stack) == 0 || (len(par.stack) == 1 && par.stack[0] == runeTypeAny)
 }
 
+func (par *Parser) StackSize() int {
+	return len(par.stack)
+}
+
 // Reset all Parser state.
 func (par *Parser) Reset() {
 	*par = Parser{
diff --git a/reencode.go b/reencode.go
index 393e8c6..b3f4d20 100644
--- a/reencode.go
+++ b/reencode.go
@@ -35,8 +35,8 @@ type ReEncoderConfig struct {
 	//
 	// Has no affect if Compact is true or Indent is empty.
 	//
-	// This has O((CompactIfUnder+1)^2) memory overhead, so set
-	// with caution.
+	// his has O(2^min(CompactIfUnder, depth)) time overhead, so
+	// set with caution.
 	CompactIfUnder int
 
 	// String to use to indent; ignored if Compact is true.
@@ -82,7 +82,7 @@ func NewReEncoder(out io.Writer, cfg ReEncoderConfig) *ReEncoder {
 // This is useful for prettifying, minifying, sanitizing, and/or
 // validating JSON.
 //
-// The memory use of a ReEncoder is O( (CompactIfUnder+1)^2 + depth).
+// The memory use of a ReEncoder is O(CompactIfUnder+depth).
 type ReEncoder struct {
 	ReEncoderConfig
 	out internal.AllWriter
@@ -111,10 +111,16 @@ type ReEncoder struct {
 }
 
 type speculation struct {
-	compactFmt ReEncoder
-	compactBuf bytes.Buffer
-	indentFmt  ReEncoder
-	indentBuf  bytes.Buffer
+	endWhenStackSize int
+	fmt              ReEncoder
+	compact          bytes.Buffer
+	buf              []inputTuple
+}
+
+type inputTuple struct {
+	c         rune
+	t         internal.RuneType
+	stackSize int
 }
 
 // public API //////////////////////////////////////////////////////////////////
@@ -207,7 +213,7 @@ func (enc *ReEncoder) Close() error {
 		}
 		return enc.err
 	}
-	if err := enc.handleRune(0, internal.RuneTypeError); err != nil {
+	if err := enc.handleRune(0, internal.RuneTypeError, enc.par.StackSize()); err != nil {
 		enc.err = &ReEncodeSyntaxError{
 			Err:    err,
 			Offset: enc.inputPos,
@@ -249,7 +255,7 @@ rehandle:
 		}
 		return enc.written, enc.err
 	}
-	enc.err = enc.handleRune(c, t)
+	enc.err = enc.handleRune(c, t, enc.par.StackSize())
 	if enc.err == nil && t == internal.RuneTypeEOF {
 		if enc.AllowMultipleValues {
 			enc.par.Reset()
@@ -269,7 +275,7 @@ rehandle:
 
 // internal ////////////////////////////////////////////////////////////////////
 
-func (enc *ReEncoder) handleRune(c rune, t internal.RuneType) error {
+func (enc *ReEncoder) handleRune(c rune, t internal.RuneType, stackSize int) error {
 	if enc.CompactIfUnder == 0 || enc.Compact || enc.Indent == "" {
 		return enc.handleRuneNoSpeculation(c, t)
 	}
@@ -282,17 +288,20 @@ func (enc *ReEncoder) handleRune(c rune, t internal.RuneType) error {
 				return err
 			}
 			specu := &speculation{
-				compactFmt: *enc,
-				indentFmt:  *enc,
+				endWhenStackSize: stackSize - 1,
+				fmt: ReEncoder{
+					ReEncoderConfig: enc.ReEncoderConfig,
+				},
 			}
-			specu.compactFmt.Compact = true
-			specu.compactFmt.out = &specu.compactBuf
-			specu.indentFmt.out = &specu.indentBuf
+			specu.fmt.Compact = true
+			specu.fmt.out = &specu.compact
 			enc.handleRuneState.specu = specu
-			if err := specu.compactFmt.handleRuneMain(c, t); err != nil {
-				return err
-			}
-			if err := specu.indentFmt.handleRuneMain(c, t); err != nil {
+			enc.handleRuneState.specu.buf = append(enc.handleRuneState.specu.buf, inputTuple{
+				c:         c,
+				t:         t,
+				stackSize: stackSize,
+			})
+			if err := specu.fmt.handleRuneMain(c, t); err != nil {
 				return err
 			}
 		default:
@@ -301,26 +310,28 @@ func (enc *ReEncoder) handleRune(c rune, t internal.RuneType) error {
 			}
 		}
 	} else { // speculating
-
-		// canCompress is whether we're 1-up from the leaf;
-		// set this *before* the calls to .handleRune.
-		canCompress := enc.handleRuneState.specu.indentFmt.handleRuneState.specu == nil
-
-		if err := enc.handleRuneState.specu.compactFmt.handleRune(c, t); err != nil {
+		enc.handleRuneState.specu.buf = append(enc.handleRuneState.specu.buf, inputTuple{
+			c:         c,
+			t:         t,
+			stackSize: stackSize,
+		})
+		if err := enc.handleRuneState.specu.fmt.handleRune(c, t, stackSize); err != nil {
 			return err
 		}
-		if err := enc.handleRuneState.specu.indentFmt.handleRune(c, t); err != nil {
-			return err
-		}
-
 		switch {
-		case enc.handleRuneState.specu.compactBuf.Len() >= enc.CompactIfUnder: // stop speculating; use indent
-			if _, err := enc.handleRuneState.specu.indentBuf.WriteTo(enc.out); err != nil {
+		case enc.handleRuneState.specu.compact.Len() >= enc.CompactIfUnder: // stop speculating; use indent
+			buf := enc.handleRuneState.specu.buf
+			enc.handleRuneState.specu = nil
+			if err := enc.handleRuneMain(buf[0].c, buf[0].t); err != nil {
 				return err
 			}
-			enc.handleRuneState = enc.handleRuneState.specu.indentFmt.handleRuneState
-		case canCompress && (t == internal.RuneTypeObjectEnd || t == internal.RuneTypeArrayEnd): // stop speculating; use compact
-			if _, err := enc.handleRuneState.specu.compactBuf.WriteTo(enc.out); err != nil {
+			for _, tuple := range buf[1:] {
+				if err := enc.handleRune(tuple.c, tuple.t, tuple.stackSize); err != nil {
+					return err
+				}
+			}
+		case stackSize == enc.handleRuneState.specu.endWhenStackSize: // stop speculating; use compact
+			if _, err := enc.handleRuneState.specu.compact.WriteTo(enc.out); err != nil {
 				return err
 			}
 			enc.handleRuneState.lastNonSpace = t