internal/impl: refactor fast-path

Move data used by the fast-path implementations into a substructure of
MessageInfo and initialize it separately.

Change-Id: Ib855ee8ea5cb0379528b52ba0e191319aa5e2dff
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/184077
Reviewed-by: Joe Tsai <thebrokentoaster@gmail.com>

internal/impl/codec_field.go

@@ -29,6 +29,8 @@ type ifaceCoderFuncs struct {
 // struct fields.
 func fieldCoder(fd pref.FieldDescriptor, ft reflect.Type) pointerCoderFuncs {
 	switch {
+	case fd.IsMap():
+		return encoderFuncsForMap(fd, ft)
 	case fd.Cardinality() == pref.Repeated && !fd.IsPacked():
 		// Repeated fields (not packed).
 		if ft.Kind() != reflect.Slice {

internal/impl/codec_message.go

@@ -0,0 +1,101 @@
+// Copyright 2019 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package impl
+
+import (
+	"reflect"
+	"sort"
+
+	"google.golang.org/protobuf/internal/encoding/wire"
+	pref "google.golang.org/protobuf/reflect/protoreflect"
+	piface "google.golang.org/protobuf/runtime/protoiface"
+)
+
+// coderMessageInfo contains per-message information used by the fast-path functions.
+// This is a different type from MessageInfo to keep MessageInfo as general-purpose as
+// possible.
+type coderMessageInfo struct {
+	orderedCoderFields []*coderFieldInfo
+	sizecacheOffset    offset
+	extensionOffset    offset
+	unknownOffset      offset
+	needsInitCheck     bool
+}
+
+type coderFieldInfo struct {
+	funcs      pointerCoderFuncs // fast-path per-field functions
+	num        pref.FieldNumber  // field number
+	offset     offset            // struct field offset
+	wiretag    uint64            // field tag (number + wire type)
+	tagsize    int               // size of the varint-encoded tag
+	isPointer  bool              // true if IsNil may be called on the struct field
+	isRequired bool              // true if field is required
+}
+
+func (mi *MessageInfo) makeMethods(t reflect.Type, si structInfo) {
+	mi.sizecacheOffset = invalidOffset
+	if fx, _ := t.FieldByName("XXX_sizecache"); fx.Type == sizecacheType {
+		mi.sizecacheOffset = offsetOf(fx)
+	}
+	mi.unknownOffset = invalidOffset
+	if fx, _ := t.FieldByName("XXX_unrecognized"); fx.Type == unknownFieldsType {
+		mi.unknownOffset = offsetOf(fx)
+	}
+	mi.extensionOffset = invalidOffset
+	if fx, _ := t.FieldByName("XXX_InternalExtensions"); fx.Type == extensionFieldsType {
+		mi.extensionOffset = offsetOf(fx)
+	} else if fx, _ = t.FieldByName("XXX_extensions"); fx.Type == extensionFieldsType {
+		mi.extensionOffset = offsetOf(fx)
+	}
+
+	for i := 0; i < mi.PBType.Descriptor().Fields().Len(); i++ {
+		fd := mi.PBType.Descriptor().Fields().Get(i)
+		if fd.ContainingOneof() != nil {
+			continue
+		}
+
+		fs := si.fieldsByNumber[fd.Number()]
+		ft := fs.Type
+		var wiretag uint64
+		if !fd.IsPacked() {
+			wiretag = wire.EncodeTag(fd.Number(), wireTypes[fd.Kind()])
+		} else {
+			wiretag = wire.EncodeTag(fd.Number(), wire.BytesType)
+		}
+		mi.orderedCoderFields = append(mi.orderedCoderFields, &coderFieldInfo{
+			num:     fd.Number(),
+			offset:  offsetOf(fs),
+			wiretag: wiretag,
+			tagsize: wire.SizeVarint(wiretag),
+			funcs:   fieldCoder(fd, ft),
+			isPointer: (fd.Cardinality() == pref.Repeated ||
+				fd.Kind() == pref.MessageKind ||
+				fd.Kind() == pref.GroupKind ||
+				fd.Syntax() != pref.Proto3),
+			isRequired: fd.Cardinality() == pref.Required,
+		})
+	}
+	for i := 0; i < mi.PBType.Descriptor().Oneofs().Len(); i++ {
+		od := mi.PBType.Descriptor().Oneofs().Get(i)
+		fs := si.oneofsByName[od.Name()]
+		mi.orderedCoderFields = append(mi.orderedCoderFields, &coderFieldInfo{
+			num:       od.Fields().Get(0).Number(),
+			offset:    offsetOf(fs),
+			funcs:     makeOneofFieldCoder(fs, od, si.fieldsByNumber, si.oneofWrappersByNumber),
+			isPointer: true,
+		})
+	}
+	sort.Slice(mi.orderedCoderFields, func(i, j int) bool {
+		return mi.orderedCoderFields[i].num < mi.orderedCoderFields[j].num
+	})
+
+	mi.needsInitCheck = needsInitCheck(mi.PBType)
+	mi.methods = piface.Methods{
+		Flags:         piface.MethodFlagDeterministicMarshal,
+		MarshalAppend: mi.marshalAppend,
+		Size:          mi.size,
+		IsInitialized: mi.isInitialized,
+	}
+}

internal/impl/encode.go

@@ -68,7 +68,7 @@ func (mi *MessageInfo) sizePointerSlow(p pointer, opts marshalOptions) (size int
 		e := p.Apply(mi.extensionOffset).Extensions()
 		size += mi.sizeExtensions(e, opts)
 	}
-	for _, f := range mi.fieldsOrdered {
+	for _, f := range mi.orderedCoderFields {
 		fptr := p.Apply(f.offset)
 		if f.isPointer && fptr.Elem().IsNil() {
 			continue
@@ -108,7 +108,7 @@ func (mi *MessageInfo) marshalAppendPointer(b []byte, p pointer, opts marshalOpt
 			return b, err
 		}
 	}
-	for _, f := range mi.fieldsOrdered {
+	for _, f := range mi.orderedCoderFields {
 		fptr := p.Apply(f.offset)
 		if f.isPointer && fptr.Elem().IsNil() {
 			continue

internal/impl/isinit.go

@@ -34,7 +34,7 @@ func (mi *MessageInfo) isInitializedPointer(p pointer) error {
 			return err
 		}
 	}
-	for _, f := range mi.fieldsOrdered {
+	for _, f := range mi.orderedCoderFields {
 		if !f.isRequired && f.funcs.isInit == nil {
 			continue
 		}

internal/impl/message.go

@@ -7,7 +7,6 @@ package impl
 import (
 	"fmt"
 	"reflect"
-	"sort"
 	"strconv"
 	"strings"
 	"sync"
@@ -33,11 +32,7 @@ type MessageInfo struct {
 	initMu   sync.Mutex // protects all unexported fields
 	initDone uint32
 
-	// Keep a separate slice of fields for efficient field encoding in tag order
-	// and because iterating over a slice is substantially faster than a map.
-	fields        map[pref.FieldNumber]*fieldInfo
-	fieldsOrdered []*fieldInfo
-
+	fields map[pref.FieldNumber]*fieldInfo
 	oneofs map[pref.Name]*oneofInfo
 
 	getUnknown func(pointer) pref.RawFields
@@ -45,12 +40,10 @@ type MessageInfo struct {
 
 	extensionMap func(pointer) *extensionMap
 
+	// Information used by the fast-path methods.
 	methods piface.Methods
+	coderMessageInfo
 
-	needsInitCheck        bool
-	sizecacheOffset       offset
-	extensionOffset       offset
-	unknownOffset         offset
 	extensionFieldInfosMu sync.RWMutex
 	extensionFieldInfos   map[pref.ExtensionType]*extensionFieldInfo
 }
@@ -102,11 +95,10 @@ func (mi *MessageInfo) initOnce() {
 	}
 
 	si := mi.makeStructInfo(t.Elem())
-	mi.needsInitCheck = needsInitCheck(mi.PBType)
 	mi.makeKnownFieldsFunc(si)
 	mi.makeUnknownFieldsFunc(t.Elem())
 	mi.makeExtensionFieldsFunc(t.Elem())
-	mi.makeMethods(t.Elem())
+	mi.makeMethods(t.Elem(), si)
 
 	atomic.StoreUint32(&mi.initDone, 1)
 }
@@ -123,27 +115,6 @@ var (
 	extensionFieldsType = reflect.TypeOf(ExtensionFields(nil))
 )
 
-func (mi *MessageInfo) makeMethods(t reflect.Type) {
-	mi.sizecacheOffset = invalidOffset
-	if fx, _ := t.FieldByName("XXX_sizecache"); fx.Type == sizecacheType {
-		mi.sizecacheOffset = offsetOf(fx)
-	}
-	mi.unknownOffset = invalidOffset
-	if fx, _ := t.FieldByName("XXX_unrecognized"); fx.Type == unknownFieldsType {
-		mi.unknownOffset = offsetOf(fx)
-	}
-	mi.extensionOffset = invalidOffset
-	if fx, _ := t.FieldByName("XXX_InternalExtensions"); fx.Type == extensionFieldsType {
-		mi.extensionOffset = offsetOf(fx)
-	} else if fx, _ = t.FieldByName("XXX_extensions"); fx.Type == extensionFieldsType {
-		mi.extensionOffset = offsetOf(fx)
-	}
-	mi.methods.Flags = piface.MethodFlagDeterministicMarshal
-	mi.methods.MarshalAppend = mi.marshalAppend
-	mi.methods.Size = mi.size
-	mi.methods.IsInitialized = mi.isInitialized
-}
-
 type structInfo struct {
 	fieldsByNumber        map[pref.FieldNumber]reflect.StructField
 	oneofsByName          map[pref.Name]reflect.StructField
@@ -204,7 +175,6 @@ fieldLoop:
 // any discrepancies.
 func (mi *MessageInfo) makeKnownFieldsFunc(si structInfo) {
 	mi.fields = map[pref.FieldNumber]*fieldInfo{}
-	mi.fieldsOrdered = make([]*fieldInfo, 0, mi.PBType.Fields().Len())
 	for i := 0; i < mi.PBType.Descriptor().Fields().Len(); i++ {
 		fd := mi.PBType.Descriptor().Fields().Get(i)
 		fs := si.fieldsByNumber[fd.Number()]
@@ -212,16 +182,6 @@ func (mi *MessageInfo) makeKnownFieldsFunc(si structInfo) {
 		switch {
 		case fd.ContainingOneof() != nil:
 			fi = fieldInfoForOneof(fd, si.oneofsByName[fd.ContainingOneof().Name()], si.oneofWrappersByNumber[fd.Number()])
-			// There is one fieldInfo for each proto message field, but only one struct
-			// field for all message fields in a oneof. We install the encoder functions
-			// on the fieldInfo for the first field in the oneof.
-			//
-			// A slightly simpler approach would be to have each fieldInfo's encoder
-			// handle the case where that field is set, but this would require more
-			// checks  against the current oneof type than a single map lookup.
-			if fd.ContainingOneof().Fields().Get(0).Name() == fd.Name() {
-				fi.funcs = makeOneofFieldCoder(si.oneofsByName[fd.ContainingOneof().Name()], fd.ContainingOneof(), si.fieldsByNumber, si.oneofWrappersByNumber)
-			}
 		case fd.IsMap():
 			fi = fieldInfoForMap(fd, fs)
 		case fd.IsList():
@@ -231,13 +191,8 @@ func (mi *MessageInfo) makeKnownFieldsFunc(si structInfo) {
 		default:
 			fi = fieldInfoForScalar(fd, fs)
 		}
-		fi.num = fd.Number()
 		mi.fields[fd.Number()] = &fi
-		mi.fieldsOrdered = append(mi.fieldsOrdered, &fi)
 	}
-	sort.Slice(mi.fieldsOrdered, func(i, j int) bool {
-		return mi.fieldsOrdered[i].num < mi.fieldsOrdered[j].num
-	})
 
 	mi.oneofs = map[pref.Name]*oneofInfo{}
 	for i := 0; i < mi.PBType.Descriptor().Oneofs().Len(); i++ {

internal/impl/message_field.go

@@ -9,7 +9,6 @@ import (
 	"math"
 	"reflect"
 
-	"google.golang.org/protobuf/internal/encoding/wire"
 	pvalue "google.golang.org/protobuf/internal/value"
 	pref "google.golang.org/protobuf/reflect/protoreflect"
 	piface "google.golang.org/protobuf/runtime/protoiface"
@@ -25,15 +24,6 @@ type fieldInfo struct {
 	set        func(pointer, pref.Value)
 	mutable    func(pointer) pref.Value
 	newMessage func() pref.Message
-
-	// These fields are used for fast-path functions.
-	funcs      pointerCoderFuncs // fast-path per-field functions
-	num        pref.FieldNumber  // field number
-	offset     offset            // struct field offset
-	wiretag    uint64            // field tag (number + wire type)
-	tagsize    int               // size of the varint-encoded tag
-	isPointer  bool              // true if IsNil may be called on the struct field
-	isRequired bool              // true if field is required
 }
 
 func fieldInfoForOneof(fd pref.FieldDescriptor, fs reflect.StructField, ot reflect.Type) fieldInfo {
@@ -118,8 +108,6 @@ func fieldInfoForOneof(fd pref.FieldDescriptor, fs reflect.StructField, ot refle
 			return conv.PBValueOf(rv)
 		},
 		newMessage: conv.NewMessage,
-		offset:     fieldOffset,
-		isPointer:  true,
 	}
 }
 
@@ -130,7 +118,6 @@ func fieldInfoForMap(fd pref.FieldDescriptor, fs reflect.StructField) fieldInfo
 	}
 	keyConv, _ := newConverter(ft.Key(), fd.MapKey().Kind())
 	valConv, _ := newConverter(ft.Elem(), fd.MapValue().Kind())
-	wiretag := wire.EncodeTag(fd.Number(), wireTypes[fd.Kind()])
 	frozenEmpty := pref.ValueOf(frozenMap{
 		pvalue.MapOf(reflect.Zero(reflect.PtrTo(fs.Type)).Interface(), keyConv, valConv),
 	})
@@ -168,11 +155,6 @@ func fieldInfoForMap(fd pref.FieldDescriptor, fs reflect.StructField) fieldInfo
 			v := p.Apply(fieldOffset).AsIfaceOf(fs.Type)
 			return pref.ValueOf(pvalue.MapOf(v, keyConv, valConv))
 		},
-		funcs:     encoderFuncsForMap(fd, ft),
-		offset:    fieldOffset,
-		wiretag:   wiretag,
-		tagsize:   wire.SizeVarint(wiretag),
-		isPointer: true,
 	}
 }
 
@@ -182,12 +164,6 @@ func fieldInfoForList(fd pref.FieldDescriptor, fs reflect.StructField) fieldInfo
 		panic(fmt.Sprintf("invalid type: got %v, want slice kind", ft))
 	}
 	conv, _ := newConverter(ft.Elem(), fd.Kind())
-	var wiretag uint64
-	if !fd.IsPacked() {
-		wiretag = wire.EncodeTag(fd.Number(), wireTypes[fd.Kind()])
-	} else {
-		wiretag = wire.EncodeTag(fd.Number(), wire.BytesType)
-	}
 	frozenEmpty := pref.ValueOf(frozenList{
 		pvalue.ListOf(reflect.Zero(reflect.PtrTo(fs.Type)).Interface(), conv),
 	})
@@ -225,11 +201,6 @@ func fieldInfoForList(fd pref.FieldDescriptor, fs reflect.StructField) fieldInfo
 			v := p.Apply(fieldOffset).AsIfaceOf(fs.Type)
 			return pref.ValueOf(pvalue.ListOf(v, conv))
 		},
-		funcs:     fieldCoder(fd, ft),
-		offset:    fieldOffset,
-		wiretag:   wiretag,
-		tagsize:   wire.SizeVarint(wiretag),
-		isPointer: true,
 	}
 }
 
@@ -237,7 +208,6 @@ var emptyBytes = reflect.ValueOf([]byte{})
 
 func fieldInfoForScalar(fd pref.FieldDescriptor, fs reflect.StructField) fieldInfo {
 	ft := fs.Type
-	funcs := fieldCoder(fd, ft)
 	nullable := fd.Syntax() == pref.Proto2
 	if nullable {
 		if ft.Kind() != reflect.Ptr && ft.Kind() != reflect.Slice {
@@ -248,7 +218,6 @@ func fieldInfoForScalar(fd pref.FieldDescriptor, fs reflect.StructField) fieldIn
 		}
 	}
 	conv, _ := newConverter(ft, fd.Kind())
-	wiretag := wire.EncodeTag(fd.Number(), wireTypes[fd.Kind()])
 
 	// TODO: Implement unsafe fast path?
 	fieldOffset := offsetOf(fs)
@@ -309,19 +278,12 @@ func fieldInfoForScalar(fd pref.FieldDescriptor, fs reflect.StructField) fieldIn
 				rv.Set(emptyBytes)
 			}
 		},
-		funcs:      funcs,
-		offset:     fieldOffset,
-		isPointer:  nullable,
-		isRequired: fd.Cardinality() == pref.Required,
-		wiretag:    wiretag,
-		tagsize:    wire.SizeVarint(wiretag),
 	}
 }
 
 func fieldInfoForMessage(fd pref.FieldDescriptor, fs reflect.StructField) fieldInfo {
 	ft := fs.Type
 	conv, _ := newConverter(ft, fd.Kind())
-	wiretag := wire.EncodeTag(fd.Number(), wireTypes[fd.Kind()])
 	frozenEmpty := pref.ValueOf(frozenMessage{conv.NewMessage()})
 
 	// TODO: Implement unsafe fast path?
@@ -364,12 +326,6 @@ func fieldInfoForMessage(fd pref.FieldDescriptor, fs reflect.StructField) fieldI
 			return conv.PBValueOf(rv)
 		},
 		newMessage: conv.NewMessage,
-		funcs:      fieldCoder(fd, ft),
-		offset:     fieldOffset,
-		isPointer:  true,
-		isRequired: fd.Cardinality() == pref.Required,
-		wiretag:    wiretag,
-		tagsize:    wire.SizeVarint(wiretag),
 	}
 }