reflect/protoreflect: add alternative message reflection API

Added API:
	Message.Len
	Message.Range
	Message.Has
	Message.Clear
	Message.Get
	Message.Set
	Message.Mutable
	Message.NewMessage
	Message.WhichOneof
	Message.GetUnknown
	Message.SetUnknown

Deprecated API (to be removed in subsequent CL):
	Message.KnownFields
	Message.UnknownFields

The primary difference with the new API is that the top-level
Message methods are keyed by FieldDescriptor rather than FieldNumber
with the following semantics:
* For known fields, the FieldDescriptor must exactly match the
field descriptor known by the message.
* For extension fields, the FieldDescriptor must implement ExtensionType,
where ContainingMessage.FullName matches the message name, and
the field number is within the message's extension range.
When setting an extension field, it automatically stores
the extension type information.
* Extension fields are always considered nullable,
implying that repeated extension fields are nullable.
That is, you can distinguish between a unpopulated list and an empty list.
* Message.Get always returns a valid Value even if unpopulated.
The behavior is already well-defined for scalars, but for unpopulated
composite types, it now returns an empty read-only version of it.

Change-Id: Ia120630b4db221aeaaf743d0f64160e1a61a0f61
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/175458
Reviewed-by: Damien Neil <dneil@google.com>
diff --git a/internal/impl/message.go b/internal/impl/message.go
index 6db60f1..c416440 100644
--- a/internal/impl/message.go
+++ b/internal/impl/message.go
@@ -40,12 +40,17 @@
 
 	oneofs map[pref.Name]*oneofInfo
 
+	getUnknown func(pointer) pref.RawFields
+	setUnknown func(pointer, pref.RawFields)
+
+	extensionMap func(pointer) *extensionMap
+
 	unknownFields   func(*messageDataType) pref.UnknownFields
 	extensionFields func(*messageDataType) pref.KnownFields
 	methods         piface.Methods
 
-	extensionOffset       offset
 	sizecacheOffset       offset
+	extensionOffset       offset
 	unknownOffset         offset
 	extensionFieldInfosMu sync.RWMutex
 	extensionFieldInfos   map[pref.ExtensionType]*extensionFieldInfo
@@ -106,23 +111,33 @@
 	atomic.StoreUint32(&mi.initDone, 1)
 }
 
-var sizecacheType = reflect.TypeOf(int32(0))
+type (
+	SizeCache       = int32
+	UnknownFields   = []byte
+	ExtensionFields = map[int32]ExtensionField
+)
+
+var (
+	sizecacheType       = reflect.TypeOf(SizeCache(0))
+	unknownFieldsType   = reflect.TypeOf(UnknownFields(nil))
+	extensionFieldsType = reflect.TypeOf(ExtensionFields(nil))
+)
 
 func (mi *MessageInfo) makeMethods(t reflect.Type) {
-	mi.extensionOffset = invalidOffset
-	if fx, _ := t.FieldByName("XXX_InternalExtensions"); fx.Type == extType {
-		mi.extensionOffset = offsetOf(fx)
-	} else if fx, _ = t.FieldByName("XXX_extensions"); fx.Type == extType {
-		mi.extensionOffset = offsetOf(fx)
-	}
 	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 == bytesType {
+	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
@@ -231,22 +246,56 @@
 }
 
 func (mi *MessageInfo) makeUnknownFieldsFunc(t reflect.Type) {
-	if f := makeLegacyUnknownFieldsFunc(t); f != nil {
-		mi.unknownFields = f
-		return
-	}
-	mi.unknownFields = func(*messageDataType) pref.UnknownFields {
-		return emptyUnknownFields{}
+	mi.unknownFields = makeLegacyUnknownFieldsFunc(t)
+
+	mi.getUnknown = func(pointer) pref.RawFields { return nil }
+	mi.setUnknown = func(pointer, pref.RawFields) { return }
+	fu, _ := t.FieldByName("XXX_unrecognized")
+	if fu.Type == unknownFieldsType {
+		fieldOffset := offsetOf(fu)
+		mi.getUnknown = func(p pointer) pref.RawFields {
+			if p.IsNil() {
+				return nil
+			}
+			rv := p.Apply(fieldOffset).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(fieldOffset).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) {
-	if f := makeLegacyExtensionFieldsFunc(t); f != nil {
-		mi.extensionFields = f
-		return
+	mi.extensionFields = makeLegacyExtensionFieldsFunc(t)
+
+	fx, _ := t.FieldByName("XXX_extensions")
+	if fx.Type != extensionFieldsType {
+		fx, _ = t.FieldByName("XXX_InternalExtensions")
 	}
-	mi.extensionFields = func(*messageDataType) pref.KnownFields {
-		return emptyExtensionFields{}
+	if fx.Type == extensionFieldsType {
+		fieldOffset := offsetOf(fx)
+		mi.extensionMap = func(p pointer) *extensionMap {
+			v := p.Apply(fieldOffset).AsValueOf(extensionFieldsType)
+			return (*extensionMap)(v.Interface().(*map[int32]ExtensionField))
+		}
+	} else {
+		mi.extensionMap = func(pointer) *extensionMap {
+			return (*extensionMap)(nil)
+		}
 	}
 }
 
@@ -295,21 +344,9 @@
 
 type messageReflectWrapper messageDataType
 
-// TODO: Remove this.
-func (m *messageReflectWrapper) Type() pref.MessageType {
-	return m.mi.PBType
-}
 func (m *messageReflectWrapper) Descriptor() pref.MessageDescriptor {
 	return m.mi.PBType.Descriptor()
 }
-func (m *messageReflectWrapper) KnownFields() pref.KnownFields {
-	m.mi.init()
-	return (*knownFields)(m)
-}
-func (m *messageReflectWrapper) UnknownFields() pref.UnknownFields {
-	m.mi.init()
-	return m.mi.unknownFields((*messageDataType)(m))
-}
 func (m *messageReflectWrapper) New() pref.Message {
 	return m.mi.PBType.New()
 }
@@ -323,6 +360,170 @@
 	return m.p.AsIfaceOf(m.mi.GoType.Elem())
 }
 
+func (m *messageReflectWrapper) Len() (cnt int) {
+	m.mi.init()
+	for _, fi := range m.mi.fields {
+		if fi.has(m.p) {
+			cnt++
+		}
+	}
+	return cnt + m.mi.extensionMap(m.p).Len()
+}
+func (m *messageReflectWrapper) Range(f func(pref.FieldDescriptor, pref.Value) bool) {
+	m.mi.init()
+	for _, fi := range m.mi.fields {
+		if fi.has(m.p) {
+			if !f(fi.fieldDesc, fi.get(m.p)) {
+				return
+			}
+		}
+	}
+	m.mi.extensionMap(m.p).Range(f)
+}
+func (m *messageReflectWrapper) Has(fd pref.FieldDescriptor) bool {
+	if fi, xt := m.checkField(fd); fi != nil {
+		return fi.has(m.p)
+	} else {
+		return m.mi.extensionMap(m.p).Has(xt)
+	}
+}
+func (m *messageReflectWrapper) Clear(fd pref.FieldDescriptor) {
+	if fi, xt := m.checkField(fd); fi != nil {
+		fi.clear(m.p)
+	} else {
+		m.mi.extensionMap(m.p).Clear(xt)
+	}
+}
+func (m *messageReflectWrapper) Get(fd pref.FieldDescriptor) pref.Value {
+	if fi, xt := m.checkField(fd); fi != nil {
+		return fi.get(m.p)
+	} else {
+		return m.mi.extensionMap(m.p).Get(xt)
+	}
+}
+func (m *messageReflectWrapper) Set(fd pref.FieldDescriptor, v pref.Value) {
+	if fi, xt := m.checkField(fd); fi != nil {
+		fi.set(m.p, v)
+	} else {
+		m.mi.extensionMap(m.p).Set(xt, v)
+	}
+}
+func (m *messageReflectWrapper) Mutable(fd pref.FieldDescriptor) pref.Value {
+	if fi, xt := m.checkField(fd); fi != nil {
+		return fi.mutable(m.p)
+	} else {
+		return m.mi.extensionMap(m.p).Mutable(xt)
+	}
+}
+func (m *messageReflectWrapper) NewMessage(fd pref.FieldDescriptor) pref.Message {
+	if fi, xt := m.checkField(fd); fi != nil {
+		return fi.newMessage()
+	} else {
+		return xt.New().Message()
+	}
+}
+func (m *messageReflectWrapper) WhichOneof(od pref.OneofDescriptor) pref.FieldDescriptor {
+	m.mi.init()
+	if oi := m.mi.oneofs[od.Name()]; oi != nil && oi.oneofDesc == od {
+		return od.Fields().ByNumber(oi.which(m.p))
+	}
+	panic("invalid oneof descriptor")
+}
+func (m *messageReflectWrapper) GetUnknown() pref.RawFields {
+	m.mi.init()
+	return m.mi.getUnknown(m.p)
+}
+func (m *messageReflectWrapper) SetUnknown(b pref.RawFields) {
+	m.mi.init()
+	m.mi.setUnknown(m.p, b)
+}
+
+// checkField verifies that the provided field descriptor is valid.
+// Exactly one of the returned values is populated.
+func (m *messageReflectWrapper) checkField(fd pref.FieldDescriptor) (*fieldInfo, pref.ExtensionType) {
+	m.mi.init()
+	if fi := m.mi.fields[fd.Number()]; fi != nil {
+		if fi.fieldDesc != fd {
+			panic("mismatching field descriptor")
+		}
+		return fi, nil
+	}
+	if fd.IsExtension() {
+		if fd.ContainingMessage().FullName() != m.mi.PBType.FullName() {
+			// TODO: Should this be exact containing message descriptor match?
+			panic("mismatching containing message")
+		}
+		if !m.mi.PBType.ExtensionRanges().Has(fd.Number()) {
+			panic("invalid extension field")
+		}
+		return nil, fd.(pref.ExtensionType)
+	}
+	panic("invalid field descriptor")
+}
+
+type extensionMap map[int32]ExtensionField
+
+func (m *extensionMap) Len() int {
+	if m != nil {
+		return len(*m)
+	}
+	return 0
+}
+func (m *extensionMap) Range(f func(pref.FieldDescriptor, pref.Value) bool) {
+	if m != nil {
+		for _, x := range *m {
+			xt := x.GetType()
+			if !f(xt, xt.ValueOf(x.GetValue())) {
+				return
+			}
+		}
+	}
+}
+func (m *extensionMap) Has(xt pref.ExtensionType) (ok bool) {
+	if m != nil {
+		_, ok = (*m)[int32(xt.Number())]
+	}
+	return ok
+}
+func (m *extensionMap) Clear(xt pref.ExtensionType) {
+	delete(*m, int32(xt.Number()))
+}
+func (m *extensionMap) Get(xt pref.ExtensionType) pref.Value {
+	if m != nil {
+		if x, ok := (*m)[int32(xt.Number())]; ok {
+			return xt.ValueOf(x.GetValue())
+		}
+	}
+	if !isComposite(xt) {
+		return defaultValueOf(xt)
+	}
+	return frozenValueOf(xt.New())
+}
+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.Number())] = x
+}
+func (m *extensionMap) Mutable(xt pref.ExtensionType) pref.Value {
+	if !isComposite(xt) {
+		panic("invalid Mutable on field with non-composite type")
+	}
+	if x, ok := (*m)[int32(xt.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()
+}
+
 var _ pvalue.Unwrapper = (*messageReflectWrapper)(nil)
 
 type messageIfaceWrapper messageDataType
@@ -348,109 +549,3 @@
 func (m *messageIfaceWrapper) size(msg pref.ProtoMessage) (size int) {
 	return m.mi.sizePointer(m.p, 0)
 }
-
-type knownFields messageDataType
-
-func (fs *knownFields) Len() (cnt int) {
-	for _, fi := range fs.mi.fields {
-		if fi.has(fs.p) {
-			cnt++
-		}
-	}
-	return cnt + fs.extensionFields().Len()
-}
-func (fs *knownFields) Has(n pref.FieldNumber) bool {
-	if fi := fs.mi.fields[n]; fi != nil {
-		return fi.has(fs.p)
-	}
-	return fs.extensionFields().Has(n)
-}
-func (fs *knownFields) Get(n pref.FieldNumber) pref.Value {
-	if fi := fs.mi.fields[n]; fi != nil {
-		return fi.get(fs.p)
-	}
-	return fs.extensionFields().Get(n)
-}
-func (fs *knownFields) Set(n pref.FieldNumber, v pref.Value) {
-	if fi := fs.mi.fields[n]; fi != nil {
-		fi.set(fs.p, v)
-		return
-	}
-	if fs.mi.PBType.Descriptor().ExtensionRanges().Has(n) {
-		fs.extensionFields().Set(n, v)
-		return
-	}
-	panic(fmt.Sprintf("invalid field: %d", n))
-}
-func (fs *knownFields) Clear(n pref.FieldNumber) {
-	if fi := fs.mi.fields[n]; fi != nil {
-		fi.clear(fs.p)
-		return
-	}
-	if fs.mi.PBType.Descriptor().ExtensionRanges().Has(n) {
-		fs.extensionFields().Clear(n)
-		return
-	}
-}
-func (fs *knownFields) WhichOneof(s pref.Name) pref.FieldNumber {
-	if oi := fs.mi.oneofs[s]; oi != nil {
-		return oi.which(fs.p)
-	}
-	return 0
-}
-func (fs *knownFields) Range(f func(pref.FieldNumber, pref.Value) bool) {
-	for n, fi := range fs.mi.fields {
-		if fi.has(fs.p) {
-			if !f(n, fi.get(fs.p)) {
-				return
-			}
-		}
-	}
-	fs.extensionFields().Range(f)
-}
-func (fs *knownFields) NewMessage(n pref.FieldNumber) pref.Message {
-	if fi := fs.mi.fields[n]; fi != nil {
-		return fi.newMessage()
-	}
-	if fs.mi.PBType.Descriptor().ExtensionRanges().Has(n) {
-		return fs.extensionFields().NewMessage(n)
-	}
-	panic(fmt.Sprintf("invalid field: %d", n))
-}
-func (fs *knownFields) ExtensionTypes() pref.ExtensionFieldTypes {
-	return fs.extensionFields().ExtensionTypes()
-}
-func (fs *knownFields) extensionFields() pref.KnownFields {
-	return fs.mi.extensionFields((*messageDataType)(fs))
-}
-
-type emptyUnknownFields struct{}
-
-func (emptyUnknownFields) Len() int                                          { return 0 }
-func (emptyUnknownFields) Get(pref.FieldNumber) pref.RawFields               { return nil }
-func (emptyUnknownFields) Set(pref.FieldNumber, pref.RawFields)              { return } // noop
-func (emptyUnknownFields) Range(func(pref.FieldNumber, pref.RawFields) bool) { return }
-func (emptyUnknownFields) IsSupported() bool                                 { return false }
-
-type emptyExtensionFields struct{}
-
-func (emptyExtensionFields) Len() int                                      { return 0 }
-func (emptyExtensionFields) Has(pref.FieldNumber) bool                     { return false }
-func (emptyExtensionFields) Get(pref.FieldNumber) pref.Value               { return pref.Value{} }
-func (emptyExtensionFields) Set(pref.FieldNumber, pref.Value)              { panic("extensions not supported") }
-func (emptyExtensionFields) Clear(pref.FieldNumber)                        { return } // noop
-func (emptyExtensionFields) WhichOneof(pref.Name) pref.FieldNumber         { return 0 }
-func (emptyExtensionFields) Range(func(pref.FieldNumber, pref.Value) bool) { return }
-func (emptyExtensionFields) NewMessage(pref.FieldNumber) pref.Message {
-	panic("extensions not supported")
-}
-func (emptyExtensionFields) ExtensionTypes() pref.ExtensionFieldTypes { return emptyExtensionTypes{} }
-
-type emptyExtensionTypes struct{}
-
-func (emptyExtensionTypes) Len() int                                     { return 0 }
-func (emptyExtensionTypes) Register(pref.ExtensionType)                  { panic("extensions not supported") }
-func (emptyExtensionTypes) Remove(pref.ExtensionType)                    { return } // noop
-func (emptyExtensionTypes) ByNumber(pref.FieldNumber) pref.ExtensionType { return nil }
-func (emptyExtensionTypes) ByName(pref.FullName) pref.ExtensionType      { return nil }
-func (emptyExtensionTypes) Range(func(pref.ExtensionType) bool)          { return }