Refactor StringN

Author: qingyang-hu

bson/raw_test.go

@@ -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 ""
 	}
 

x/bsonx/bsoncore/array.go

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

x/bsonx/bsoncore/array_test.go

@@ -506,7 +506,7 @@ func TestArray_StringN(t *testing.T) {
 		},
 		{
 			description: "n>0, array EQ n",
-			n:           22,
+			n:           20,
 			values: []Value{
 				{
 					Type: TypeInt32,
@@ -576,7 +576,7 @@ func TestArray_StringN(t *testing.T) {
 			got := Array(BuildArray(nil, tc.values...)).StringN(tc.n)
 			assert.Equal(t, tc.want, got)
 			if tc.n >= 0 {
-				assert.LessOrEqual(t, len(got), tc.n)
+				assert.LessOrEqual(t, len(got), tc.n, "got %v, want %v", got, tc.want)
 			}
 		})
 	}

x/bsonx/bsoncore/document.go

@@ -10,11 +10,8 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"math"
 	"strconv"
 	"strings"
-
-	"go.mongodb.org/mongo-driver/v2/internal/bsoncoreutil"
 )
 
 // ValidationError is an error type returned when attempting to validate a document or array.
@@ -264,58 +261,73 @@ 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 {
-	return d.StringN(math.MaxInt)
+	str, _ := d.stringN(0)
+	return str
 }
 
 // StringN stringifies a document upto N bytes
 func (d Document) StringN(n int) string {
-	if len(d) < 5 || n <= 0 {
+	if n <= 0 {
 		return ""
 	}
+	str, _ := d.stringN(n)
+	return str
+}
 
-	var buf strings.Builder
+// 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) {
+	if len(d) < 5 {
+		return "", false
+	}
 
+	var buf strings.Builder
 	buf.WriteByte('{')
 
 	length, rem, _ := ReadLength(d)
-	length -= 4
+	length -= (4 /* length bytes */ + 1 /* final null byte */)
 
+	var truncated bool
 	var elem Element
 	var ok bool
-
+	var str string
 	first := true
-	truncated := false
-
-	if n > 0 {
-		for length > 1 {
-			if !first {
-				buf.WriteByte(',')
-			}
-			elem, rem, ok = ReadElement(rem)
-			length -= int32(len(elem))
-			if !ok {
-				return ""
-			}
-
-			str := elem.StringN(n)
-			if buf.Len()+len(str) > n {
-				truncatedStr := bsoncoreutil.Truncate(str, n-buf.Len())
-				buf.WriteString(truncatedStr)
-
+	for length > 0 && !truncated {
+		l := 0
+		if n > 0 {
+			if buf.Len() >= n {
 				truncated = true
 				break
 			}
+			l = n - buf.Len()
+		}
+		if !first {
+			buf.WriteByte(',')
+			if l > 0 {
+				l--
+				if l == 0 {
+					truncated = true
+					break
+				}
+			}
+		}
 
-			buf.WriteString(str)
-			first = false
+		elem, rem, ok = ReadElement(rem)
+		length -= int32(len(elem))
+		if !ok || length < 0 {
+			return "", true
 		}
+
+		str, truncated = elem.stringN(l)
+		buf.WriteString(str)
+
+		first = false
 	}
 
-	if !truncated {
+	if n <= 0 || (buf.Len() < n && !truncated) {
 		buf.WriteByte('}')
 	}
 
-	return buf.String()
+	return buf.String(), truncated
 }
 
 // Elements returns this document as a slice of elements. The returned slice will contain valid

x/bsonx/bsoncore/document_test.go

@@ -532,7 +532,7 @@ func TestDocument_StringN(t *testing.T) {
 			got := bs.StringN(tc.n)
 			assert.Equal(t, tc.want, got)
 			if tc.n >= 0 {
-				assert.LessOrEqual(t, len(got), tc.n)
+				assert.LessOrEqual(t, len(got), tc.n, "got %v, want %v", got, tc.want)
 			}
 		})
 	}

x/bsonx/bsoncore/element.go

@@ -9,7 +9,9 @@ package bsoncore
 import (
 	"bytes"
 	"fmt"
-	"math"
+	"strings"
+
+	"go.mongodb.org/mongo-driver/v2/internal/bsoncoreutil"
 )
 
 // MalformedElementError represents a class of errors that RawElement methods return.
@@ -115,35 +117,78 @@ func (e Element) ValueErr() (Value, error) {
 
 // String implements the fmt.String interface. The output will be in extended JSON format.
 func (e Element) String() string {
-	return e.StringN(math.MaxInt)
+	str, _ := e.stringN(0)
+	return str
 }
 
 // StringN implements the fmt.String interface for upto N bytes. The output will be in extended JSON format.
 func (e Element) StringN(n int) string {
-	if len(e) == 0 {
+	if n <= 0 {
 		return ""
 	}
+	str, _ := e.stringN(n)
+	return str
+}
+
+// stringN stringify an element. If N is larger than 0, it will truncate the string to N bytes.
+func (e Element) stringN(n int) (string, bool) {
+	if len(e) == 0 {
+		return "", false
+	}
+	if n == 1 {
+		return `"`, true
+	}
+
 	t := Type(e[0])
 	idx := bytes.IndexByte(e[1:], 0x00)
-	if idx == -1 {
-		return ""
-	}
-	key, valBytes := []byte(e[1:idx+1]), []byte(e[idx+2:])
-	val, _, valid := ReadValue(valBytes, t)
+	if idx <= 0 {
+		return "", false
+	}
+	key := e[1 : idx+1]
+
+	var buf strings.Builder
+	buf.WriteByte('"')
+	const postfix = `": `
+	switch {
+	case n <= 0 || idx <= n-4:
+		buf.Write(key)
+		buf.WriteString(postfix)
+	case idx < n:
+		buf.Write(key)
+		buf.WriteString(postfix[:n-idx])
+	default:
+		buf.WriteString(bsoncoreutil.Truncate(string(key), n-1))
+	}
+
+	l := 0
+	if n > 0 {
+		if buf.Len() >= n {
+			return buf.String(), true
+		}
+		l = n - buf.Len()
+	}
+
+	val, _, valid := ReadValue(e[idx+2:], t)
 	if !valid {
-		return ""
+		return "", false
 	}
 
 	var str string
+	var truncated bool
 	if _, ok := val.StringValueOK(); ok {
-		str = val.StringN(n)
+		str, truncated = val.stringN(l)
 	} else if arr, ok := val.ArrayOK(); ok {
-		str = arr.StringN(n)
+		str, truncated = arr.stringN(l)
 	} else {
 		str = val.String()
+		if l > 0 && len(str) > l {
+			truncated = true
+			str = bsoncoreutil.Truncate(str, l)
+		}
 	}
 
-	return "\"" + string(key) + "\": " + str
+	buf.WriteString(str)
+	return buf.String(), truncated
 }
 
 // DebugString outputs a human readable version of RawElement. It will attempt to stringify the