GODRIVER-3605 Refactor StringN (#2128)

Author: qingyang-hu

bson/raw_test.go

@@ -450,7 +450,7 @@ func BenchmarkRawString(b *testing.B) {
 			b.ReportAllocs()
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				_ = bsoncore.Document(bs).StringN(1024) // Assuming you want to limit to 1024 bytes for this benchmark
+				_, _ = bsoncore.Document(bs).StringN(1024) // Assuming you want to limit to 1024 bytes for this benchmark
 			}
 		})
 	}
@@ -473,7 +473,7 @@ func TestComplexDocuments_StringN(t *testing.T) {
 			bson, _ := Marshal(tc.doc)
 			bsonDoc := bsoncore.Document(bson)
 
-			got := bsonDoc.StringN(tc.n)
+			got, _ := bsonDoc.StringN(tc.n)
 			assert.Equal(t, tc.n, len(got))
 		})
 	}
@@ -518,7 +518,7 @@ func createMassiveArraysDocument(arraySize int) D {
 func createUniqueVoluminousDocument(t *testing.T, size int) bsoncore.Document {
 	t.Helper()
 
-	docs := make(D, size)
+	docs := make(D, 0, size)
 
 	for i := 0; i < size; i++ {
 		docs = append(docs, E{
@@ -561,7 +561,7 @@ func createLargeSingleDoc(t *testing.T) bsoncore.Document {
 func createVoluminousArrayDocuments(t *testing.T, size int) bsoncore.Document {
 	t.Helper()
 
-	docs := make(D, size)
+	docs := make(D, 0, size)
 
 	for i := 0; i < size; i++ {
 		docs = append(docs, E{

internal/bsoncoreutil/bsoncoreutil.go

@@ -8,7 +8,7 @@ package bsoncoreutil
 
 // Truncate truncates a given string for a certain width
 func Truncate(str string, width int) string {
-	if width == 0 {
+	if width <= 0 {
 		return ""
 	}
 

internal/logger/logger.go

@@ -241,10 +241,9 @@ func FormatDocument(msg bson.Raw, width uint) string {
 		return "{}"
 	}
 
-	str := bsoncore.Document(msg).StringN(int(width))
+	str, truncated := bsoncore.Document(msg).StringN(int(width))
 
-	// If the last byte is not a closing bracket, then the document was truncated
-	if len(str) > 0 && str[len(str)-1] != '}' {
+	if truncated {
 		str += TruncationSuffix
 	}
 

x/bsonx/bsoncore/array.go

@@ -9,7 +9,6 @@ package bsoncore
 import (
 	"fmt"
 	"io"
-	"math"
 	"strconv"
 	"strings"
 )
@@ -83,55 +82,79 @@ 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 {
-	return a.StringN(math.MaxInt)
+	str, _ := a.StringN(-1)
+	return str
 }
 
-// StringN stringifies an array upto N bytes
-func (a Array) StringN(n int) string {
-	if lens, _, _ := ReadLength(a); lens < 5 || n <= 0 {
-		return ""
+// 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) {
+	length, rem, ok := ReadLength(a)
+	if !ok || length < 5 {
+		return "", false
+	}
+	length -= 4 // length bytes
+	length--    // final null byte
+
+	if n == 0 {
+		return "", true
 	}
 
 	var buf strings.Builder
 	buf.WriteByte('[')
 
-	length, rem, _ := ReadLength(a) // We know we have enough bytes to read the length
-	length -= 4
-
+	var truncated bool
 	var elem Element
-	var ok bool
-
-	if n > 0 {
-		for length > 1 {
-			elem, rem, ok = ReadElement(rem)
-
-			length -= int32(len(elem))
-			if !ok {
-				return ""
-			}
-
-			str := elem.Value().StringN(n - buf.Len())
-
-			buf.WriteString(str)
-
-			if buf.Len() == n {
-				return buf.String()
+	var str string
+	first := true
+	for length > 0 && !truncated {
+		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
 			}
+			needStrLen = n - buf.Len()
+		}
 
-			if length > 1 {
-				buf.WriteByte(',')
+		// Append a comma if this is not the first element.
+		if !first {
+			buf.WriteByte(',')
+			// If we are truncating, we need to account for the comma in the length.
+			if needStrLen > 0 {
+				needStrLen--
+				if needStrLen == 0 {
+					truncated = true
+					break
+				}
 			}
 		}
-		if length != 1 { // Missing final null byte or inaccurate length
-			return ""
+
+		elem, rem, ok = ReadElement(rem)
+		length -= int32(len(elem))
+		// Exit on malformed element.
+		if !ok || length < 0 {
+			return "", false
 		}
+
+		// Delegate to StringN() on the element.
+		str, truncated = elem.Value().StringN(needStrLen)
+		buf.WriteString(str)
+
+		first = false
 	}
 
-	if buf.Len()+1 <= n {
+	if n <= 0 || (buf.Len() < n && !truncated) {
 		buf.WriteByte(']')
+	} else {
+		truncated = true
 	}
 
-	return buf.String()
+	return buf.String(), truncated
 }
 
 // Values returns this array as a slice of values. The returned slice will contain valid values.