Reduce number of allocations during serializing

Author: inteon

Diff for: fieldpath/path.go

@@ -80,7 +80,7 @@ func (fp Path) Copy() Path {
 
 // MakePath constructs a Path. The parts may be PathElements, ints, strings.
 func MakePath(parts ...interface{}) (Path, error) {
-	var fp Path
+	fp := make(Path, 0, len(parts))
 	for _, p := range parts {
 		switch t := p.(type) {
 		case PathElement:

Diff for: fieldpath/serialize-pe.go

@@ -19,10 +19,9 @@ package fieldpath
 import (
 	"errors"
 	"fmt"
-	"io"
 	"strconv"
-	"strings"
 
+	"sigs.k8s.io/structured-merge-diff/v4/internal/builder"
 	"sigs.k8s.io/structured-merge-diff/v4/value"
 )
 
@@ -57,7 +56,7 @@ var (
 func DeserializePathElement(s string) (PathElement, error) {
 	b := []byte(s)
 	if len(b) < 2 {
-		return PathElement{}, errors.New("key must be 2 characters long:")
+		return PathElement{}, errors.New("key must be 2 characters long")
 	}
 	typeSep, b := b[:2], b[2:]
 	if typeSep[1] != peSepBytes[0] {
@@ -99,41 +98,37 @@ func DeserializePathElement(s string) (PathElement, error) {
 
 // SerializePathElement serializes a path element
 func SerializePathElement(pe PathElement) (string, error) {
-	buf := strings.Builder{}
+	buf := builder.JSONBuilder{}
 	err := serializePathElementToWriter(&buf, pe)
 	return buf.String(), err
 }
 
-func serializePathElementToWriter(w io.Writer, pe PathElement) error {
+func serializePathElementToWriter(w *builder.JSONBuilder, pe PathElement) error {
 	switch {
 	case pe.FieldName != nil:
 		if _, err := w.Write(peFieldSepBytes); err != nil {
 			return err
 		}
-		fmt.Fprintf(w, "%s", *pe.FieldName)
+		w.WriteString(*pe.FieldName)
 	case pe.Key != nil:
 		if _, err := w.Write(peKeySepBytes); err != nil {
 			return err
 		}
-		jsonVal, err := value.FieldListToJSON(*pe.Key)
-		if err != nil {
+		if err := value.FieldListToJSON(*pe.Key, w); err != nil {
 			return err
 		}
-		fmt.Fprintf(w, "%s", jsonVal)
 	case pe.Value != nil:
 		if _, err := w.Write(peValueSepBytes); err != nil {
 			return err
 		}
-		jsonVal, err := value.ToJSON(*pe.Value)
-		if err != nil {
+		if err := value.ToJSON(*pe.Value, w); err != nil {
 			return err
 		}
-		fmt.Fprintf(w, "%s", jsonVal)
 	case pe.Index != nil:
 		if _, err := w.Write(peIndexSepBytes); err != nil {
 			return err
 		}
-		fmt.Fprintf(w, "%d", *pe.Index)
+		w.WriteString(strconv.Itoa(*pe.Index))
 	default:
 		return errors.New("invalid PathElement")
 	}

Diff for: fieldpath/serialize.go

@@ -17,34 +17,36 @@ limitations under the License.
 package fieldpath
 
 import (
-	"bytes"
-	gojson "encoding/json"
 	"fmt"
 	"io"
+	"sort"
 
 	json "sigs.k8s.io/json"
+	"sigs.k8s.io/structured-merge-diff/v4/internal/builder"
 )
 
 func (s *Set) ToJSON() ([]byte, error) {
-	buf := bytes.Buffer{}
-	err := s.ToJSONStream(&buf)
-	if err != nil {
+	buf := builder.JSONBuilder{}
+	if err := s.emitContentsV1(false, &buf, &builder.ReusableBuilder{}); err != nil {
 		return nil, err
 	}
 	return buf.Bytes(), nil
 }
 
 func (s *Set) ToJSONStream(w io.Writer) error {
-	err := s.emitContentsV1(false, w)
-	if err != nil {
+	buf := builder.JSONBuilder{}
+	if err := s.emitContentsV1(false, &buf, &builder.ReusableBuilder{}); err != nil {
 		return err
 	}
-	return nil
+	_, err := buf.WriteTo(w)
+	return err
 }
 
 type orderedMapItemWriter struct {
-	w        io.Writer
+	w        *builder.JSONBuilder
 	hasItems bool
+
+	builder *builder.ReusableBuilder
 }
 
 // writeKey writes a key to the writer, including a leading comma if necessary.
@@ -74,16 +76,24 @@ func (om *orderedMapItemWriter) writeKey(key []byte) error {
 // After writing the key, the caller should write the encoded value, e.g. using
 // writeEmptyValue or by directly writing the value to the writer.
 func (om *orderedMapItemWriter) writePathKey(pe PathElement) error {
-	pev, err := SerializePathElement(pe)
-	if err != nil {
+	if om.hasItems {
+		if _, err := om.w.Write([]byte{','}); err != nil {
+			return err
+		}
+	}
+
+	if err := serializePathElementToWriter(om.builder.Reset(), pe); err != nil {
 		return err
 	}
-	key, err := gojson.Marshal(pev)
-	if err != nil {
+	if err := om.w.WriteJSON(om.builder.UnsafeString()); err != nil {
 		return err
 	}
 
-	return om.writeKey(key)
+	if _, err := om.w.Write([]byte{':'}); err != nil {
+		return err
+	}
+	om.hasItems = true
+	return nil
 }
 
 // writeEmptyValue writes an empty JSON object to the writer.
@@ -95,11 +105,11 @@ func (om orderedMapItemWriter) writeEmptyValue() error {
 	return nil
 }
 
-func (s *Set) emitContentsV1(includeSelf bool, w io.Writer) error {
-	om := orderedMapItemWriter{w: w}
+func (s *Set) emitContentsV1(includeSelf bool, w *builder.JSONBuilder, r *builder.ReusableBuilder) error {
+	om := orderedMapItemWriter{w: w, builder: r}
 	mi, ci := 0, 0
 
-	if _, err := om.w.Write([]byte{'{'}); err != nil {
+	if _, err := w.Write([]byte{'{'}); err != nil {
 		return err
 	}
 
@@ -129,7 +139,7 @@ func (s *Set) emitContentsV1(includeSelf bool, w io.Writer) error {
 			if err := om.writePathKey(cpe); err != nil {
 				return err
 			}
-			if err := s.Children.members[ci].set.emitContentsV1(c == 0, om.w); err != nil {
+			if err := s.Children.members[ci].set.emitContentsV1(c == 0, w, r); err != nil {
 				return err
 			}
 
@@ -160,14 +170,14 @@ func (s *Set) emitContentsV1(includeSelf bool, w io.Writer) error {
 		if err := om.writePathKey(cpe); err != nil {
 			return err
 		}
-		if err := s.Children.members[ci].set.emitContentsV1(false, om.w); err != nil {
+		if err := s.Children.members[ci].set.emitContentsV1(false, w, r); err != nil {
 			return err
 		}
 
 		ci++
 	}
 
-	if _, err := om.w.Write([]byte{'}'}); err != nil {
+	if _, err := w.Write([]byte{'}'}); err != nil {
 		return err
 	}
 
@@ -237,34 +247,33 @@ func readIterV1(data []byte) (children *Set, isMember bool, err error) {
 				children = &Set{}
 			}
 
+			// Append the member to the members list, we will sort it later
 			m := &children.Members.members
-			// Since we expect that most of the time these will have been
-			// serialized in the right order, we just verify that and append.
-			appendOK := len(*m) == 0 || (*m)[len(*m)-1].Less(pe)
-			if appendOK {
-				*m = append(*m, pe)
-			} else {
-				children.Members.Insert(pe)
-			}
+			*m = append(*m, pe)
 		}
 
 		if v.target != nil {
 			if children == nil {
 				children = &Set{}
 			}
 
-			// Since we expect that most of the time these will have been
-			// serialized in the right order, we just verify that and append.
+			// Append the child to the children list, we will sort it later
 			m := &children.Children.members
-			appendOK := len(*m) == 0 || (*m)[len(*m)-1].pathElement.Less(pe)
-			if appendOK {
-				*m = append(*m, setNode{pe, v.target})
-			} else {
-				*children.Children.Descend(pe) = *v.target
-			}
+			*m = append(*m, setNode{pe, v.target})
 		}
 	}
 
+	// Sort the members and children
+	if children != nil {
+		sort.Slice(children.Members.members, func(i, j int) bool {
+			return children.Members.members[i].Less(children.Members.members[j])
+		})
+
+		sort.Slice(children.Children.members, func(i, j int) bool {
+			return children.Children.members[i].pathElement.Less(children.Children.members[j].pathElement)
+		})
+	}
+
 	if children == nil {
 		isMember = true
 	}

Diff for: internal/builder/builder.go

@@ -0,0 +1,67 @@
+package builder
+
+import (
+	"bytes"
+	gojson "encoding/json"
+	"runtime"
+	"unsafe"
+)
+
+type ReusableBuilder struct {
+	JSONBuilder
+}
+
+func (r *ReusableBuilder) Reset() *JSONBuilder {
+	r.JSONBuilder.Reset()
+	return &r.JSONBuilder
+}
+
+func (r *ReusableBuilder) UnsafeString() string {
+	b := r.Bytes()
+	return *(*string)(unsafe.Pointer(&b))
+}
+
+type JSONBuilder struct {
+	initialised bool
+	bytes.Buffer
+	gojson.Encoder
+}
+
+type noNewlineWriter struct {
+	*bytes.Buffer
+}
+
+func (w noNewlineWriter) Write(p []byte) (n int, err error) {
+	if len(p) > 0 && p[len(p)-1] == '\n' {
+		p = p[:len(p)-1]
+	}
+	return w.Buffer.Write(p)
+}
+
+// noescape hides a pointer from escape analysis. It is the identity function
+// but escape analysis doesn't think the output depends on the input.
+// noescape is inlined and currently compiles down to zero instructions.
+// USE CAREFULLY!
+// This was copied from the runtime; see issues 23382 and 7921.
+//
+//go:nosplit
+//go:nocheckptr
+func noescape(p unsafe.Pointer) unsafe.Pointer {
+	x := uintptr(p)
+	return unsafe.Pointer(x ^ 0) //nolint:unsafeptr
+}
+
+func (r *JSONBuilder) WriteJSON(v interface{}) error {
+	if !r.initialised {
+		r.Encoder = *gojson.NewEncoder(noNewlineWriter{&r.Buffer})
+		r.initialised = true
+	}
+
+	if err := r.Encoder.Encode((*interface{})(noescape(unsafe.Pointer(&v)))); err != nil {
+		return err
+	}
+
+	runtime.KeepAlive(v)
+
+	return nil
+}

Diff for: value/fields.go

@@ -17,11 +17,11 @@ limitations under the License.
 package value
 
 import (
-	gojson "encoding/json"
 	"sort"
 	"strings"
 
 	"sigs.k8s.io/json"
+	"sigs.k8s.io/structured-merge-diff/v4/internal/builder"
 )
 
 // Field is an individual key-value pair.
@@ -50,12 +50,22 @@ func FieldListFromJSON(input []byte) (FieldList, error) {
 }
 
 // FieldListToJSON is a helper function for producing a JSON document.
-func FieldListToJSON(v FieldList) ([]byte, error) {
-	m := make(map[string]interface{}, len(v))
-	for _, f := range v {
-		m[f.Name] = f.Value.Unstructured()
+func FieldListToJSON(v FieldList, w *builder.JSONBuilder) error {
+	w.WriteByte('{')
+	for i, f := range v {
+		if err := w.WriteJSON(f.Name); err != nil {
+			return err
+		}
+		w.WriteByte(':')
+		if err := w.WriteJSON(f.Value.Unstructured()); err != nil {
+			return err
+		}
+		if i < len(v)-1 {
+			w.WriteByte(',')
+		}
 	}
-	return gojson.Marshal(m)
+	w.WriteByte('}')
+	return nil
 }
 
 // Sort sorts the field list by Name.

Diff for: value/value.go

@@ -17,12 +17,12 @@ limitations under the License.
 package value
 
 import (
-	gojson "encoding/json"
 	"fmt"
 	"strings"
 
 	"gopkg.in/yaml.v2"
 	"sigs.k8s.io/json"
+	"sigs.k8s.io/structured-merge-diff/v4/internal/builder"
 )
 
 // A Value corresponds to an 'atom' in the schema. It should return true
@@ -85,8 +85,8 @@ func FromJSON(input []byte) (Value, error) {
 }
 
 // ToJSON is a helper function for producing a JSON document.
-func ToJSON(v Value) ([]byte, error) {
-	return gojson.Marshal(v.Unstructured())
+func ToJSON(v Value, w *builder.JSONBuilder) error {
+	return w.WriteJSON(v.Unstructured())
 }
 
 // ToYAML marshals a value as YAML.