Repository: fxamacker/cbor
Pull request #519 (2024-04-12 19:12:15)

Encode structs directly to output buffer.
<!--
Thank you for your interest in contributing to fxamacker/cbor!
-->

### Description

For variable-length structs (structs with omitempty fields), encoding to the unused capacity at the
end of the output buffer while counting nonempty items is cheaper than using a separate temporary
buffer (no pool interactions and better spatial locality). Copying the items can be avoided entirely
by reserving space in the output buffer for the head if the encoded length of the head can be
predicted before checking optional fields.

```
                                                                     │ before.txt  │              after.txt              │
                                                                     │   sec/op    │   sec/op     vs base                │
Marshal/Go_struct_to_CBOR_map                                          1.404µ ± 0%   1.408µ ± 1%        ~ (p=0.170 n=10)
Marshal/Go_struct_many_fields_all_omitempty_all_empty_to_CBOR_map      443.8n ± 0%   430.6n ± 0%   -2.99% (p=0.000 n=10)
Marshal/Go_struct_some_fields_all_omitempty_all_empty_to_CBOR_map      181.7n ± 0%   163.5n ± 0%  -10.04% (p=0.000 n=10)
Marshal/Go_struct_many_fields_all_omitempty_all_nonempty_to_CBOR_map   813.5n ± 0%   784.8n ± 0%   -3.53% (p=0.000 n=10)
Marshal/Go_struct_some_fields_all_omitempty_all_nonempty_to_CBOR_map   300.8n ± 0%   275.4n ± 0%   -8.43% (p=0.000 n=10)
Marshal/Go_struct_many_fields_one_omitempty_to_CBOR_map                763.8n ± 0%   727.7n ± 0%   -4.73% (p=0.000 n=10)
Marshal/Go_struct_some_fields_one_omitempty_to_CBOR_map                284.2n ± 0%   257.6n ± 0%   -9.36% (p=0.000 n=10)
Marshal/Go_struct_keyasint_to_CBOR_map                                 1.422µ ± 0%   1.414µ ± 1%   -0.56% (p=0.029 n=10)
Marshal/Go_struct_toarray_to_CBOR_array                                1.341µ ± 1%   1.338µ ± 1%        ~ (p=0.340 n=10)
MarshalCanonical/Go_struct_to_CBOR_map                                 386.4n ± 0%   392.4n ± 0%   +1.57% (p=0.000 n=10)
MarshalCanonical/Go_struct_to_CBOR_map_canonical                       386.9n ± 0%   384.8n ± 0%   -0.52% (p=0.001 n=10)
geomean                                                                560.5n        540.4n        -3.59%

                                                                     │ before.txt │              after.txt              │
                                                                     │    B/op    │    B/op     vs base                 │
Marshal/Go_struct_to_CBOR_map                                          208.0 ± 0%   208.0 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_many_fields_all_omitempty_all_empty_to_CBOR_map      1.000 ± 0%   1.000 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_some_fields_all_omitempty_all_empty_to_CBOR_map      1.000 ± 0%   1.000 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_many_fields_all_omitempty_all_nonempty_to_CBOR_map   176.0 ± 0%   176.0 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_some_fields_all_omitempty_all_nonempty_to_CBOR_map   48.00 ± 0%   48.00 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_many_fields_one_omitempty_to_CBOR_map                160.0 ± 0%   160.0 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_some_fields_one_omitempty_to_CBOR_map                48.00 ± 0%   48.00 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_keyasint_to_CBOR_map                                 192.0 ± 0%   192.0 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_toarray_to_CBOR_array                                192.0 ± 0%   192.0 ± 0%       ~ (p=1.000 n=10) ¹
MarshalCanonical/Go_struct_to_CBOR_map                                 64.00 ± 0%   64.00 ± 0%       ~ (p=1.000 n=10) ¹
MarshalCanonical/Go_struct_to_CBOR_map_canonical                       64.00 ± 0%   64.00 ± 0%       ~ (p=1.000 n=10) ¹
geomean                                                                46.18        46.18       +0.00%
¹ all samples are equal

                                                                     │ before.txt │              after.txt              │
                                                                     │ allocs/op  │ allocs/op   vs base                 │
Marshal/Go_struct_to_CBOR_map                                          1.000 ± 0%   1.000 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_many_fields_all_omitempty_all_empty_to_CBOR_map      1.000 ± 0%   1.000 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_some_fields_all_omitempty_all_empty_to_CBOR_map      1.000 ± 0%   1.000 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_many_fields_all_omitempty_all_nonempty_to_CBOR_map   1.000 ± 0%   1.000 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_some_fields_all_omitempty_all_nonempty_to_CBOR_map   1.000 ± 0%   1.000 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_many_fields_one_omitempty_to_CBOR_map                1.000 ± 0%   1.000 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_some_fields_one_omitempty_to_CBOR_map                1.000 ± 0%   1.000 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_keyasint_to_CBOR_map                                 1.000 ± 0%   1.000 ± 0%       ~ (p=1.000 n=10) ¹
Marshal/Go_struct_toarray_to_CBOR_array                                1.000 ± 0%   1.000 ± 0%       ~ (p=1.000 n=10) ¹
MarshalCanonical/Go_struct_to_CBOR_map                                 1.000 ± 0%   1.000 ± 0%       ~ (p=1.000 n=10) ¹
MarshalCanonical/Go_struct_to_CBOR_map_canonical                       1.000 ± 0%   1.000 ± 0%       ~ (p=1.000 n=10) ¹
geomean                                                                1.000        1.000       +0.00%
¹ all samples are equal
```

