internal/impl: refactor MessageInfo

Changes made:
* Move MessageInfo.{methods,extensionFieldInfosMu,extensionFieldInfos}
to the coderMessageInfo type since they are fields all related to the
fast-path implementation, which is what the coderMessageInfo is for.
* Rename message_field.go -> message_reflect_field.go to make it obvious
from the file name that this only deals with the reflection implementation.
* Rename message_test.go -> message_reflect_test.go.
* Move reflection-specific implementation functions from message.go
to message_reflect.go. The intention is to make it such that message.go
is the entry point to message implementations and is agnostic towards
whether we are implementing reflection or the table-driven fast path.

Overall, there is no semantic changes. Just code movement.

Change-Id: I7743c39ba84dc63cd2a02934c3319285e16d6b1c
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/190100
Reviewed-by: Herbie Ong <herbie@google.com>

internal/impl/codec_message.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"reflect"
 	"sort"
+	"sync"
 
 	"google.golang.org/protobuf/internal/encoding/messageset"
 	"google.golang.org/protobuf/internal/encoding/wire"
@@ -19,6 +20,8 @@ import (
 // This is a different type from MessageInfo to keep MessageInfo as general-purpose as
 // possible.
 type coderMessageInfo struct {
+	methods piface.Methods
+
 	orderedCoderFields []*coderFieldInfo
 	denseCoderFields   []*coderFieldInfo
 	coderFields        map[wire.Number]*coderFieldInfo
@@ -26,6 +29,9 @@ type coderMessageInfo struct {
 	unknownOffset      offset
 	extensionOffset    offset
 	needsInitCheck     bool
+
+	extensionFieldInfosMu sync.RWMutex
+	extensionFieldInfos   map[pref.ExtensionType]*extensionFieldInfo
 }
 
 type coderFieldInfo struct {
@@ -38,7 +44,7 @@ type coderFieldInfo struct {
 	isRequired bool              // true if field is required
 }
 
-func (mi *MessageInfo) makeMethods(t reflect.Type, si structInfo) {
+func (mi *MessageInfo) makeCoderMethods(t reflect.Type, si structInfo) {
 	mi.sizecacheOffset = si.sizecacheOffset
 	mi.unknownOffset = si.unknownOffset
 	mi.extensionOffset = si.extensionOffset

internal/impl/message.go

@@ -20,13 +20,14 @@ import (
 // MessageInfo provides protobuf related functionality for a given Go type
 // that represents a message. A given instance of MessageInfo is tied to
 // exactly one Go type, which must be a pointer to a struct type.
+//
+// The exported fields must be populated before any methods are called
+// and cannot be mutated after set.
 type MessageInfo struct {
 	// GoReflectType is the underlying message Go type and must be populated.
-	// Once set, this field must never be mutated.
 	GoReflectType reflect.Type // pointer to struct
 
 	// Desc is the underlying message descriptor type and must be populated.
-	// Once set, this field must never be mutated.
 	Desc pref.MessageDescriptor
 
 	// Exporter must be provided in a purego environment in order to provide
@@ -39,25 +40,8 @@ type MessageInfo struct {
 	initMu   sync.Mutex // protects all unexported fields
 	initDone uint32
 
-	reflectMessageInfo
-
-	// Information used by the fast-path methods.
-	methods piface.Methods
-	coderMessageInfo
-
-	extensionFieldInfosMu sync.RWMutex
-	extensionFieldInfos   map[pref.ExtensionType]*extensionFieldInfo
-}
-
-type reflectMessageInfo struct {
-	fields map[pref.FieldNumber]*fieldInfo
-	oneofs map[pref.Name]*oneofInfo
-
-	getUnknown   func(pointer) pref.RawFields
-	setUnknown   func(pointer, pref.RawFields)
-	extensionMap func(pointer) *extensionMap
-
-	nilMessage atomicNilMessage
+	reflectMessageInfo // for reflection implementation
+	coderMessageInfo   // for fast-path method implementations
 }
 
 // exporter is a function that returns a reference to the ith field of v,
@@ -81,8 +65,8 @@ func getMessageInfo(mt reflect.Type) *MessageInfo {
 }
 
 func (mi *MessageInfo) init() {
-	// This function is called in the hot path. Inline the sync.Once
-	// logic, since allocating a closure for Once.Do is expensive.
+	// This function is called in the hot path. Inline the sync.Once logic,
+	// since allocating a closure for Once.Do is expensive.
 	// Keep init small to ensure that it can be inlined.
 	if atomic.LoadUint32(&mi.initDone) == 0 {
 		mi.initOnce()
@@ -100,12 +84,11 @@ func (mi *MessageInfo) initOnce() {
 	if t.Kind() != reflect.Ptr && t.Elem().Kind() != reflect.Struct {
 		panic(fmt.Sprintf("got %v, want *struct kind", t))
 	}
+	t = t.Elem()
 
-	si := mi.makeStructInfo(t.Elem())
-	mi.makeKnownFieldsFunc(si)
-	mi.makeUnknownFieldsFunc(t.Elem(), si)
-	mi.makeExtensionFieldsFunc(t.Elem(), si)
-	mi.makeMethods(t.Elem(), si)
+	si := mi.makeStructInfo(t)
+	mi.makeReflectFuncs(t, si)
+	mi.makeCoderMethods(t, si)
 
 	atomic.StoreUint32(&mi.initDone, 1)
 }
@@ -215,89 +198,6 @@ fieldLoop:
 	return si
 }
 
-// makeKnownFieldsFunc generates functions for operations that can be performed
-// on each protobuf message field. It takes in a reflect.Type representing the
-// Go struct and matches message fields with struct fields.
-//
-// This code assumes that the struct is well-formed and panics if there are
-// any discrepancies.
-func (mi *MessageInfo) makeKnownFieldsFunc(si structInfo) {
-	mi.fields = map[pref.FieldNumber]*fieldInfo{}
-	md := mi.Desc
-	for i := 0; i < md.Fields().Len(); i++ {
-		fd := md.Fields().Get(i)
-		fs := si.fieldsByNumber[fd.Number()]
-		var fi fieldInfo
-		switch {
-		case fd.ContainingOneof() != nil:
-			fi = fieldInfoForOneof(fd, si.oneofsByName[fd.ContainingOneof().Name()], mi.Exporter, si.oneofWrappersByNumber[fd.Number()])
-		case fd.IsMap():
-			fi = fieldInfoForMap(fd, fs, mi.Exporter)
-		case fd.IsList():
-			fi = fieldInfoForList(fd, fs, mi.Exporter)
-		case fd.IsWeak():
-			fi = fieldInfoForWeakMessage(fd, si.weakOffset)
-		case fd.Kind() == pref.MessageKind || fd.Kind() == pref.GroupKind:
-			fi = fieldInfoForMessage(fd, fs, mi.Exporter)
-		default:
-			fi = fieldInfoForScalar(fd, fs, mi.Exporter)
-		}
-		mi.fields[fd.Number()] = &fi
-	}
-
-	mi.oneofs = map[pref.Name]*oneofInfo{}
-	for i := 0; i < md.Oneofs().Len(); i++ {
-		od := md.Oneofs().Get(i)
-		mi.oneofs[od.Name()] = makeOneofInfo(od, si.oneofsByName[od.Name()], mi.Exporter, si.oneofWrappersByType)
-	}
-}
-
-func (mi *MessageInfo) makeUnknownFieldsFunc(t reflect.Type, si structInfo) {
-	mi.getUnknown = func(pointer) pref.RawFields { return nil }
-	mi.setUnknown = func(pointer, pref.RawFields) { return }
-	if si.unknownOffset.IsValid() {
-		mi.getUnknown = func(p pointer) pref.RawFields {
-			if p.IsNil() {
-				return nil
-			}
-			rv := p.Apply(si.unknownOffset).AsValueOf(unknownFieldsType)
-			return pref.RawFields(*rv.Interface().(*[]byte))
-		}
-		mi.setUnknown = func(p pointer, b pref.RawFields) {
-			if p.IsNil() {
-				panic("invalid SetUnknown on nil Message")
-			}
-			rv := p.Apply(si.unknownOffset).AsValueOf(unknownFieldsType)
-			*rv.Interface().(*[]byte) = []byte(b)
-		}
-	} else {
-		mi.getUnknown = func(pointer) pref.RawFields {
-			return nil
-		}
-		mi.setUnknown = func(p pointer, _ pref.RawFields) {
-			if p.IsNil() {
-				panic("invalid SetUnknown on nil Message")
-			}
-		}
-	}
-}
-
-func (mi *MessageInfo) makeExtensionFieldsFunc(t reflect.Type, si structInfo) {
-	if si.extensionOffset.IsValid() {
-		mi.extensionMap = func(p pointer) *extensionMap {
-			if p.IsNil() {
-				return (*extensionMap)(nil)
-			}
-			v := p.Apply(si.extensionOffset).AsValueOf(extensionFieldsType)
-			return (*extensionMap)(v.Interface().(*map[int32]ExtensionField))
-		}
-	} else {
-		mi.extensionMap = func(pointer) *extensionMap {
-			return (*extensionMap)(nil)
-		}
-	}
-}
-
 func (mi *MessageInfo) GoType() reflect.Type {
 	return mi.GoReflectType
 }

internal/impl/message_reflect.go

@@ -12,6 +12,159 @@ import (
 	pref "google.golang.org/protobuf/reflect/protoreflect"
 )
 
+type reflectMessageInfo struct {
+	fields map[pref.FieldNumber]*fieldInfo
+	oneofs map[pref.Name]*oneofInfo
+
+	getUnknown   func(pointer) pref.RawFields
+	setUnknown   func(pointer, pref.RawFields)
+	extensionMap func(pointer) *extensionMap
+
+	nilMessage atomicNilMessage
+}
+
+// makeReflectFuncs generates the set of functions to support reflection.
+func (mi *MessageInfo) makeReflectFuncs(t reflect.Type, si structInfo) {
+	mi.makeKnownFieldsFunc(si)
+	mi.makeUnknownFieldsFunc(t, si)
+	mi.makeExtensionFieldsFunc(t, si)
+}
+
+// makeKnownFieldsFunc generates functions for operations that can be performed
+// on each protobuf message field. It takes in a reflect.Type representing the
+// Go struct and matches message fields with struct fields.
+//
+// This code assumes that the struct is well-formed and panics if there are
+// any discrepancies.
+func (mi *MessageInfo) makeKnownFieldsFunc(si structInfo) {
+	mi.fields = map[pref.FieldNumber]*fieldInfo{}
+	md := mi.Desc
+	for i := 0; i < md.Fields().Len(); i++ {
+		fd := md.Fields().Get(i)
+		fs := si.fieldsByNumber[fd.Number()]
+		var fi fieldInfo
+		switch {
+		case fd.ContainingOneof() != nil:
+			fi = fieldInfoForOneof(fd, si.oneofsByName[fd.ContainingOneof().Name()], mi.Exporter, si.oneofWrappersByNumber[fd.Number()])
+		case fd.IsMap():
+			fi = fieldInfoForMap(fd, fs, mi.Exporter)
+		case fd.IsList():
+			fi = fieldInfoForList(fd, fs, mi.Exporter)
+		case fd.IsWeak():
+			fi = fieldInfoForWeakMessage(fd, si.weakOffset)
+		case fd.Kind() == pref.MessageKind || fd.Kind() == pref.GroupKind:
+			fi = fieldInfoForMessage(fd, fs, mi.Exporter)
+		default:
+			fi = fieldInfoForScalar(fd, fs, mi.Exporter)
+		}
+		mi.fields[fd.Number()] = &fi
+	}
+
+	mi.oneofs = map[pref.Name]*oneofInfo{}
+	for i := 0; i < md.Oneofs().Len(); i++ {
+		od := md.Oneofs().Get(i)
+		mi.oneofs[od.Name()] = makeOneofInfo(od, si.oneofsByName[od.Name()], mi.Exporter, si.oneofWrappersByType)
+	}
+}
+
+func (mi *MessageInfo) makeUnknownFieldsFunc(t reflect.Type, si structInfo) {
+	mi.getUnknown = func(pointer) pref.RawFields { return nil }
+	mi.setUnknown = func(pointer, pref.RawFields) { return }
+	if si.unknownOffset.IsValid() {
+		mi.getUnknown = func(p pointer) pref.RawFields {
+			if p.IsNil() {
+				return nil
+			}
+			rv := p.Apply(si.unknownOffset).AsValueOf(unknownFieldsType)
+			return pref.RawFields(*rv.Interface().(*[]byte))
+		}
+		mi.setUnknown = func(p pointer, b pref.RawFields) {
+			if p.IsNil() {
+				panic("invalid SetUnknown on nil Message")
+			}
+			rv := p.Apply(si.unknownOffset).AsValueOf(unknownFieldsType)
+			*rv.Interface().(*[]byte) = []byte(b)
+		}
+	} else {
+		mi.getUnknown = func(pointer) pref.RawFields {
+			return nil
+		}
+		mi.setUnknown = func(p pointer, _ pref.RawFields) {
+			if p.IsNil() {
+				panic("invalid SetUnknown on nil Message")
+			}
+		}
+	}
+}
+
+func (mi *MessageInfo) makeExtensionFieldsFunc(t reflect.Type, si structInfo) {
+	if si.extensionOffset.IsValid() {
+		mi.extensionMap = func(p pointer) *extensionMap {
+			if p.IsNil() {
+				return (*extensionMap)(nil)
+			}
+			v := p.Apply(si.extensionOffset).AsValueOf(extensionFieldsType)
+			return (*extensionMap)(v.Interface().(*map[int32]ExtensionField))
+		}
+	} else {
+		mi.extensionMap = func(pointer) *extensionMap {
+			return (*extensionMap)(nil)
+		}
+	}
+}
+
+type extensionMap map[int32]ExtensionField
+
+func (m *extensionMap) Range(f func(pref.FieldDescriptor, pref.Value) bool) {
+	if m != nil {
+		for _, x := range *m {
+			xt := x.GetType()
+			if !f(xt.Descriptor(), xt.ValueOf(x.GetValue())) {
+				return
+			}
+		}
+	}
+}
+func (m *extensionMap) Has(xt pref.ExtensionType) (ok bool) {
+	if m != nil {
+		_, ok = (*m)[int32(xt.Descriptor().Number())]
+	}
+	return ok
+}
+func (m *extensionMap) Clear(xt pref.ExtensionType) {
+	delete(*m, int32(xt.Descriptor().Number()))
+}
+func (m *extensionMap) Get(xt pref.ExtensionType) pref.Value {
+	xd := xt.Descriptor()
+	if m != nil {
+		if x, ok := (*m)[int32(xd.Number())]; ok {
+			return xt.ValueOf(x.GetValue())
+		}
+	}
+	return xt.Zero()
+}
+func (m *extensionMap) Set(xt pref.ExtensionType, v pref.Value) {
+	if *m == nil {
+		*m = make(map[int32]ExtensionField)
+	}
+	var x ExtensionField
+	x.SetType(xt)
+	x.SetEagerValue(xt.InterfaceOf(v))
+	(*m)[int32(xt.Descriptor().Number())] = x
+}
+func (m *extensionMap) Mutable(xt pref.ExtensionType) pref.Value {
+	xd := xt.Descriptor()
+	if xd.Kind() != pref.MessageKind && xd.Kind() != pref.GroupKind && !xd.IsList() && !xd.IsMap() {
+		panic("invalid Mutable on field with non-composite type")
+	}
+	if x, ok := (*m)[int32(xd.Number())]; ok {
+		return xt.ValueOf(x.GetValue())
+	}
+	v := xt.New()
+	m.Set(xt, v)
+	return v
+}
+
 // MessageState is a data structure that is nested as the first field in a
 // concrete message. It provides a way to implement the ProtoReflect method
 // in an allocation-free way without needing to have a shadow Go type generated
@@ -115,62 +268,6 @@ func (m *messageIfaceWrapper) ProtoUnwrap() interface{} {
 	return m.p.AsIfaceOf(m.mi.GoReflectType.Elem())
 }
 
-type extensionMap map[int32]ExtensionField
-
-func (m *extensionMap) Range(f func(pref.FieldDescriptor, pref.Value) bool) {
-	if m != nil {
-		for _, x := range *m {
-			xt := x.GetType()
-			if !f(xt.Descriptor(), xt.ValueOf(x.GetValue())) {
-				return
-			}
-		}
-	}
-}
-func (m *extensionMap) Has(xt pref.ExtensionType) (ok bool) {
-	if m != nil {
-		_, ok = (*m)[int32(xt.Descriptor().Number())]
-	}
-	return ok
-}
-func (m *extensionMap) Clear(xt pref.ExtensionType) {
-	delete(*m, int32(xt.Descriptor().Number()))
-}
-func (m *extensionMap) Get(xt pref.ExtensionType) pref.Value {
-	xd := xt.Descriptor()
-	if m != nil {
-		if x, ok := (*m)[int32(xd.Number())]; ok {
-			return xt.ValueOf(x.GetValue())
-		}
-	}
-	return xt.Zero()
-}
-func (m *extensionMap) Set(xt pref.ExtensionType, v pref.Value) {
-	if *m == nil {
-		*m = make(map[int32]ExtensionField)
-	}
-	var x ExtensionField
-	x.SetType(xt)
-	x.SetEagerValue(xt.InterfaceOf(v))
-	(*m)[int32(xt.Descriptor().Number())] = x
-}
-func (m *extensionMap) Mutable(xt pref.ExtensionType) pref.Value {
-	xd := xt.Descriptor()
-	if !isComposite(xd) {
-		panic("invalid Mutable on field with non-composite type")
-	}
-	if x, ok := (*m)[int32(xd.Number())]; ok {
-		return xt.ValueOf(x.GetValue())
-	}
-	v := xt.New()
-	m.Set(xt, v)
-	return v
-}
-
-func isComposite(fd pref.FieldDescriptor) bool {
-	return fd.Kind() == pref.MessageKind || fd.Kind() == pref.GroupKind || fd.IsList() || fd.IsMap()
-}
-
 // checkField verifies that the provided field descriptor is valid.
 // Exactly one of the returned values is populated.
 func (mi *MessageInfo) checkField(fd pref.FieldDescriptor) (*fieldInfo, pref.ExtensionType) {