encoding/json: fix Indent with a non-space prefix or indent

src/encoding/json/scanner_test.go

@@ -110,6 +110,28 @@ func TestCompactAndIndent(t *testing.T) {
 	}
 }
 
+func TestIndentPrefix(t *testing.T) {
+	tests := []struct {
+		CaseName
+		src, prefix, indent string
+		want                string
+	}{
+		{Name(""), "[1,[2]]", ">", "--", "[\n>--1,\n>--[\n>----2\n>--]\n>]"},
+		{Name(""), "[1,2]\n  ", ">", "--", "[\n>--1,\n>--2\n>]\n  "},
+		{Name(""), "1\n  ", ">", "", "1\n  "}, // See golang.org/issue/81547
+	}
+	for _, tt := range tests {
+		t.Run(tt.Name, func(t *testing.T) {
+			var buf bytes.Buffer
+			if err := Indent(&buf, []byte(tt.src), tt.prefix, tt.indent); err != nil {
+				t.Errorf("%s: Indent error: %v", tt.Where, err)
+			} else if got := buf.String(); got != tt.want {
+				t.Errorf("%s: Indent:\n\tgot:  %q\n\twant: %q", tt.Where, got, tt.want)
+			}
+		})
+	}
+}
+
 func TestCompactSeparators(t *testing.T) {
 	// U+2028 and U+2029 should be escaped inside strings.
 	// They should not appear outside strings.

src/encoding/json/v2_indent.go

@@ -91,24 +91,12 @@ func Indent(dst *bytes.Buffer, src []byte, prefix, indent string) error {
 func appendIndent(dst, src []byte, prefix, indent string) ([]byte, error) {
 	// In v2, only spaces and tabs are allowed, while v1 allowed any character.
 	dstLen := len(dst)
-	if len(strings.Trim(prefix, " \t"))+len(strings.Trim(indent, " \t")) > 0 {
+	invalidPrefix, invalidIndent := prefix, indent
+	replaceIndent := len(strings.Trim(prefix, " \t"))+len(strings.Trim(indent, " \t")) > 0
+	if replaceIndent {
 		// Use placeholder spaces of correct length, and replace afterwards.
-		invalidPrefix, invalidIndent := prefix, indent
 		prefix = strings.Repeat(" ", len(prefix))
 		indent = strings.Repeat(" ", len(indent))
-		defer func() {
-			b := dst[dstLen:]
-			for i := bytes.IndexByte(b, '\n'); i >= 0; i = bytes.IndexByte(b, '\n') {
-				b = b[i+len("\n"):]
-				n := len(b) - len(bytes.TrimLeft(b, " ")) // len(prefix)+n*len(indent)
-				spaces := b[:n]
-				spaces = spaces[copy(spaces, invalidPrefix):]
-				for len(spaces) > 0 {
-					spaces = spaces[copy(spaces, invalidIndent):]
-				}
-				b = b[n:]
-			}
-		}()
 	}
 
 	dst, err := jsontext.AppendFormat(dst, src,
@@ -123,6 +111,21 @@ func appendIndent(dst, src []byte, prefix, indent string) ([]byte, error) {
 		return dst[:dstLen], transformSyntacticError(err)
 	}
 
+	// Replace the placeholder spaces with the original prefix and indent.
+	if replaceIndent {
+		b := dst[dstLen:]
+		for i := bytes.IndexByte(b, '\n'); i >= 0; i = bytes.IndexByte(b, '\n') {
+			b = b[i+len("\n"):]
+			n := len(b) - len(bytes.TrimLeft(b, " ")) // len(prefix)+n*len(indent)
+			spaces := b[:n]
+			spaces = spaces[copy(spaces, invalidPrefix):]
+			for len(spaces) > 0 {
+				spaces = spaces[copy(spaces, invalidIndent):]
+			}
+			b = b[n:]
+		}
+	}
+
 	// In v2, trailing whitespace is discarded, while v1 preserved it.
 	if n := len(src) - len(bytes.TrimRight(src, " \n\r\t")); n > 0 {
 		dst = append(dst, src[len(src)-n:]...)

src/encoding/json/v2_scanner_test.go

@@ -110,6 +110,28 @@ func TestCompactAndIndent(t *testing.T) {
 	}
 }
 
+func TestIndentPrefix(t *testing.T) {
+	tests := []struct {
+		CaseName
+		src, prefix, indent string
+		want                string
+	}{
+		{Name(""), "[1,[2]]", ">", "--", "[\n>--1,\n>--[\n>----2\n>--]\n>]"},
+		{Name(""), "[1,2]\n  ", ">", "--", "[\n>--1,\n>--2\n>]\n  "},
+		{Name(""), "1\n  ", ">", "", "1\n  "}, // See golang.org/issue/81547
+	}
+	for _, tt := range tests {
+		t.Run(tt.Name, func(t *testing.T) {
+			var buf bytes.Buffer
+			if err := Indent(&buf, []byte(tt.src), tt.prefix, tt.indent); err != nil {
+				t.Errorf("%s: Indent error: %v", tt.Where, err)
+			} else if got := buf.String(); got != tt.want {
+				t.Errorf("%s: Indent:\n\tgot:  %q\n\twant: %q", tt.Where, got, tt.want)
+			}
+		})
+	}
+}
+
 func TestCompactSeparators(t *testing.T) {
 	// U+2028 and U+2029 should be escaped inside strings.
 	// They should not appear outside strings.