Merge branch 'lukeshu/fix'

1 files changed, 34 insertions, 16 deletions

diff --git a/reencode.go b/reencode.go
index b135fb8..3e9cf37 100644
--- a/reencode.go
+++ b/reencode.go
@@ -108,8 +108,13 @@ type ReEncoder struct {
 	expZero            bool
 	specu              *speculation
 
-	// state: .pushBarrier and .popBarrier
-	stackInputPos []int64
+	// state: .pushWriteBarrier and .popWriteBarrier
+	barriers []barrier
+}
+
+type barrier struct {
+	inputPos  int64
+	stackSize int
 }
 
 type speculation struct {
@@ -224,15 +229,17 @@ func (enc *ReEncoder) Close() error {
 		}
 		return enc.err
 	}
-	if err := enc.handleRune(0, internal.RuneTypeError, enc.par.StackSize()); err != nil {
-		enc.err = &ReEncodeSyntaxError{
-			Err:    err,
-			Offset: enc.inputPos,
+	if len(enc.barriers) == 0 {
+		if err := enc.handleRune(0, internal.RuneTypeError, enc.stackSize()); err != nil {
+			enc.err = &ReEncodeSyntaxError{
+				Err:    err,
+				Offset: enc.inputPos,
+			}
+			return enc.err
+		}
+		if enc.AllowMultipleValues {
+			enc.par.Reset()
 		}
-		return enc.err
-	}
-	if enc.AllowMultipleValues && len(enc.stackInputPos) == 0 {
-		enc.par.Reset()
 	}
 	return nil
 }
@@ -266,9 +273,9 @@ rehandle:
 		}
 		return enc.written, enc.err
 	}
-	enc.err = enc.handleRune(c, t, enc.par.StackSize())
+	enc.err = enc.handleRune(c, t, enc.stackSize())
 	if enc.err == nil && t == internal.RuneTypeEOF {
-		if enc.AllowMultipleValues && len(enc.stackInputPos) == 0 {
+		if enc.AllowMultipleValues && len(enc.barriers) == 0 {
 			enc.par.Reset()
 			goto rehandle
 		} else {
@@ -287,20 +294,31 @@ rehandle:
 // semi-public API /////////////////////////////////////////////////////////////
 
 func (enc *ReEncoder) pushWriteBarrier() {
+	enc.barriers = append(enc.barriers, barrier{
+		inputPos:  enc.inputPos,
+		stackSize: enc.stackSize(),
+	})
 	enc.par.PushWriteBarrier()
-	enc.stackInputPos = append(enc.stackInputPos, enc.inputPos)
 	enc.inputPos = 0
 }
 
 func (enc *ReEncoder) popWriteBarrier() {
 	enc.par.PopBarrier()
-	enc.inputPos += enc.stackInputPos[len(enc.stackInputPos)-1]
-	enc.stackInputPos = enc.stackInputPos[:len(enc.stackInputPos)-1]
+	enc.inputPos += enc.barriers[len(enc.barriers)-1].inputPos
+	enc.barriers = enc.barriers[:len(enc.barriers)-1]
 	enc.lastNonSpace = enc.lastNonSpaceNonEOF
 }
 
 // internal ////////////////////////////////////////////////////////////////////
 
+func (enc *ReEncoder) stackSize() int {
+	sz := enc.par.StackSize()
+	for _, barrier := range enc.barriers {
+		sz += barrier.stackSize
+	}
+	return sz
+}
+
 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)
@@ -522,7 +540,7 @@ func (enc *ReEncoder) handleRuneMain(c rune, t internal.RuneType) error {
 	case internal.RuneTypeEOF: // EOF implied by the start of the next top-level value
 		enc.wasNumber = enc.lastNonSpace.IsNumber()
 		switch {
-		case enc.ForceTrailingNewlines && len(enc.stackInputPos) == 0:
+		case enc.ForceTrailingNewlines && len(enc.barriers) == 0:
 			t = internal.RuneTypeError // enc.lastNonSpace : an NL isn't needed (we already printed one)
 			err = enc.emitByte('\n')
 		default: