Swap type assert and fix it (gogo#418)

* swap type assert back

* type assert rather than check for marshal function at compile time

* Fix swap type assert to only call Size when Marshal is generated.

* added Sizer type assert at cachedsize method. The sizecache is not calculated if the struct conforms to the Sizer interface

* fixed go vet issue. removed err variable shadowing

* refactor pointer interface type conversion out in order to implement both reflect and unsafe cases

* added minimal test example

* added sizer option when computing the marshalinfo and generator now either use MarshalTo or default marhsaling. hassizer property now checked in the marshalinfo size method if the type has a Size() method. removed the Sizer interface check in the generated XXX_Size methods.

* remove unnecessary variable and just return size.

* added a check if the message have a protosize method. changed the size and cached size methods to just return the message size without looking at any other struct fields.

* removed the sizer checks in the caching method. only use size mehtods of the message is also a marshaler

* revert the bench result accidental update

* do a generate time check for XXX_unmarshal.

Special thanks to @jmarais for doing all the work for this commit.

Makefile

@@ -119,6 +119,7 @@ regenerate:
 	make -C test/issue330 regenerate
 	make -C test/importcustom-issue389 regenerate
 	make -C test/merge regenerate
+	make -C test/cachedsize regenerate
 	make gofmt
 
 tests:

proto/decode.go

@@ -334,12 +334,12 @@ type newUnmarshaler interface {
 // to preserve and append to existing data.
 func Unmarshal(buf []byte, pb Message) error {
 	pb.Reset()
-	if u, ok := pb.(Unmarshaler); ok {
-		return u.Unmarshal(buf)
-	}
 	if u, ok := pb.(newUnmarshaler); ok {
 		return u.XXX_Unmarshal(buf)
 	}
+	if u, ok := pb.(Unmarshaler); ok {
+		return u.Unmarshal(buf)
+	}
 	return NewBuffer(buf).Unmarshal(pb)
 }
 
@@ -350,6 +350,9 @@ func Unmarshal(buf []byte, pb Message) error {
 // UnmarshalMerge merges into existing data in pb.
 // Most code should use Unmarshal instead.
 func UnmarshalMerge(buf []byte, pb Message) error {
+	if u, ok := pb.(newUnmarshaler); ok {
+		return u.XXX_Unmarshal(buf)
+	}
 	if u, ok := pb.(Unmarshaler); ok {
 		// NOTE: The history of proto have unfortunately been inconsistent
 		// whether Unmarshaler should or should not implicitly clear itself.
@@ -359,9 +362,6 @@ func UnmarshalMerge(buf []byte, pb Message) error {
 		// See https://github.com/golang/protobuf/issues/424
 		return u.Unmarshal(buf)
 	}
-	if u, ok := pb.(newUnmarshaler); ok {
-		return u.XXX_Unmarshal(buf)
-	}
 	return NewBuffer(buf).Unmarshal(pb)
 }
 
@@ -396,6 +396,11 @@ func (p *Buffer) DecodeGroup(pb Message) error {
 // Unlike proto.Unmarshal, this does not reset pb before starting to unmarshal.
 func (p *Buffer) Unmarshal(pb Message) error {
 	// If the object can unmarshal itself, let it.
+	if u, ok := pb.(newUnmarshaler); ok {
+		err := u.XXX_Unmarshal(p.buf[p.index:])
+		p.index = len(p.buf)
+		return err
+	}
 	if u, ok := pb.(Unmarshaler); ok {
 		// NOTE: The history of proto have unfortunately been inconsistent
 		// whether Unmarshaler should or should not implicitly clear itself.
@@ -407,11 +412,6 @@ func (p *Buffer) Unmarshal(pb Message) error {
 		p.index = len(p.buf)
 		return err
 	}
-	if u, ok := pb.(newUnmarshaler); ok {
-		err := u.XXX_Unmarshal(p.buf[p.index:])
-		p.index = len(p.buf)
-		return err
-	}
 
 	// Slow workaround for messages that aren't Unmarshalers.
 	// This includes some hand-coded .pb.go files and

proto/lib_gogo.go

@@ -37,6 +37,10 @@ type Sizer interface {
 	Size() int
 }
 
+type ProtoSizer interface {
+	ProtoSize() int
+}
+
 func MarshalJSONEnum(m map[int32]string, value int32) ([]byte, error) {
 	s, ok := m[value]
 	if !ok {

proto/properties_gogo.go

@@ -0,0 +1,36 @@
+// Protocol Buffers for Go with Gadgets
+//
+// Copyright (c) 2018, The GoGo Authors. All rights reserved.
+// http://github.com/gogo/protobuf
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are
+// met:
+//
+//     * Redistributions of source code must retain the above copyright
+// notice, this list of conditions and the following disclaimer.
+//     * Redistributions in binary form must reproduce the above
+// copyright notice, this list of conditions and the following disclaimer
+// in the documentation and/or other materials provided with the
+// distribution.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+package proto
+
+import (
+	"reflect"
+)
+
+var sizerType = reflect.TypeOf((*Sizer)(nil)).Elem()
+var protosizerType = reflect.TypeOf((*ProtoSizer)(nil)).Elem()

proto/table_marshal.go

@@ -66,6 +66,9 @@ type marshalInfo struct {
 	sync.RWMutex                            // protect extElems map, also for initialization
 	extElems     map[int32]*marshalElemInfo // info of extension elements
 
+	hassizer      bool // has custom sizer
+	hasprotosizer bool // has custom protosizer
+
 	bytesExtensions field // offset of XXX_extensions where the field type is []byte
 }
 
@@ -170,6 +173,17 @@ func (u *marshalInfo) size(ptr pointer) int {
 	// If the message can marshal itself, let it do it, for compatibility.
 	// NOTE: This is not efficient.
 	if u.hasmarshaler {
+		// Uses the message's Size method if available
+		if u.hassizer {
+			s := ptr.asPointerTo(u.typ).Interface().(Sizer)
+			return s.Size()
+		}
+		// Uses the message's ProtoSize method if available
+		if u.hasprotosizer {
+			s := ptr.asPointerTo(u.typ).Interface().(ProtoSizer)
+			return s.ProtoSize()
+		}
+
 		m := ptr.asPointerTo(u.typ).Interface().(Marshaler)
 		b, _ := m.Marshal()
 		return len(b)
@@ -203,6 +217,7 @@ func (u *marshalInfo) size(ptr pointer) int {
 		s := *ptr.offset(u.unrecognized).toBytes()
 		n += len(s)
 	}
+
 	// cache the result for use in marshal
 	if u.sizecache.IsValid() {
 		atomic.StoreInt32(ptr.offset(u.sizecache).toInt32(), int32(n))
@@ -313,6 +328,12 @@ func (u *marshalInfo) computeMarshalInfo() {
 	u.sizecache = invalidField
 	isOneofMessage := false
 
+	if reflect.PtrTo(t).Implements(sizerType) {
+		u.hassizer = true
+	}
+	if reflect.PtrTo(t).Implements(protosizerType) {
+		u.hasprotosizer = true
+	}
 	// If the message can marshal itself, let it do it, for compatibility.
 	// NOTE: This is not efficient.
 	if reflect.PtrTo(t).Implements(marshalerType) {
@@ -2686,15 +2707,15 @@ type newMarshaler interface {
 // Size returns the encoded size of a protocol buffer message.
 // This is the main entry point.
 func Size(pb Message) int {
+	if m, ok := pb.(newMarshaler); ok {
+		return m.XXX_Size()
+	}
 	if m, ok := pb.(Marshaler); ok {
 		// If the message can marshal itself, let it do it, for compatibility.
 		// NOTE: This is not efficient.
 		b, _ := m.Marshal()
 		return len(b)
 	}
-	if m, ok := pb.(newMarshaler); ok {
-		return m.XXX_Size()
-	}
 	// in case somehow we didn't generate the wrapper
 	if pb == nil {
 		return 0
@@ -2707,16 +2728,16 @@ func Size(pb Message) int {
 // and encodes it into the wire format, returning the data.
 // This is the main entry point.
 func Marshal(pb Message) ([]byte, error) {
-	if m, ok := pb.(Marshaler); ok {
-		// If the message can marshal itself, let it do it, for compatibility.
-		// NOTE: This is not efficient.
-		return m.Marshal()
-	}
 	if m, ok := pb.(newMarshaler); ok {
 		siz := m.XXX_Size()
 		b := make([]byte, 0, siz)
 		return m.XXX_Marshal(b, false)
 	}
+	if m, ok := pb.(Marshaler); ok {
+		// If the message can marshal itself, let it do it, for compatibility.
+		// NOTE: This is not efficient.
+		return m.Marshal()
+	}
 	// in case somehow we didn't generate the wrapper
 	if pb == nil {
 		return nil, ErrNil
@@ -2733,21 +2754,21 @@ func Marshal(pb Message) ([]byte, error) {
 // This is an alternative entry point. It is not necessary to use
 // a Buffer for most applications.
 func (p *Buffer) Marshal(pb Message) error {
-	if m, ok := pb.(Marshaler); ok {
-		// If the message can marshal itself, let it do it, for compatibility.
-		// NOTE: This is not efficient.
-		b, err := m.Marshal()
-		p.buf = append(p.buf, b...)
-		return err
-	}
 	var err error
 	if m, ok := pb.(newMarshaler); ok {
 		siz := m.XXX_Size()
-		// make sure buf has enough capacity
 		p.grow(siz) // make sure buf has enough capacity
 		p.buf, err = m.XXX_Marshal(p.buf, p.deterministic)
 		return err
 	}
+	if m, ok := pb.(Marshaler); ok {
+		// If the message can marshal itself, let it do it, for compatibility.
+		// NOTE: This is not efficient.
+		var b []byte
+		b, err = m.Marshal()
+		p.buf = append(p.buf, b...)
+		return err
+	}
 	// in case somehow we didn't generate the wrapper
 	if pb == nil {
 		return ErrNil

protoc-gen-gogo/generator/generator.go

@@ -2332,13 +2332,29 @@ func (g *Generator) generateMessage(message *Descriptor) {
 	// Wrapper for table-driven marshaling and unmarshaling.
 	g.P("func (m *", ccTypeName, ") XXX_Unmarshal(b []byte) error {")
 	g.In()
-	g.P("return xxx_messageInfo_", ccTypeName, ".Unmarshal(m, b)")
+	if gogoproto.IsUnmarshaler(g.file.FileDescriptorProto, message.DescriptorProto) {
+		g.P("return m.Unmarshal(b)")
+	} else {
+		g.P("return xxx_messageInfo_", ccTypeName, ".Unmarshal(m, b)")
+	}
 	g.Out()
 	g.P("}")
 
 	g.P("func (m *", ccTypeName, ") XXX_Marshal(b []byte, deterministic bool) ([]byte, error) {")
 	g.In()
-	g.P("return xxx_messageInfo_", ccTypeName, ".Marshal(b, m, deterministic)")
+	if gogoproto.IsMarshaler(g.file.FileDescriptorProto, message.DescriptorProto) ||
+		gogoproto.IsUnsafeMarshaler(g.file.FileDescriptorProto, message.DescriptorProto) {
+		g.P("b = b[:cap(b)]")
+		g.P("n, err := m.MarshalTo(b)")
+		g.P("if err != nil {")
+		g.In()
+		g.P("return nil, err")
+		g.Out()
+		g.P("}")
+		g.P("return b[:n], nil")
+	} else {
+		g.P("return xxx_messageInfo_", ccTypeName, ".Marshal(b, m, deterministic)")
+	}
 	g.Out()
 	g.P("}")
 
@@ -2350,7 +2366,17 @@ func (g *Generator) generateMessage(message *Descriptor) {
 
 	g.P("func (m *", ccTypeName, ") XXX_Size() int {") // avoid name clash with "Size" field in some message
 	g.In()
-	g.P("return xxx_messageInfo_", ccTypeName, ".Size(m)")
+	if (gogoproto.IsMarshaler(g.file.FileDescriptorProto, message.DescriptorProto) ||
+		gogoproto.IsUnsafeMarshaler(g.file.FileDescriptorProto, message.DescriptorProto)) &&
+		gogoproto.IsSizer(g.file.FileDescriptorProto, message.DescriptorProto) {
+		g.P("return m.Size()")
+	} else if (gogoproto.IsMarshaler(g.file.FileDescriptorProto, message.DescriptorProto) ||
+		gogoproto.IsUnsafeMarshaler(g.file.FileDescriptorProto, message.DescriptorProto)) &&
+		gogoproto.IsProtoSizer(g.file.FileDescriptorProto, message.DescriptorProto) {
+		g.P("return m.ProtoSize()")
+	} else {
+		g.P("return xxx_messageInfo_", ccTypeName, ".Size(m)")
+	}
 	g.Out()
 	g.P("}")
 

test/cachedsize/Makefile

@@ -0,0 +1,30 @@
+# Protocol Buffers for Go with Gadgets
+#
+# Copyright (c) 2018, The GoGo Authors. All rights reserved.
+# http://github.com/gogo/protobuf
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are
+# met:
+#
+#     * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+#     * Redistributions in binary form must reproduce the above
+# copyright notice, this list of conditions and the following disclaimer
+# in the documentation and/or other materials provided with the
+# distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+regenerate:
+	(protoc --proto_path=../../../../../:../../protobuf/:. --gogo_out=. cachedsize.proto)