size.go 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193
  1. // Go support for Protocol Buffers - Google's data interchange format
  2. //
  3. // Copyright 2012 The Go Authors. All rights reserved.
  4. // http://code.google.com/p/goprotobuf/
  5. //
  6. // Redistribution and use in source and binary forms, with or without
  7. // modification, are permitted provided that the following conditions are
  8. // met:
  9. //
  10. // * Redistributions of source code must retain the above copyright
  11. // notice, this list of conditions and the following disclaimer.
  12. // * Redistributions in binary form must reproduce the above
  13. // copyright notice, this list of conditions and the following disclaimer
  14. // in the documentation and/or other materials provided with the
  15. // distribution.
  16. // * Neither the name of Google Inc. nor the names of its
  17. // contributors may be used to endorse or promote products derived from
  18. // this software without specific prior written permission.
  19. //
  20. // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  21. // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  22. // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  23. // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  24. // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  25. // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  26. // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  27. // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  28. // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  29. // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  30. // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  31. // Functions to determine the size of an encoded protocol buffer.
  32. package proto
  33. import (
  34. "log"
  35. "reflect"
  36. "strings"
  37. )
  38. // Size returns the encoded size of a protocol buffer.
  39. // This function is expensive enough to be avoided unless proven worthwhile with instrumentation.
  40. func Size(pb Message) int {
  41. in := reflect.ValueOf(pb)
  42. if in.IsNil() {
  43. return 0
  44. }
  45. return sizeStruct(in.Elem())
  46. }
  47. func sizeStruct(x reflect.Value) (n int) {
  48. sprop := GetProperties(x.Type())
  49. for _, prop := range sprop.Prop {
  50. if strings.HasPrefix(prop.Name, "XXX_") { // handled below
  51. continue
  52. }
  53. fi, _ := sprop.decoderTags.get(prop.Tag)
  54. f := x.Field(fi)
  55. switch f.Kind() {
  56. case reflect.Ptr:
  57. if f.IsNil() {
  58. continue
  59. }
  60. n += len(prop.tagcode)
  61. f = f.Elem() // avoid a recursion in sizeField
  62. case reflect.Slice:
  63. if f.IsNil() {
  64. continue
  65. }
  66. if f.Len() == 0 && f.Type().Elem().Kind() != reflect.Uint8 {
  67. // short circuit for empty repeated fields.
  68. // []byte isn't a repeated field.
  69. continue
  70. }
  71. default:
  72. log.Printf("proto: unknown struct field type %v", f.Type())
  73. continue
  74. }
  75. n += sizeField(f, prop)
  76. }
  77. if em, ok := x.Addr().Interface().(extendableProto); ok {
  78. for _, ext := range em.ExtensionMap() {
  79. ms := len(ext.enc)
  80. if ext.enc == nil {
  81. props := new(Properties)
  82. props.Init(reflect.TypeOf(ext.desc.ExtensionType), "x", ext.desc.Tag, nil)
  83. ms = len(props.tagcode) + sizeField(reflect.ValueOf(ext.value), props)
  84. }
  85. n += ms
  86. }
  87. }
  88. if uf := x.FieldByName("XXX_unrecognized"); uf.IsValid() {
  89. n += uf.Len()
  90. }
  91. return n
  92. }
  93. func sizeField(x reflect.Value, prop *Properties) (n int) {
  94. if x.Type().Kind() == reflect.Slice {
  95. n := x.Len()
  96. et := x.Type().Elem()
  97. if et.Kind() == reflect.Uint8 {
  98. // []byte is easy.
  99. return len(prop.tagcode) + sizeVarint(uint64(n)) + n
  100. }
  101. var nb int
  102. // []bool and repeated fixed integer types are easy.
  103. switch {
  104. case et.Kind() == reflect.Bool:
  105. nb += n
  106. case prop.WireType == WireFixed64:
  107. nb += n * 8
  108. case prop.WireType == WireFixed32:
  109. nb += n * 4
  110. default:
  111. for i := 0; i < n; i++ {
  112. nb += sizeField(x.Index(i), prop)
  113. }
  114. }
  115. // Non-packed repeated fields have a per-element header of the tagcode.
  116. // Packed repeated fields only have a single header: the tag code plus a varint of the number of bytes.
  117. if !prop.Packed {
  118. nb += len(prop.tagcode) * n
  119. } else {
  120. nb += len(prop.tagcode) + sizeVarint(uint64(nb))
  121. }
  122. return nb
  123. }
  124. // easy scalars
  125. switch prop.WireType {
  126. case WireFixed64:
  127. return 8
  128. case WireFixed32:
  129. return 4
  130. }
  131. switch x.Kind() {
  132. case reflect.Bool:
  133. return 1
  134. case reflect.Int32, reflect.Int64:
  135. if prop.Wire == "varint" {
  136. return sizeVarint(uint64(x.Int()))
  137. } else if prop.Wire == "zigzag32" || prop.Wire == "zigzag64" {
  138. return sizeZigZag(uint64(x.Int()))
  139. }
  140. case reflect.Ptr:
  141. return sizeField(x.Elem(), prop)
  142. case reflect.String:
  143. n := x.Len()
  144. return sizeVarint(uint64(n)) + n
  145. case reflect.Struct:
  146. nb := sizeStruct(x)
  147. if prop.Wire == "group" {
  148. // Groups have start and end tags instead of a start tag and a length.
  149. return nb + len(prop.tagcode)
  150. }
  151. return sizeVarint(uint64(nb)) + nb
  152. case reflect.Uint32, reflect.Uint64:
  153. if prop.Wire == "varint" {
  154. return sizeVarint(uint64(x.Uint()))
  155. } else if prop.Wire == "zigzag32" || prop.Wire == "zigzag64" {
  156. return sizeZigZag(uint64(x.Int()))
  157. }
  158. default:
  159. log.Printf("proto.sizeField: unhandled kind %v", x.Kind())
  160. }
  161. // unknown type, so not a protocol buffer
  162. log.Printf("proto: don't know size of %v", x.Type())
  163. return 0
  164. }
  165. func sizeVarint(x uint64) (n int) {
  166. for {
  167. n++
  168. x >>= 7
  169. if x == 0 {
  170. break
  171. }
  172. }
  173. return n
  174. }
  175. func sizeZigZag(x uint64) (n int) {
  176. return sizeVarint(uint64((x << 1) ^ uint64((int64(x) >> 63))))
  177. }