goprotobuf: Add a Size function.

R=r
CC=golang-dev
http://codereview.appspot.com/6812097

proto/size.go

@@ -0,0 +1,184 @@
+// Go support for Protocol Buffers - Google's data interchange format
+//
+// Copyright 2012 Google Inc.  All rights reserved.
+// http://code.google.com/p/goprotobuf/
+//
+// 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.
+//     * Neither the name of Google Inc. nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// 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.
+
+// Functions to determine the size of an encoded protocol buffer.
+
+package proto
+
+import (
+	"log"
+	"reflect"
+	"strings"
+)
+
+// Size returns the encoded size of a protocol buffer.
+// This function is expensive enough to be avoided unless proven worthwhile with instrumentation.
+func Size(pb Message) int {
+	in := reflect.ValueOf(pb)
+	if in.IsNil() {
+		return 0
+	}
+	return sizeStruct(in.Elem())
+}
+
+func sizeStruct(x reflect.Value) (n int) {
+	sprop := GetProperties(x.Type())
+	for _, prop := range sprop.Prop {
+		if strings.HasPrefix(prop.Name, "XXX_") { // handled below
+			continue
+		}
+		fi, _ := sprop.decoderTags.get(prop.Tag)
+		f := x.Field(fi)
+		switch f.Kind() {
+		case reflect.Ptr:
+			if f.IsNil() {
+				continue
+			}
+			n += len(prop.tagcode)
+			f = f.Elem() // avoid a recursion in sizeField
+		case reflect.Slice:
+			if f.IsNil() {
+				continue
+			}
+			if f.Len() == 0 && f.Type().Elem().Kind() != reflect.Uint8 {
+				// short circuit for empty repeated fields.
+				// []byte isn't a repeated field.
+				continue
+			}
+		default:
+			log.Printf("proto: unknown struct field type %v", f.Type())
+			continue
+		}
+		n += sizeField(f, prop)
+	}
+
+	if em, ok := x.Addr().Interface().(extendableProto); ok {
+		for _, ext := range em.ExtensionMap() {
+			ms := len(ext.enc)
+			if ext.enc == nil {
+				props := new(Properties)
+				props.Init(reflect.TypeOf(ext.desc.ExtensionType), "x", ext.desc.Tag, nil)
+				ms = len(props.tagcode) + sizeField(reflect.ValueOf(ext.value), props)
+			}
+			n += ms
+		}
+	}
+
+	// Groups don't have XXX_unrecognized fields.
+	if uf := x.FieldByName("XXX_unrecognized"); uf.IsValid() {
+		n += uf.Len()
+	}
+
+	return n
+}
+
+func sizeField(x reflect.Value, prop *Properties) (n int) {
+	// Handle easy ones first.
+	switch prop.WireType {
+	case WireFixed64:
+		return 8
+	case WireFixed32:
+		return 4
+	}
+
+	if x.Type().Kind() == reflect.Slice {
+		n := x.Len()
+		et := x.Type().Elem()
+		if et.Kind() == reflect.Uint8 {
+			// []byte is easy.
+			return len(prop.tagcode) + sizeVarint(uint64(n)) + n
+		}
+
+		// Non-packed repeated fields have a per-element header of the tagcode.
+		// Packed repeated fields only have a single header: the tag code plus a varint of the number of bytes.
+		var nb int
+		if !prop.Packed {
+			nb = len(prop.tagcode) * n
+		} else {
+			nb = len(prop.tagcode) + sizeVarint(uint64(n))
+		}
+
+		if et.Kind() == reflect.Bool {
+			// []bool is easy to compute: each element requires one byte.
+			return nb + n
+		}
+		for i := 0; i < n; i++ {
+			nb += sizeField(x.Index(i), prop)
+		}
+		return nb
+	}
+
+	switch x.Kind() {
+	case reflect.Bool:
+		return 1
+	case reflect.Int32, reflect.Int64:
+		if prop.Wire == "varint" {
+			return sizeVarint(uint64(x.Int()))
+		} else if prop.Wire == "zigzag32" || prop.Wire == "zigzag64" {
+			return sizeZigZag(uint64(x.Int()))
+		}
+	case reflect.Ptr:
+		return sizeField(x.Elem(), prop)
+	case reflect.String:
+		n := x.Len()
+		return sizeVarint(uint64(n)) + n
+	case reflect.Struct:
+		nb := sizeStruct(x)
+		return sizeVarint(uint64(nb)) + nb
+	case reflect.Uint32, reflect.Uint64:
+		if prop.Wire == "varint" {
+			return sizeVarint(uint64(x.Uint()))
+		} else if prop.Wire == "zigzag32" || prop.Wire == "zigzag64" {
+			return sizeZigZag(uint64(x.Int()))
+		}
+	default:
+		log.Printf("proto.sizeField: unhandled kind %v", x.Kind())
+	}
+
+	// unknown type, so not a protocol buffer
+	log.Printf("proto: don't know size of %v", x.Type())
+	return 0
+}
+
+func sizeVarint(x uint64) (n int) {
+	for {
+		n++
+		x >>= 7
+		if x == 0 {
+			break
+		}
+	}
+	return n
+}
+
+func sizeZigZag(x uint64) (n int) {
+	return sizeVarint(uint64((x << 1) ^ uint64((int64(x) >> 63))))
+}

