Merge branch 'lukeshu/fix'

5 files changed, 137 insertions, 35 deletions

diff --git a/ReleaseNotes.md b/ReleaseNotes.md
index 76a1a5b..174ecaf 100644
--- a/ReleaseNotes.md
+++ b/ReleaseNotes.md
@@ -1,3 +1,17 @@
+# v0.3.3 (TBD)
+
+  Theme: Bugfix
+
+  User-facing changes:
+
+   - ReEncoder: Fixes a regression in v0.3.1 where it erroneously
+     enters compact mode when CompactIfUnder is set and write barriers
+     are in use.
+
+   - ReEncoder: Fixes a regression in v0.3.1 where it sometimes emits
+     extra (but syntactically valid) newlines when write barriers are
+     in use.
+
 # v0.3.2 (2023-02-03)
 
   Theme: Bugfix
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:
diff --git a/roundtrip_test.go b/roundtrip_test.go
index 71ca6d0..01d6f13 100644
--- a/roundtrip_test.go
+++ b/roundtrip_test.go
@@ -6,9 +6,9 @@ package lowmemjson_test
 
 import (
 	"bytes"
+	"encoding/json"
 	"os"
 	"path/filepath"
-	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -16,7 +16,7 @@ import (
 	"git.lukeshu.com/go/lowmemjson"
 )
 
-type ScanDevicesResult map[string]ScanOneDeviceResult
+type ScanDevicesResult map[uint64]ScanOneDeviceResult
 
 type ScanOneDeviceResult struct {
 	FoundExtentCSums []SysExtentCSum
@@ -27,27 +27,82 @@ type SysExtentCSum struct {
 	Sums       SumRun
 }
 
+type Optional[T any] struct {
+	OK  bool
+	Val T
+}
+
+var (
+	_ json.Marshaler   = Optional[bool]{}
+	_ json.Unmarshaler = (*Optional[bool])(nil)
+)
+
+func (o Optional[T]) MarshalJSON() ([]byte, error) {
+	if o.OK {
+		return json.Marshal(o.Val)
+	} else {
+		return []byte("null"), nil
+	}
+}
+
+func (o *Optional[T]) UnmarshalJSON(dat []byte) error {
+	if string(dat) == "null" {
+		*o = Optional[T]{}
+		return nil
+	}
+	o.OK = true
+	return json.Unmarshal(dat, &o.Val)
+}
+
+type QualifiedPhysicalAddr struct {
+	Dev  uint64
+	Addr int64
+}
+
+type Mapping struct {
+	LAddr      int64
+	PAddr      QualifiedPhysicalAddr
+	Size       int64
+	SizeLocked bool             `json:",omitempty"`
+	Flags      Optional[uint64] `json:",omitempty"`
+}
+
 func TestRoundTrip(t *testing.T) {
 	t.Parallel()
-	dents, err := os.ReadDir(filepath.Join("testdata", "roundtrip"))
-	require.NoError(t, err)
-	for _, dent := range dents {
-		filename := dent.Name()
-		if !strings.HasSuffix(filename, ".json") {
-			continue
-		}
-		t.Run(strings.TrimSuffix(filename, ".json"), func(t *testing.T) {
+
+	type testcase struct {
+		ObjPtr any
+		Cfg    lowmemjson.ReEncoderConfig
+	}
+	testcases := map[string]testcase{
+		"scandevices": {
+			ObjPtr: new(ScanDevicesResult),
+			Cfg: lowmemjson.ReEncoderConfig{
+				Indent:                "\t",
+				ForceTrailingNewlines: true,
+				CompactIfUnder:        16,
+			},
+		},
+		"mappings": {
+			ObjPtr: new([]Mapping),
+			Cfg: lowmemjson.ReEncoderConfig{
+				Indent:                "\t",
+				ForceTrailingNewlines: true,
+				CompactIfUnder:        120,
+			},
+		},
+	}
+
+	for tcName, tc := range testcases {
+		tcName := tcName
+		tc := tc
+		t.Run(tcName, func(t *testing.T) {
 			t.Parallel()
-			inBytes, err := os.ReadFile(filepath.Join("testdata", "roundtrip", filename))
+			inBytes, err := os.ReadFile(filepath.Join("testdata", "roundtrip", tcName+".json")) // #nosec G304
 			require.NoError(t, err)
-			var obj ScanDevicesResult
-			require.NoError(t, lowmemjson.NewDecoder(bytes.NewReader(inBytes)).DecodeThenEOF(&obj))
+			require.NoError(t, lowmemjson.NewDecoder(bytes.NewReader(inBytes)).DecodeThenEOF(tc.ObjPtr))
 			var outBytes bytes.Buffer
-			require.NoError(t, lowmemjson.NewEncoder(lowmemjson.NewReEncoder(&outBytes, lowmemjson.ReEncoderConfig{
-				Indent:                "\t",
-				ForceTrailingNewlines: true,
-				CompactIfUnder:        16, //nolint:gomnd // This is what looks nice.
-			})).Encode(obj))
+			require.NoError(t, lowmemjson.NewEncoder(lowmemjson.NewReEncoder(&outBytes, tc.Cfg)).Encode(tc.ObjPtr))
 			require.Equal(t, string(inBytes), outBytes.String())
 		})
 	}
diff --git a/testdata/roundtrip/mappings.json b/testdata/roundtrip/mappings.json
new file mode 100644
index 0000000..69e596f
--- /dev/null
+++ b/testdata/roundtrip/mappings.json
@@ -0,0 +1,15 @@
+[
+	{"LAddr":1048576,"PAddr":{"Dev":1,"Addr":1048576},"Size":4194304,"Flags":2},
+	{"LAddr":22020096,"PAddr":{"Dev":1,"Addr":22020096},"Size":1073741824,"Flags":4},
+	{"LAddr":1095761920,"PAddr":{"Dev":1,"Addr":1095761920},"Size":1073741824,"Flags":1},
+	{"LAddr":2169503744,"PAddr":{"Dev":1,"Addr":2169503744},"Size":1073741824,"Flags":1},
+	{"LAddr":3243245568,"PAddr":{"Dev":1,"Addr":3243245568},"Size":1073741824,"Flags":1},
+	{"LAddr":4316987392,"PAddr":{"Dev":1,"Addr":4316987392},"Size":1073741824,"Flags":1},
+	{"LAddr":5390729216,"PAddr":{"Dev":1,"Addr":5390729216},"Size":1073741824,"Flags":1},
+	{"LAddr":6464471040,"PAddr":{"Dev":1,"Addr":6464471040},"Size":1073741824,"Flags":1},
+	{"LAddr":7538212864,"PAddr":{"Dev":1,"Addr":7538212864},"Size":1073741824,"Flags":1},
+	{"LAddr":8611954688,"PAddr":{"Dev":1,"Addr":8611954688},"Size":1073741824,"Flags":1},
+	{"LAddr":9685696512,"PAddr":{"Dev":1,"Addr":9685696512},"Size":1073741824,"Flags":1},
+	{"LAddr":413900537856,"PAddr":{"Dev":1,"Addr":227655286784},"Size":1073741824,"Flags":1},
+	{"LAddr":414974279680,"PAddr":{"Dev":1,"Addr":243761414144},"Size":1073741824,"Flags":1}
+]
diff --git a/testdata/roundtrip/scandevices.json b/testdata/roundtrip/scandevices.json
index c62f7df..f56a312 100644
--- a/testdata/roundtrip/scandevices.json
+++ b/testdata/roundtrip/scandevices.json
@@ -1,5 +1,5 @@
 {
-	"dev-1234": {
+	"1234": {
 		"FoundExtentCSums": [
 			{
 				"Generation": 6596005,