blob: c84523f0a36c77f4fa9084cb54584ab3795b192e [file] [log] [blame]
Scott Baker4a35a702019-11-26 08:17:33 -08001package codec
2
3import "github.com/golang/protobuf/proto"
4
5// EncodeVarint writes a varint-encoded integer to the Buffer.
6// This is the format for the
7// int32, int64, uint32, uint64, bool, and enum
8// protocol buffer types.
9func (cb *Buffer) EncodeVarint(x uint64) error {
10 for x >= 1<<7 {
11 cb.buf = append(cb.buf, uint8(x&0x7f|0x80))
12 x >>= 7
13 }
14 cb.buf = append(cb.buf, uint8(x))
15 return nil
16}
17
18// EncodeTagAndWireType encodes the given field tag and wire type to the
19// buffer. This combines the two values and then writes them as a varint.
20func (cb *Buffer) EncodeTagAndWireType(tag int32, wireType int8) error {
21 v := uint64((int64(tag) << 3) | int64(wireType))
22 return cb.EncodeVarint(v)
23}
24
25// EncodeFixed64 writes a 64-bit integer to the Buffer.
26// This is the format for the
27// fixed64, sfixed64, and double protocol buffer types.
28func (cb *Buffer) EncodeFixed64(x uint64) error {
29 cb.buf = append(cb.buf,
30 uint8(x),
31 uint8(x>>8),
32 uint8(x>>16),
33 uint8(x>>24),
34 uint8(x>>32),
35 uint8(x>>40),
36 uint8(x>>48),
37 uint8(x>>56))
38 return nil
39}
40
41// EncodeFixed32 writes a 32-bit integer to the Buffer.
42// This is the format for the
43// fixed32, sfixed32, and float protocol buffer types.
44func (cb *Buffer) EncodeFixed32(x uint64) error {
45 cb.buf = append(cb.buf,
46 uint8(x),
47 uint8(x>>8),
48 uint8(x>>16),
49 uint8(x>>24))
50 return nil
51}
52
53// EncodeZigZag64 does zig-zag encoding to convert the given
54// signed 64-bit integer into a form that can be expressed
55// efficiently as a varint, even for negative values.
56func EncodeZigZag64(v int64) uint64 {
57 return (uint64(v) << 1) ^ uint64(v>>63)
58}
59
60// EncodeZigZag32 does zig-zag encoding to convert the given
61// signed 32-bit integer into a form that can be expressed
62// efficiently as a varint, even for negative values.
63func EncodeZigZag32(v int32) uint64 {
64 return uint64((uint32(v) << 1) ^ uint32((v >> 31)))
65}
66
67// EncodeRawBytes writes a count-delimited byte buffer to the Buffer.
68// This is the format used for the bytes protocol buffer
69// type and for embedded messages.
70func (cb *Buffer) EncodeRawBytes(b []byte) error {
71 if err := cb.EncodeVarint(uint64(len(b))); err != nil {
72 return err
73 }
74 cb.buf = append(cb.buf, b...)
75 return nil
76}
77
78// EncodeMessage writes the given message to the buffer.
79func (cb *Buffer) EncodeMessage(pm proto.Message) error {
80 bytes, err := marshalMessage(cb.buf, pm, cb.deterministic)
81 if err != nil {
82 return err
83 }
84 cb.buf = bytes
85 return nil
86}
87
88// EncodeDelimitedMessage writes the given message to the buffer with a
89// varint-encoded length prefix (the delimiter).
90func (cb *Buffer) EncodeDelimitedMessage(pm proto.Message) error {
91 bytes, err := marshalMessage(cb.tmp, pm, cb.deterministic)
92 if err != nil {
93 return err
94 }
95 // save truncated buffer if it was grown (so we can re-use it and
96 // curtail future allocations)
97 if cap(bytes) > cap(cb.tmp) {
98 cb.tmp = bytes[:0]
99 }
100 return cb.EncodeRawBytes(bytes)
101}
102
103func marshalMessage(b []byte, pm proto.Message, deterministic bool) ([]byte, error) {
104 // we try to use the most efficient way to marshal to existing slice
105 nm, ok := pm.(interface {
106 // this interface is implemented by generated messages
107 XXX_Size() int
108 XXX_Marshal(b []byte, deterministic bool) ([]byte, error)
109 })
110 if ok {
111 sz := nm.XXX_Size()
112 if cap(b) < len(b)+sz {
113 // re-allocate to fit
114 bytes := make([]byte, len(b), len(b)+sz)
115 copy(bytes, b)
116 b = bytes
117 }
118 return nm.XXX_Marshal(b, deterministic)
119 }
120
121 if deterministic {
122 // see if the message has custom deterministic methods, preferring an
123 // "append" method over one that must always re-allocate
124 madm, ok := pm.(interface {
125 MarshalAppendDeterministic(b []byte) ([]byte, error)
126 })
127 if ok {
128 return madm.MarshalAppendDeterministic(b)
129 }
130
131 mdm, ok := pm.(interface {
132 MarshalDeterministic() ([]byte, error)
133 })
134 if ok {
135 bytes, err := mdm.MarshalDeterministic()
136 if err != nil {
137 return nil, err
138 }
139 if len(b) == 0 {
140 return bytes, nil
141 }
142 return append(b, bytes...), nil
143 }
144 }
145
146 mam, ok := pm.(interface {
147 // see if we can append the message, vs. having to re-allocate
148 MarshalAppend(b []byte) ([]byte, error)
149 })
150 if ok {
151 return mam.MarshalAppend(b)
152 }
153
154 // lowest common denominator
155 bytes, err := proto.Marshal(pm)
156 if err != nil {
157 return nil, err
158 }
159 if len(b) == 0 {
160 return bytes, nil
161 }
162 return append(b, bytes...), nil
163}