Update StringN() behavior.

Author: qingyang-hu

x/bsonx/bsoncore/array.go

@@ -82,28 +82,24 @@ func (a Array) DebugString() string {
 // String outputs an ExtendedJSON version of Array. If the Array is not valid, this method
 // returns an empty string.
 func (a Array) String() string {
-	str, _ := a.stringN(0)
+	str, _ := a.StringN(-1)
 	return str
 }
 
-// StringN stringifies an array upto N bytes
+// StringN stringifies an array. If N is non-negative, it will truncate the string to N bytes.
+// Otherwise, it will return the full string representation. The second return value indicates
+// whether the string was truncated or not.
 func (a Array) StringN(n int) (string, bool) {
-	if n <= 0 {
-		if l, _, ok := ReadLength(a); !ok || l < 5 {
-			return "", false
-		}
-		return "", true
-	}
-	return a.stringN(n)
-}
-
-// stringN stringify an array. If N is larger than 0, it will truncate the string to N bytes.
-func (a Array) stringN(n int) (string, bool) {
 	length, rem, ok := ReadLength(a)
 	if !ok || length < 5 {
 		return "", false
 	}
-	length -= (4 /* length bytes */ + 1 /* final null byte */)
+	length -= 4 // length bytes
+	length--    // final null byte
+
+	if n == 0 {
+		return "", true
+	}
 
 	var buf strings.Builder
 	buf.WriteByte('[')
@@ -113,19 +109,25 @@ func (a Array) stringN(n int) (string, bool) {
 	var str string
 	first := true
 	for length > 0 && !truncated {
-		l := 0
+		needStrLen := -1
+		// Set needStrLen if n is positive, meaning we want to limit the string length.
 		if n > 0 {
+			// Stop stringifying if we reach the limit, that also ensures needStrLen is
+			// greater than 0 if we need to limit the length.
 			if buf.Len() >= n {
 				truncated = true
 				break
 			}
-			l = n - buf.Len()
+			needStrLen = n - buf.Len()
 		}
+
+		// Append a comma if this is not the first element.
 		if !first {
 			buf.WriteByte(',')
-			if l > 0 {
-				l--
-				if l == 0 {
+			// If we are truncating, we need to account for the comma in the length.
+			if needStrLen > 0 {
+				needStrLen--
+				if needStrLen == 0 {
 					truncated = true
 					break
 				}
@@ -134,18 +136,22 @@ func (a Array) stringN(n int) (string, bool) {
 
 		elem, rem, ok = ReadElement(rem)
 		length -= int32(len(elem))
+		// Exit on malformed element.
 		if !ok || length < 0 {
 			return "", false
 		}
 
-		str, truncated = elem.Value().stringN(l)
+		// Delegate to StringN() on the element.
+		str, truncated = elem.Value().StringN(needStrLen)
 		buf.WriteString(str)
 
 		first = false
 	}
 
 	if n <= 0 || (buf.Len() < n && !truncated) {
 		buf.WriteByte(']')
+	} else {
+		truncated = true
 	}
 
 	return buf.String(), truncated

x/bsonx/bsoncore/array_test.go

@@ -349,76 +349,78 @@ func TestArray(t *testing.T) {
 	})
 }
 
-func TestArray_Stringer(t *testing.T) {
-	testCases := []struct {
-		description string
-		array       Array
-		want        string
-	}{
-		{
-			description: "empty array",
-			array:       BuildArray(nil),
-			want:        `[]`,
-		},
-		{
-			description: "array with 1 element",
-			array: BuildArray(nil, Value{
+var arrayStringTestCases = []struct {
+	description string
+	array       Array
+	want        string
+}{
+	{
+		description: "empty array",
+		array:       BuildArray(nil),
+		want:        `[]`,
+	},
+	{
+		description: "array with 1 element",
+		array: BuildArray(nil, Value{
+			Type: TypeInt32,
+			Data: AppendInt32(nil, 123),
+		}),
+		want: `[{"$numberInt":"123"}]`,
+	},
+	{
+		description: "nested array",
+		array: BuildArray(nil, Value{
+			Type: TypeArray,
+			Data: BuildArray(nil, Value{
+				Type: TypeString,
+				Data: AppendString(nil, "abc"),
+			}),
+		}),
+		want: `[["abc"]]`,
+	},
+	{
+		description: "array with mixed types",
+		array: BuildArray(nil,
+			Value{
+				Type: TypeString,
+				Data: AppendString(nil, "abc"),
+			},
+			Value{
 				Type: TypeInt32,
 				Data: AppendInt32(nil, 123),
-			}),
-			want: `[{"$numberInt":"123"}]`,
-		},
-		{
-			description: "nested array",
-			array: BuildArray(nil, Value{
-				Type: TypeArray,
-				Data: BuildArray(nil, Value{
-					Type: TypeString,
-					Data: AppendString(nil, "abc"),
-				}),
-			}),
-			want: `[["abc"]]`,
-		},
-		{
-			description: "array with mixed types",
-			array: BuildArray(nil,
-				Value{
-					Type: TypeString,
-					Data: AppendString(nil, "abc"),
-				},
-				Value{
-					Type: TypeInt32,
-					Data: AppendInt32(nil, 123),
-				},
-				Value{
-					Type: TypeBoolean,
-					Data: AppendBoolean(nil, true),
-				},
-			),
-			want: `["abc",{"$numberInt":"123"},true]`,
-		},
-	}
+			},
+			Value{
+				Type: TypeBoolean,
+				Data: AppendBoolean(nil, true),
+			},
+		),
+		want: `["abc",{"$numberInt":"123"},true]`,
+	},
+}
 
-	for _, tc := range testCases {
-		t.Run(fmt.Sprintf("String %s", tc.description), func(t *testing.T) {
+func TestArray_String(t *testing.T) {
+	for _, tc := range arrayStringTestCases {
+		t.Run(tc.description, func(t *testing.T) {
 			got := tc.array.String()
-			assert.Equal(t, tc.want, got)
+			assert.Equal(t, tc.want, got, "expected string %s, got %s", tc.want, got)
 		})
 	}
+}
 
-	for _, tc := range testCases {
+func TestArray_StringN(t *testing.T) {
+	for _, tc := range arrayStringTestCases {
 		for n := -1; n <= len(tc.want)+1; n++ {
-			t.Run(fmt.Sprintf("StringN %s n==%d", tc.description, n), func(t *testing.T) {
-				got, _ := tc.array.StringN(n)
+			t.Run(fmt.Sprintf("%s n==%d", tc.description, n), func(t *testing.T) {
+				got, truncated := tc.array.StringN(n)
 				l := n
-				if l < 0 {
-					l = 0
-				}
-				if l > len(tc.want) {
+				toBeTruncated := true
+				if l >= len(tc.want) || l < 0 {
 					l = len(tc.want)
+					toBeTruncated = false
 				}
 				want := tc.want[:l]
-				assert.Equal(t, want, got, "got %v, want %v", got, want)
+				assert.Equal(t, want, got, "expected truncated string %s, got %s", want, got)
+				assert.Equal(t, toBeTruncated, truncated, "expected truncated to be %t, got %t", toBeTruncated, truncated)
 			})
 		}
 	}

x/bsonx/bsoncore/document.go

@@ -261,28 +261,24 @@ func (d Document) DebugString() string {
 // String outputs an ExtendedJSON version of Document. If the document is not valid, this method
 // returns an empty string.
 func (d Document) String() string {
-	str, _ := d.stringN(0)
+	str, _ := d.StringN(-1)
 	return str
 }
 
-// StringN stringifies a document upto N bytes
+// StringN stringifies a document. If N is non-negative, it will truncate the string to N bytes.
+// Otherwise, it will return the full string representation. The second return value indicates
+// whether the string was truncated or not.
 func (d Document) StringN(n int) (string, bool) {
-	if n <= 0 {
-		if l, _, ok := ReadLength(d); !ok || l < 5 {
-			return "", false
-		}
-		return "", true
-	}
-	return d.stringN(n)
-}
-
-// stringN stringify a document. If N is larger than 0, it will truncate the string to N bytes.
-func (d Document) stringN(n int) (string, bool) {
 	length, rem, ok := ReadLength(d)
 	if !ok || length < 5 {
 		return "", false
 	}
-	length -= (4 /* length bytes */ + 1 /* final null byte */)
+	length -= 4 // length bytes
+	length--    // final null byte
+
+	if n == 0 {
+		return "", true
+	}
 
 	var buf strings.Builder
 	buf.WriteByte('{')
@@ -292,19 +288,25 @@ func (d Document) stringN(n int) (string, bool) {
 	var str string
 	first := true
 	for length > 0 && !truncated {
-		l := 0
+		needStrLen := -1
+		// Set needStrLen if n is positive, meaning we want to limit the string length.
 		if n > 0 {
+			// Stop stringifying if we reach the limit, that also ensures needStrLen is
+			// greater than 0 if we need to limit the length.
 			if buf.Len() >= n {
 				truncated = true
 				break
 			}
-			l = n - buf.Len()
+			needStrLen = n - buf.Len()
 		}
+
+		// Append a comma if this is not the first element.
 		if !first {
 			buf.WriteByte(',')
-			if l > 0 {
-				l--
-				if l == 0 {
+			// If we are truncating, we need to account for the comma in the length.
+			if needStrLen > 0 {
+				needStrLen--
+				if needStrLen == 0 {
 					truncated = true
 					break
 				}
@@ -313,18 +315,22 @@ func (d Document) stringN(n int) (string, bool) {
 
 		elem, rem, ok = ReadElement(rem)
 		length -= int32(len(elem))
+		// Exit on malformed element.
 		if !ok || length < 0 {
 			return "", false
 		}
 
-		str, truncated = elem.stringN(l)
+		// Delegate to StringN() on the element.
+		str, truncated = elem.StringN(needStrLen)
 		buf.WriteString(str)
 
 		first = false
 	}
 
 	if n <= 0 || (buf.Len() < n && !truncated) {
 		buf.WriteByte('}')
+	} else {
+		truncated = true
 	}
 
 	return buf.String(), truncated