size.go

  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
  5package proto
  6
  7import (
  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.
 19func 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.
 28func (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.
 40func (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
 61func (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
 73func (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
 85func (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
102func (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}