github.com/hoveychen/protoreflect@v1.4.7-0.20221103114119-0b4b3385ec76/codec/encode.go (about)

     1  package codec
     2  
     3  import "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.
     9  func (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.
    20  func (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.
    28  func (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.
    44  func (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.
    56  func 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.
    63  func 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.
    70  func (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.
    79  func (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).
    90  func (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  
   103  func 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  }