Encode structs directly to output buffer.

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 Signed-off-by: Ben Luddy <[email protected]>
fxamacker · Apr 29, 2024 · d981dec · d981dec
1 parent da82dfa
commit d981dec
Show file tree

Hide file tree

Showing 4 changed files with 181 additions and 14 deletions.
diff --git 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/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
@@ -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()

diff --git 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,