reencode: Fix .stackSize

2 files changed, 18 insertions, 2 deletions

diff --git a/reencode.go b/reencode.go
index a33cc8f..fd848f8 100644
--- a/reencode.go
+++ b/reencode.go
@@ -364,8 +364,8 @@ func (enc *ReEncoder) popWriteBarrier() {
 
 func (enc *ReEncoder) stackSize() int {
 	sz := enc.par.StackSize()
-	for _, barrier := range enc.barriers {
-		sz += barrier.stackSize
+	if len(enc.barriers) > 0 {
+		sz += enc.barriers[len(enc.barriers)-1].stackSize
 	}
 	return sz
 }
diff --git a/reencode_test.go b/reencode_test.go
index 82a1861..bc6d246 100644
--- a/reencode_test.go
+++ b/reencode_test.go
@@ -9,6 +9,8 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
+	"git.lukeshu.com/go/lowmemjson/internal/fastio"
 )
 
 func TestReEncode(t *testing.T) {
@@ -222,3 +224,17 @@ func TestReEncodeWriteSize(t *testing.T) {
 		assert.Equal(t, input, out.String())
 	})
 }
+
+func TestReEncoderStackSize(t *testing.T) {
+	t.Parallel()
+
+	enc := NewReEncoder(fastio.Discard, ReEncoderConfig{})
+	assert.Equal(t, 0, enc.stackSize())
+
+	for i := 0; i < 5; i++ {
+		assert.NoError(t, enc.WriteByte('['))
+		assert.Equal(t, i+1, enc.stackSize())
+		enc.pushWriteBarrier()
+		assert.Equal(t, i+2, enc.stackSize())
+	}
+}