go-bot

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
func Size(m Message) int {
16
	return MarshalOptions{}.Size(m)
17
}
18

19
// Size returns the size in bytes of the wire-format encoding of m.
20
func (o MarshalOptions) Size(m Message) int {
21
	// Treat a nil message interface as an empty message; nothing to output.
22
	if m == nil {
23
		return 0
24
	}
25

26
	return o.size(m.ProtoReflect())
27
}
28

29
// size is a centralized function that all size operations go through.
30
// For profiling purposes, avoid changing the name of this function or
31
// introducing other code paths for size that do not go through this.
32
func (o MarshalOptions) size(m protoreflect.Message) (size int) {
33
	methods := protoMethods(m)
34
	if methods != nil && methods.Size != nil {
35
		out := methods.Size(protoiface.SizeInput{
36
			Message: m,
37
		})
38
		return out.Size
39
	}
40
	if methods != nil && methods.Marshal != nil {
41
		// This is not efficient, but we don't have any choice.
42
		// This case is mainly used for legacy types with a Marshal method.
43
		out, _ := methods.Marshal(protoiface.MarshalInput{
44
			Message: m,
45
		})
46
		return len(out.Buf)
47
	}
48
	return o.sizeMessageSlow(m)
49
}
50

51
func (o MarshalOptions) sizeMessageSlow(m protoreflect.Message) (size int) {
52
	if messageset.IsMessageSet(m.Descriptor()) {
53
		return o.sizeMessageSet(m)
54
	}
55
	m.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool {
56
		size += o.sizeField(fd, v)
57
		return true
58
	})
59
	size += len(m.GetUnknown())
60
	return size
61
}
62

63
func (o MarshalOptions) sizeField(fd protoreflect.FieldDescriptor, value protoreflect.Value) (size int) {
64
	num := fd.Number()
65
	switch {
66
	case fd.IsList():
67
		return o.sizeList(num, fd, value.List())
68
	case fd.IsMap():
69
		return o.sizeMap(num, fd, value.Map())
70
	default:
71
		return protowire.SizeTag(num) + o.sizeSingular(num, fd.Kind(), value)
72
	}
73
}
74

75
func (o MarshalOptions) sizeList(num protowire.Number, fd protoreflect.FieldDescriptor, list protoreflect.List) (size int) {
76
	if fd.IsPacked() && list.Len() > 0 {
77
		content := 0
78
		for i, llen := 0, list.Len(); i < llen; i++ {
79
			content += o.sizeSingular(num, fd.Kind(), list.Get(i))
80
		}
81
		return protowire.SizeTag(num) + protowire.SizeBytes(content)
82
	}
83

84
	for i, llen := 0, list.Len(); i < llen; i++ {
85
		size += protowire.SizeTag(num) + o.sizeSingular(num, fd.Kind(), list.Get(i))
86
	}
87
	return size
88
}
89

90
func (o MarshalOptions) sizeMap(num protowire.Number, fd protoreflect.FieldDescriptor, mapv protoreflect.Map) (size int) {
91
	mapv.Range(func(key protoreflect.MapKey, value protoreflect.Value) bool {
92
		size += protowire.SizeTag(num)
93
		size += protowire.SizeBytes(o.sizeField(fd.MapKey(), key.Value()) + o.sizeField(fd.MapValue(), value))
94
		return true
95
	})
96
	return size
97
}
98