size.go raw

   1  // Copyright 2019 The Go Authors. All rights reserved.
   2  // Use of this source code is governed by a BSD-style
   3  // license that can be found in the LICENSE file.
   4  
   5  package proto
   6  
   7  import (
   8  	"google.golang.org/protobuf/encoding/protowire"
   9  	"google.golang.org/protobuf/internal/encoding/messageset"
  10  	"google.golang.org/protobuf/reflect/protoreflect"
  11  	"google.golang.org/protobuf/runtime/protoiface"
  12  )
  13  
  14  // Size returns the size in bytes of the wire-format encoding of m.
  15  //
  16  // Note that Size might return more bytes than Marshal will write in the case of
  17  // lazily decoded messages that arrive in non-minimal wire format: see
  18  // https://protobuf.dev/reference/go/size/ for more details.
  19  func Size(m Message) int {
  20  	return MarshalOptions{}.Size(m)
  21  }
  22  
  23  // Size returns the size in bytes of the wire-format encoding of m.
  24  //
  25  // Note that Size might return more bytes than Marshal will write in the case of
  26  // lazily decoded messages that arrive in non-minimal wire format: see
  27  // https://protobuf.dev/reference/go/size/ for more details.
  28  func (o MarshalOptions) Size(m Message) int {
  29  	// Treat a nil message interface as an empty message; nothing to output.
  30  	if m == nil {
  31  		return 0
  32  	}
  33  
  34  	return o.size(m.ProtoReflect())
  35  }
  36  
  37  // size is a centralized function that all size operations go through.
  38  // For profiling purposes, avoid changing the name of this function or
  39  // introducing other code paths for size that do not go through this.
  40  func (o MarshalOptions) size(m protoreflect.Message) (size int) {
  41  	methods := protoMethods(m)
  42  	if methods != nil && methods.Size != nil {
  43  		out := methods.Size(protoiface.SizeInput{
  44  			Message: m,
  45  			Flags:   o.flags(),
  46  		})
  47  		return out.Size
  48  	}
  49  	if methods != nil && methods.Marshal != nil {
  50  		// This is not efficient, but we don't have any choice.
  51  		// This case is mainly used for legacy types with a Marshal method.
  52  		out, _ := methods.Marshal(protoiface.MarshalInput{
  53  			Message: m,
  54  			Flags:   o.flags(),
  55  		})
  56  		return len(out.Buf)
  57  	}
  58  	return o.sizeMessageSlow(m)
  59  }
  60  
  61  func (o MarshalOptions) sizeMessageSlow(m protoreflect.Message) (size int) {
  62  	if messageset.IsMessageSet(m.Descriptor()) {
  63  		return o.sizeMessageSet(m)
  64  	}
  65  	m.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool {
  66  		size += o.sizeField(fd, v)
  67  		return true
  68  	})
  69  	size += len(m.GetUnknown())
  70  	return size
  71  }
  72  
  73  func (o MarshalOptions) sizeField(fd protoreflect.FieldDescriptor, value protoreflect.Value) (size int) {
  74  	num := fd.Number()
  75  	switch {
  76  	case fd.IsList():
  77  		return o.sizeList(num, fd, value.List())
  78  	case fd.IsMap():
  79  		return o.sizeMap(num, fd, value.Map())
  80  	default:
  81  		return protowire.SizeTag(num) + o.sizeSingular(num, fd.Kind(), value)
  82  	}
  83  }
  84  
  85  func (o MarshalOptions) sizeList(num protowire.Number, fd protoreflect.FieldDescriptor, list protoreflect.List) (size int) {
  86  	sizeTag := protowire.SizeTag(num)
  87  
  88  	if fd.IsPacked() && list.Len() > 0 {
  89  		content := 0
  90  		for i, llen := 0, list.Len(); i < llen; i++ {
  91  			content += o.sizeSingular(num, fd.Kind(), list.Get(i))
  92  		}
  93  		return sizeTag + protowire.SizeBytes(content)
  94  	}
  95  
  96  	for i, llen := 0, list.Len(); i < llen; i++ {
  97  		size += sizeTag + o.sizeSingular(num, fd.Kind(), list.Get(i))
  98  	}
  99  	return size
 100  }
 101  
 102  func (o MarshalOptions) sizeMap(num protowire.Number, fd protoreflect.FieldDescriptor, mapv protoreflect.Map) (size int) {
 103  	sizeTag := protowire.SizeTag(num)
 104  
 105  	mapv.Range(func(key protoreflect.MapKey, value protoreflect.Value) bool {
 106  		size += sizeTag
 107  		size += protowire.SizeBytes(o.sizeField(fd.MapKey(), key.Value()) + o.sizeField(fd.MapValue(), value))
 108  		return true
 109  	})
 110  	return size
 111  }
 112