compat/json: Indent: Preserve trailing whitespace

3 files changed, 27 insertions, 10 deletions

diff --git a/compat/json/compat.go b/compat/json/compat.go
index edc6908..d33f278 100644
--- a/compat/json/compat.go
+++ b/compat/json/compat.go
@@ -188,6 +188,15 @@ func Compact(dst *bytes.Buffer, src []byte) error {
 	return err
 }
 
+func isSpace(c byte) bool {
+	switch c {
+	case 0x0020, 0x000A, 0x000D, 0x0009:
+		return true
+	default:
+		return false
+	}
+}
+
 func Indent(dst *bytes.Buffer, src []byte, prefix, indent string) error {
 	start := dst.Len()
 	err := reencode(dst, src, lowmemjson.ReEncoderConfig{
@@ -198,8 +207,18 @@ func Indent(dst *bytes.Buffer, src []byte, prefix, indent string) error {
 	})
 	if err != nil {
 		dst.Truncate(start)
+		return err
 	}
-	return err
+
+	// Preserve trailing whitespace.
+	lastNonWS := len(src) - 1
+	for ; lastNonWS >= 0 && isSpace(src[lastNonWS]); lastNonWS-- {
+	}
+	if _, err := dst.Write(src[lastNonWS+1:]); err != nil {
+		return err
+	}
+
+	return nil
 }
 
 func Valid(data []byte) bool {
diff --git a/compat/json/compat_test.go b/compat/json/compat_test.go
index 0c14a60..c83ca7e 100644
--- a/compat/json/compat_test.go
+++ b/compat/json/compat_test.go
@@ -98,6 +98,13 @@ func TestCompatIndent(t *testing.T) {
 		"object":   {In: `{}`, Out: `{}`},
 		"non-utf8": {In: "\"\x85\xcd\"", Out: "\"\x85\xcd\""},
 		"float":    {In: `1.200e003`, Out: `1.200e003`},
+		"tailws0":  {In: `0`, Out: `0`},
+		"tailws1":  {In: `0 `, Out: `0 `},
+		"tailws2":  {In: `0  `, Out: `0  `},
+		"tailws3":  {In: "0\n", Out: "0\n"},
+		"headws1":  {In: ` 0`, Out: `0`},
+		"objws1":   {In: `{"a"  :  1}`, Out: "{\n>.\"a\": 1\n>}"},
+		"objws2":   {In: "{\"a\"\n:\n1}", Out: "{\n>.\"a\": 1\n>}"},
 	}
 	for tcName, tc := range testcases {
 		tc := tc
diff --git a/compat/json/testcompat_test.go b/compat/json/testcompat_test.go
index e89b4b4..73153d9 100644
--- a/compat/json/testcompat_test.go
+++ b/compat/json/testcompat_test.go
@@ -46,15 +46,6 @@ const (
 	startDetectingCyclesAfter = 1000
 )
 
-func isSpace(c byte) bool {
-	switch c {
-	case 0x0020, 0x000A, 0x000D, 0x0009:
-		return true
-	default:
-		return false
-	}
-}
-
 type encodeState struct {
 	bytes.Buffer
 }