Encode structs directly to output buffer.

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},

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
 }

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()

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,