<!-- For code contributions, please complete all the items below this line. -->
<!-- For documentation-only contributions, please delete everything below this line. -->

#### PR Was Proposed and Welcomed in Currently Open Issue

- [ ] This PR was proposed and welcomed by maintainer(s) in issue #___
- [ ] Closes or Updates Issue #___

#### Checklist (for code PR only, ignore for docs PR)

- [x] Include unit tests that cover the new code
- [x] Pass all unit tests 
- [x] Pass all lint checks in CI (goimports, gosec, staticcheck, etc.)
- [x] Sign each commit with your real name and email.  
      Last line of each commit message should be in this format:  
      Signed-off-by: Firstname Lastname <firstname.lastname@example.com>
- [x] Certify the Developer's Certificate of Origin 1.1
      (see next section).

#### Certify the Developer's Certificate of Origin 1.1

- [x] By marking this item as completed, I certify 
      the Developer Certificate of Origin 1.1.

```
Developer Certificate of Origin
Version 1.1

Copyright (C) 2004, 2006 The Linux Foundation and its contributors.
660 York Street, Suite 102,
San Francisco, CA 94110 USA

Everyone is permitted to copy and distribute verbatim copies of this
license document, but changing it is not allowed.

Developer's Certificate of Origin 1.1

By making a contribution to this project, I certify that:

(a) The contribution was created in whole or in part by me and I
    have the right to submit it under the open source license
    indicated in the file; or

(b) The contribution is based upon previous work that, to the best
    of my knowledge, is covered under an appropriate open source
    license and I have the right under that license to submit that
    work with modifications, whether created in whole or in part
    by me, under the same open source license (unless I am
    permitted to submit under a different license), as indicated
    in the file; or

(c) The contribution was provided directly to me by some other
    person who certified (a), (b) or (c) and I have not modified
    it.

(d) I understand and agree that this project and the contribution
    are public and that a record of the contribution (including all
    personal information I submit with it, including my sign-off) is
    maintained indefinitely and may be redistributed consistent with
    this project or the open source license(s) involved.
```



benluddy: Instead of making two passes, it now encodes the items to the output buffer while counting, encodes the head at the end, and uses excess capacity in the output buffer to swap the positions of the encoded head and the encoded items. This turned out to be faster. Then I realized that you'll usually have a variable-length struct whose head encodes to the same number of bytes regardless of the number of items (e.g. any struct with fewer than 24 fields). In that case, it reserves space in the output buffer for the head, encodes the items, then overwrites the head bytes in the output buffer at the end, once it knows the actual map size.