proto/size2_test.go

@@ -0,0 +1,63 @@
+// Go support for Protocol Buffers - Google's data interchange format
+//
+// Copyright 2012 Google Inc.  All rights reserved.
+// http://code.google.com/p/goprotobuf/
+//
+// 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.
+//     * Neither the name of Google Inc. nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// 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 (
+	"testing"
+)
+
+// This is a separate file and package from size_test.go because that one uses
+// generated messages and thus may not be in package proto without having a circular
+// dependency, whereas this file tests unexported details of size.go.
+
+func TestVarintSize(t *testing.T) {
+	// Check the edge cases carefully.
+	testCases := []struct {
+		n    uint64
+		size int
+	}{
+		{0, 1},
+		{1, 1},
+		{127, 1},
+		{128, 2},
+		{16383, 2},
+		{16384, 3},
+		{1<<63 - 1, 9},
+		{1 << 63, 10},
+	}
+	for _, tc := range testCases {
+		size := sizeVarint(tc.n)
+		if size != tc.size {
+			t.Errorf("sizeVarint(%d) = %d, want %d", tc.n, size, tc.size)
+		}
+	}
+}

proto/size_test.go

@@ -0,0 +1,107 @@
+// Go support for Protocol Buffers - Google's data interchange format
+//
+// Copyright 2012 Google Inc.  All rights reserved.
+// http://code.google.com/p/goprotobuf/
+//
+// 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.
+//     * Neither the name of Google Inc. nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// 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_test
+
+import (
+	"log"
+	"testing"
+
+	pb "./testdata"
+	. "code.google.com/p/goprotobuf/proto"
+)
+
+var messageWithExtension1 = &pb.MyMessage{Count: Int32(7)}
+var messageWithExtension3 = &pb.MyMessage{Count: Int32(8)}
+
+func init() {
+	if err := SetExtension(messageWithExtension1, pb.E_Ext_More, &pb.Ext{Data: String("Abbott")}); err != nil {
+		log.Panicf("SetExtension: %v", err)
+	}
+	if err := SetExtension(messageWithExtension3, pb.E_Ext_More, &pb.Ext{Data: String("Costello")}); err != nil {
+		log.Panicf("SetExtension: %v", err)
+	}
+
+	// Force messageWithExtension3 to have the extension encoded.
+	Marshal(messageWithExtension3)
+}
+
+var SizeTests = []struct {
+	desc string
+	pb   Message
+}{
+	{"empty", &pb.OtherMessage{}},
+	// Basic types.
+	{"bool", &pb.Defaults{F_Bool: Bool(true)}},
+	{"int32", &pb.Defaults{F_Int32: Int32(12)}},
+	{"small int64", &pb.Defaults{F_Int64: Int64(1)}},
+	{"big int64", &pb.Defaults{F_Int64: Int64(1 << 20)}},
+	{"fixed32", &pb.Defaults{F_Fixed32: Uint32(71)}},
+	{"fixed64", &pb.Defaults{F_Fixed64: Uint64(72)}},
+	{"uint32", &pb.Defaults{F_Uint32: Uint32(123)}},
+	{"uint64", &pb.Defaults{F_Uint64: Uint64(124)}},
+	{"float", &pb.Defaults{F_Float: Float32(12.6)}},
+	{"double", &pb.Defaults{F_Double: Float64(13.9)}},
+	{"string", &pb.Defaults{F_String: String("niles")}},
+	{"bytes", &pb.Defaults{F_Bytes: []byte("wowsa")}},
+	{"bytes, empty", &pb.Defaults{F_Bytes: []byte{}}},
+	{"sint32", &pb.Defaults{F_Sint32: Int32(65)}},
+	{"sint64", &pb.Defaults{F_Sint64: Int64(67)}},
+	{"enum", &pb.Defaults{F_Enum: pb.Defaults_BLUE.Enum()}},
+	// Repeated.
+	{"empty repeated bool", &pb.MoreRepeated{Bools: []bool{}}},
+	{"repeated bool", &pb.MoreRepeated{Bools: []bool{false, true, true, false}}},
+	{"packed repeated bool", &pb.MoreRepeated{BoolsPacked: []bool{false, true, true, false, true, true, true}}},
+	{"repeated int32", &pb.MoreRepeated{Ints: []int32{1, 12203, 1729}}},
+	{"repeated int32 packed", &pb.MoreRepeated{IntsPacked: []int32{1, 12203, 1729}}},
+	{"repeated string", &pb.MoreRepeated{Strings: []string{"r", "ken", "gri"}}},
+	// Nested.
+	{"nested", &pb.OldMessage{Nested: &pb.OldMessage_Nested{Name: String("whatever")}}},
+	// Other things.
+	{"unrecognized", &pb.MoreRepeated{XXX_unrecognized: []byte{13<<3 | 0, 4}}},
+	{"extension (unencoded)", messageWithExtension1},
+	{"extension (encoded)", messageWithExtension3},
+}
+
+func TestSize(t *testing.T) {
+	for _, tc := range SizeTests {
+		size := Size(tc.pb)
+		b, err := Marshal(tc.pb)
+		if err != nil {
+			t.Errorf("%v: Marshal failed: %v", tc.desc, err)
+			continue
+		}
+		if size != len(b) {
+			t.Errorf("%v: Size(%v) = %d, want %d", tc.desc, tc.pb, size, len(b))
+			t.Logf("%v: bytes: %#v", tc.desc, b)
+		}
+	}
+}