Patch:
diff --git a/cache.go b/cache.go
index dca581a..24559cc 100644
--- a/cache.go
+++ b/cache.go
@@ -171,6 +171,7 @@ type encodingStructType struct {
 	err                error
 	toArray            bool
 	fixedLength        bool // Struct type doesn't have any omitempty or anonymous fields.
+	maxHeadLen         int
 }
 
 func (st *encodingStructType) getFields(em *encMode) fields {
@@ -231,6 +232,8 @@ func getEncodingStructType(t reflect.Type) (*encodingStructType, error) {
 		return getEncodingStructToArrayType(t, flds)
 	}
 
+	nOptional := 0
+
 	var err error
 	var hasKeyAsInt bool
 	var hasKeyAsStr bool
@@ -293,6 +296,10 @@ func getEncodingStructType(t reflect.Type) (*encodingStructType, error) {
 			fixedLength = false
 			omitEmptyIdx = append(omitEmptyIdx, i)
 		}
+
+		if len(flds[i].idx) > 1 || flds[i].omitEmpty {
+			nOptional++
+		}
 	}
 	putEncoderBuffer(e)
 
@@ -320,7 +327,9 @@ func getEncodingStructType(t reflect.Type) (*encodingStructType, error) {
 		lengthFirstFields:  lengthFirstFields,
 		omitEmptyFieldsIdx: omitEmptyIdx,
 		fixedLength:        fixedLength,
+		maxHeadLen:         encodedHeadLen(uint64(len(flds))),
 	}
+
 	encodingStructTypeCache.Store(t, structType)
 	return structType, structType.err
 }
diff --git a/encode.go b/encode.go
index a58c3a5..a4fdfeb 100644
--- a/encode.go
+++ b/encode.go
@@ -1297,12 +1297,17 @@ func encodeStruct(e *encoderBuffer, em *encMode, v reflect.Value) (err error) {
 		start = rand.Intn(len(flds)) //nolint:gosec // Don't need a CSPRNG for deck cutting.
 	}
 
-	kve := getEncoderBuffer() // encode key-value pairs based on struct field tag options
-	kvcount := 0
+	if b := em.encTagBytes(v.Type()); b != nil {
+		e.Write(b)
+	}
+
+	// Reserve space in the output buffer for the head if its encoded size is fixed.
+	encodeHead(e, byte(cborTypeMap), uint64(len(flds)))
 
+	kvbegin := e.Len()
+	kvcount := 0
 	for offset := 0; offset < len(flds); offset++ {
-		i := (start + offset) % len(flds)
-		f := flds[i]
+		f := flds[(start+offset)%len(flds)]
 
 		var fv reflect.Value
 		if len(f.idx) == 1 {
@@ -1320,7 +1325,6 @@ func encodeStruct(e *encoderBuffer, em *encMode, v reflect.Value) (err error) {
 		if f.omitEmpty {
 			empty, err := f.ief(em, fv)
 			if err != nil {
-				putEncoderBuffer(kve)
 				return err
 			}
 			if empty {
@@ -1329,26 +1333,42 @@ func encodeStruct(e *encoderBuffer, em *encMode, v reflect.Value) (err error) {
 		}
 
 		if !f.keyAsInt && em.fieldName == FieldNameToByteString {
-			kve.Write(f.cborNameByteString)
+			e.Write(f.cborNameByteString)
 		} else { // int or text string
-			kve.Write(f.cborName)
+			e.Write(f.cborName)
 		}
 
-		if err := f.ef(kve, em, fv); err != nil {
-			putEncoderBuffer(kve)
+		if err := f.ef(e, em, fv); err != nil {
 			return err
 		}
+
 		kvcount++
 	}
 
-	if b := em.encTagBytes(v.Type()); b != nil {
-		e.Write(b)
+	// Overwrite the bytes that were reserved for the head before encoding the map entries.
+	{
+		headbuf := encoderBuffer{Buffer: *bytes.NewBuffer(e.Bytes()[kvbegin-structType.maxHeadLen : kvbegin-structType.maxHeadLen : kvbegin])}
+		encodeHead(&headbuf, byte(cborTypeMap), uint64(kvcount))
+	}
+
+	actualHeadLen := encodedHeadLen(uint64(kvcount))
+	if actualHeadLen == structType.maxHeadLen {
+		// The bytes reserved for the encoded head were exactly the right size, so the
+		// encoded entries are already in their final positions.
+		return nil
 	}
 
-	encodeHead(e, byte(cborTypeMap), uint64(kvcount))
-	e.Write(kve.Bytes())
+	// We reserved more bytes than needed for the encoded head, based on the number of fields
+	// encoded. The encoded entries are offset to the right by the number of excess reserved
+	// bytes. Shift the entries left to remove the gap.
+	excessReservedBytes := structType.maxHeadLen - actualHeadLen
+	dst := e.Bytes()[kvbegin-excessReservedBytes : e.Len()-excessReservedBytes]
+	src := e.Bytes()[kvbegin:e.Len()]
+	copy(dst, src)
 
-	putEncoderBuffer(kve)
+	// After shifting, the excess bytes are at the end of the output buffer and they are
+	// garbage.
+	e.Truncate(e.Len() - excessReservedBytes)
 	return nil
 }
 
@@ -1527,6 +1547,22 @@ func encodeHead(e *encoderBuffer, t byte, n uint64) {
 	e.Write(e.scratch[:9])
 }
 
+// encodedHeadLen returns the number of bytes that will be written by a call to encodeHead with the
+// given argument. This must be kept in sync with encodeHead.
+func encodedHeadLen(arg uint64) int {
+	switch {
+	case arg <= 23:
+		return 1
+	case arg <= math.MaxUint8:
+		return 2
+	case arg <= math.MaxUint16:
+		return 3
+	case arg <= math.MaxUint32:
+		return 5
+	}
+	return 9
+}
+
 var (
 	typeMarshaler       = reflect.TypeOf((*Marshaler)(nil)).Elem()
 	typeBinaryMarshaler = reflect.TypeOf((*encoding.BinaryMarshaler)(nil)).Elem()


Test patch:
diff --git a/bench_test.go b/bench_test.go
index 47d3c5e..b88249b 100644
--- a/bench_test.go
+++ b/bench_test.go
@@ -107,6 +107,65 @@ type T3 struct {
 	Mss  map[string]string
 }
 
+type ManyFieldsOneOmitEmpty struct {
+	F01, F02, F03, F04, F05, F06, F07, F08, F09, F10, F11, F12, F13, F14, F15, F16 int
+	F17, F18, F19, F20, F21, F22, F23, F24, F25, F26, F27, F28, F29, F30, F31      int
+
+	F32 int `cbor:",omitempty"`
+}
+
+type SomeFieldsOneOmitEmpty struct {
+	F01, F02, F03, F04, F05, F06, F07 int
+
+	F08 int `cbor:",omitempty"`
+}
+
+type ManyFieldsAllOmitEmpty struct {
+	F01 int `cbor:",omitempty"`
+	F02 int `cbor:",omitempty"`
+	F03 int `cbor:",omitempty"`
+	F04 int `cbor:",omitempty"`
+	F05 int `cbor:",omitempty"`
+	F06 int `cbor:",omitempty"`
+	F07 int `cbor:",omitempty"`
+	F08 int `cbor:",omitempty"`
+	F09 int `cbor:",omitempty"`
+	F10 int `cbor:",omitempty"`
+	F11 int `cbor:",omitempty"`
+	F12 int `cbor:",omitempty"`
+	F13 int `cbor:",omitempty"`
+	F14 int `cbor:",omitempty"`
+	F15 int `cbor:",omitempty"`
+	F16 int `cbor:",omitempty"`
+	F17 int `cbor:",omitempty"`
+	F18 int `cbor:",omitempty"`
+	F19 int `cbor:",omitempty"`
+	F20 int `cbor:",omitempty"`
+	F21 int `cbor:",omitempty"`
+	F22 int `cbor:",omitempty"`
+	F23 int `cbor:",omitempty"`
+	F24 int `cbor:",omitempty"`
+	F25 int `cbor:",omitempty"`
+	F26 int `cbor:",omitempty"`
+	F27 int `cbor:",omitempty"`
+	F28 int `cbor:",omitempty"`
+	F29 int `cbor:",omitempty"`
+	F30 int `cbor:",omitempty"`
+	F31 int `cbor:",omitempty"`
+	F32 int `cbor:",omitempty"`
+}
+
+type SomeFieldsAllOmitEmpty struct {
+	F01 int `cbor:",omitempty"`
+	F02 int `cbor:",omitempty"`
+	F03 int `cbor:",omitempty"`
+	F04 int `cbor:",omitempty"`
+	F05 int `cbor:",omitempty"`
+	F06 int `cbor:",omitempty"`
+	F07 int `cbor:",omitempty"`
+	F08 int `cbor:",omitempty"`
+}
+
 var decodeBenchmarks = []struct {
 	name          string
 	data          []byte
@@ -396,6 +455,17 @@ func BenchmarkMarshal(b *testing.B) {
 	}{
 		{"Go map[string]interface{} to CBOR map", m1},
 		{"Go struct to CBOR map", v1},
+		{"Go struct many fields all omitempty all empty to CBOR map", ManyFieldsAllOmitEmpty{}},
+		{"Go struct some fields all omitempty all empty to CBOR map", SomeFieldsAllOmitEmpty{}},
+		{"Go struct many fields all omitempty all nonempty to CBOR map", ManyFieldsAllOmitEmpty{
+			F01: 1, F02: 1, F03: 1, F04: 1, F05: 1, F06: 1, F07: 1, F08: 1, F09: 1, F10: 1, F11: 1, F12: 1, F13: 1, F14: 1, F15: 1, F16: 1,
+			F17: 1, F18: 1, F19: 1, F20: 1, F21: 1, F22: 1, F23: 1, F24: 1, F25: 1, F26: 1, F27: 1, F28: 1, F29: 1, F30: 1, F31: 1, F32: 1,
+		}},
+		{"Go struct some fields all omitempty all nonempty to CBOR map", SomeFieldsAllOmitEmpty{
+			F01: 1, F02: 1, F03: 1, F04: 1, F05: 1, F06: 1, F07: 1, F08: 1,
+		}},
+		{"Go struct many fields one omitempty to CBOR map", ManyFieldsOneOmitEmpty{}},
+		{"Go struct some fields one omitempty to CBOR map", SomeFieldsOneOmitEmpty{}},
 		{"Go map[int]interface{} to CBOR map", m2},
 		{"Go struct keyasint to CBOR map", v2},
 		{"Go []interface{} to CBOR map", slc},
diff --git a/encode_test.go b/encode_test.go
index 2588973..15ca2bf 100644
--- a/encode_test.go
+++ b/encode_test.go
@@ -550,6 +550,58 @@ func TestMarshalStruct(t *testing.T) {
 		t.Errorf("Unmarshal() = %v, want %v", v2, unmarshalWant)
 	}
 }
+
+// TestMarshalStructVariableLength tests marshaling structs that can encode to CBOR maps of varying
+// size depending on their field contents.
+func TestMarshalStructVariableLength(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		in   interface{}
+		want []byte
+	}{
+		{
+			name: "zero out of one items",
+			in: struct {
+				F int `cbor:",omitempty"`
+			}{},
+			want: hexDecode("a0"),
+		},
+		{
+			name: "one out of one items",
+			in: struct {
+				F int `cbor:",omitempty"`
+			}{F: 1},
+			want: hexDecode("a1614601"),
+		},
+		{
+			name: "23 out of 24 items",
+			in: struct {
+				A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P, Q, R, S, T, U, V, W int
+				X                                                                   int `cbor:",omitempty"`
+			}{},
+			want: hexDecode("b7614100614200614300614400614500614600614700614800614900614a00614b00614c00614d00614e00614f00615000615100615200615300615400615500615600615700"),
+		},
+		{
+			name: "24 out of 24 items",
+			in: struct {
+				A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P, Q, R, S, T, U, V, W int
+				X                                                                   int `cbor:",omitempty"`
+			}{X: 1},
+			want: hexDecode("b818614100614200614300614400614500614600614700614800614900614a00614b00614c00614d00614e00614f00615000615100615200615300615400615500615600615700615801"),
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			got, err := Marshal(tc.in)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if !bytes.Equal(tc.want, got) {
+				t.Errorf("want 0x%x but got 0x%x", tc.want, got)
+			}
+		})
+	}
+}
+
 func TestMarshalStructCanonical(t *testing.T) {
 	v := outer{
 		IntField:          123,