proto/testdata/test.pb.go

@@ -685,13 +685,10 @@ func (this *GoTest) GetOptionalgroup() *GoTest_OptionalGroup {
 }
 
 type GoTest_RequiredGroup struct {
-	RequiredField    *string `protobuf:"bytes,71,req" json:"RequiredField,omitempty"`
-	XXX_unrecognized []byte  `json:"-"`
+	RequiredField *string `protobuf:"bytes,71,req" json:"RequiredField,omitempty"`
 }
 
-func (this *GoTest_RequiredGroup) Reset()         { *this = GoTest_RequiredGroup{} }
-func (this *GoTest_RequiredGroup) String() string { return proto.CompactTextString(this) }
-func (*GoTest_RequiredGroup) ProtoMessage()       {}
+func (this *GoTest_RequiredGroup) Reset() { *this = GoTest_RequiredGroup{} }
 
 func (this *GoTest_RequiredGroup) GetRequiredField() string {
 	if this != nil && this.RequiredField != nil {
@@ -701,13 +698,10 @@ func (this *GoTest_RequiredGroup) GetRequiredField() string {
 }
 
 type GoTest_RepeatedGroup struct {
-	RequiredField    *string `protobuf:"bytes,81,req" json:"RequiredField,omitempty"`
-	XXX_unrecognized []byte  `json:"-"`
+	RequiredField *string `protobuf:"bytes,81,req" json:"RequiredField,omitempty"`
 }
 
-func (this *GoTest_RepeatedGroup) Reset()         { *this = GoTest_RepeatedGroup{} }
-func (this *GoTest_RepeatedGroup) String() string { return proto.CompactTextString(this) }
-func (*GoTest_RepeatedGroup) ProtoMessage()       {}
+func (this *GoTest_RepeatedGroup) Reset() { *this = GoTest_RepeatedGroup{} }
 
 func (this *GoTest_RepeatedGroup) GetRequiredField() string {
 	if this != nil && this.RequiredField != nil {
@@ -717,13 +711,10 @@ func (this *GoTest_RepeatedGroup) GetRequiredField() string {
 }
 
 type GoTest_OptionalGroup struct {
-	RequiredField    *string `protobuf:"bytes,91,req" json:"RequiredField,omitempty"`
-	XXX_unrecognized []byte  `json:"-"`
+	RequiredField *string `protobuf:"bytes,91,req" json:"RequiredField,omitempty"`
 }
 
-func (this *GoTest_OptionalGroup) Reset()         { *this = GoTest_OptionalGroup{} }
-func (this *GoTest_OptionalGroup) String() string { return proto.CompactTextString(this) }
-func (*GoTest_OptionalGroup) ProtoMessage()       {}
+func (this *GoTest_OptionalGroup) Reset() { *this = GoTest_OptionalGroup{} }
 
 func (this *GoTest_OptionalGroup) GetRequiredField() string {
 	if this != nil && this.RequiredField != nil {
@@ -781,14 +772,11 @@ func (this *GoSkipTest) GetSkipgroup() *GoSkipTest_SkipGroup {
 }
 
 type GoSkipTest_SkipGroup struct {
-	GroupInt32       *int32  `protobuf:"varint,16,req,name=group_int32" json:"group_int32,omitempty"`
-	GroupString      *string `protobuf:"bytes,17,req,name=group_string" json:"group_string,omitempty"`
-	XXX_unrecognized []byte  `json:"-"`
+	GroupInt32  *int32  `protobuf:"varint,16,req,name=group_int32" json:"group_int32,omitempty"`
+	GroupString *string `protobuf:"bytes,17,req,name=group_string" json:"group_string,omitempty"`
 }
 
-func (this *GoSkipTest_SkipGroup) Reset()         { *this = GoSkipTest_SkipGroup{} }
-func (this *GoSkipTest_SkipGroup) String() string { return proto.CompactTextString(this) }
-func (*GoSkipTest_SkipGroup) ProtoMessage()       {}
+func (this *GoSkipTest_SkipGroup) Reset() { *this = GoSkipTest_SkipGroup{} }
 
 func (this *GoSkipTest_SkipGroup) GetGroupInt32() int32 {
 	if this != nil && this.GroupInt32 != nil {
@@ -1059,13 +1047,10 @@ func (this *MyMessage) GetSomegroup() *MyMessage_SomeGroup {
 }
 
 type MyMessage_SomeGroup struct {
-	GroupField       *int32 `protobuf:"varint,9,opt,name=group_field" json:"group_field,omitempty"`
-	XXX_unrecognized []byte `json:"-"`
+	GroupField *int32 `protobuf:"varint,9,opt,name=group_field" json:"group_field,omitempty"`
 }
 
-func (this *MyMessage_SomeGroup) Reset()         { *this = MyMessage_SomeGroup{} }
-func (this *MyMessage_SomeGroup) String() string { return proto.CompactTextString(this) }
-func (*MyMessage_SomeGroup) ProtoMessage()       {}
+func (this *MyMessage_SomeGroup) Reset() { *this = MyMessage_SomeGroup{} }
 
 func (this *MyMessage_SomeGroup) GetGroupField() int32 {
 	if this != nil && this.GroupField != nil {
@@ -1124,14 +1109,11 @@ func (this *MessageList) String() string { return proto.CompactTextString(this)
 func (*MessageList) ProtoMessage()       {}
 
 type MessageList_Message struct {
-	Name             *string `protobuf:"bytes,2,req,name=name" json:"name,omitempty"`
-	Count            *int32  `protobuf:"varint,3,req,name=count" json:"count,omitempty"`
-	XXX_unrecognized []byte  `json:"-"`
+	Name  *string `protobuf:"bytes,2,req,name=name" json:"name,omitempty"`
+	Count *int32  `protobuf:"varint,3,req,name=count" json:"count,omitempty"`
 }
 
-func (this *MessageList_Message) Reset()         { *this = MessageList_Message{} }
-func (this *MessageList_Message) String() string { return proto.CompactTextString(this) }
-func (*MessageList_Message) ProtoMessage()       {}
+func (this *MessageList_Message) Reset() { *this = MessageList_Message{} }
 
 func (this *MessageList_Message) GetName() string {
 	if this != nil && this.Name != nil {
@@ -1371,6 +1353,19 @@ func (this *RepeatedEnum) Reset()         { *this = RepeatedEnum{} }
 func (this *RepeatedEnum) String() string { return proto.CompactTextString(this) }
 func (*RepeatedEnum) ProtoMessage()       {}
 
+type MoreRepeated struct {
+	Bools            []bool   `protobuf:"varint,1,rep,name=bools" json:"bools,omitempty"`
+	BoolsPacked      []bool   `protobuf:"varint,2,rep,packed,name=bools_packed" json:"bools_packed,omitempty"`
+	Ints             []int32  `protobuf:"varint,3,rep,name=ints" json:"ints,omitempty"`
+	IntsPacked       []int32  `protobuf:"varint,4,rep,packed,name=ints_packed" json:"ints_packed,omitempty"`
+	Strings          []string `protobuf:"bytes,5,rep,name=strings" json:"strings,omitempty"`
+	XXX_unrecognized []byte   `json:"-"`
+}
+
+func (this *MoreRepeated) Reset()         { *this = MoreRepeated{} }
+func (this *MoreRepeated) String() string { return proto.CompactTextString(this) }
+func (*MoreRepeated) ProtoMessage()       {}
+
 var E_Greeting = &proto.ExtensionDesc{
 	ExtendedType:  (*MyMessage)(nil),
 	ExtensionType: ([]string)(nil),

proto/testdata/test.proto

@@ -324,3 +324,11 @@ message RepeatedEnum {
   }
   repeated Color color = 1;
 }
+
+message MoreRepeated {
+  repeated bool bools = 1;
+  repeated bool bools_packed = 2 [packed=true];
+  repeated int32 ints = 3;
+  repeated int32 ints_packed = 4 [packed=true];
+  repeated string strings = 5;
+}