xref: /aosp_15_r20/external/pigweed/pw_protobuf/public/pw_protobuf/encoder.h (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1 // Copyright 2021 The Pigweed Authors
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License"); you may not
4 // use this file except in compliance with the License. You may obtain a copy of
5 // the License at
6 //
7 //     https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11 // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12 // License for the specific language governing permissions and limitations under
13 // the License.
14 #pragma once
15 
16 #include <algorithm>
17 #include <array>
18 #include <cstddef>
19 #include <cstring>
20 #include <string_view>
21 
22 #include "pw_assert/assert.h"
23 #include "pw_bytes/bit.h"
24 #include "pw_bytes/endian.h"
25 #include "pw_bytes/span.h"
26 #include "pw_containers/vector.h"
27 #include "pw_protobuf/config.h"
28 #include "pw_protobuf/internal/codegen.h"
29 #include "pw_protobuf/wire_format.h"
30 #include "pw_span/span.h"
31 #include "pw_status/status.h"
32 #include "pw_status/try.h"
33 #include "pw_stream/memory_stream.h"
34 #include "pw_stream/stream.h"
35 #include "pw_toolchain/internal/sibling_cast.h"
36 #include "pw_varint/varint.h"
37 
38 namespace pw::protobuf {
39 
40 // Provides a size estimate to help with sizing buffers passed to
41 // StreamEncoder and MemoryEncoder objects.
42 //
43 // Args:
44 //   max_message_size: For MemoryEncoder objects, this is the max expected size
45 //     of the final proto. For StreamEncoder objects, this should be the max
46 //     size of any nested proto submessage that will be built with this encoder
47 //     (recursively accumulating the size from the root submessage). If your
48 //     proto will encode many large submessages, this value should just be the
49 //     size of the largest one.
50 //  max_nested_depth: The max number of nested submessage encoders that are
51 //     expected to be open simultaneously to encode this proto message.
MaxScratchBufferSize(size_t max_message_size,size_t max_nested_depth)52 constexpr size_t MaxScratchBufferSize(size_t max_message_size,
53                                       size_t max_nested_depth) {
54   return max_message_size + max_nested_depth * config::kMaxVarintSize;
55 }
56 
57 // Write a varint value to the writer.
58 //
59 // Args:
60 //   value: The value of the varint to write
61 //   writer: The writer for writing to output.
62 //
63 // Returns:
64 // OK - varint is written successfully
65 //
66 // Errors encountered by the `writer` will be returned as it is.
WriteVarint(uint64_t value,stream::Writer & writer)67 inline Status WriteVarint(uint64_t value, stream::Writer& writer) {
68   std::array<std::byte, varint::kMaxVarint64SizeBytes> varint_encode_buffer;
69   const size_t varint_size =
70       pw::varint::EncodeLittleEndianBase128(value, varint_encode_buffer);
71   return writer.Write(span(varint_encode_buffer).first(varint_size));
72 }
73 
74 // Write the field key and length prefix for a length-delimited field. It is
75 // up to the caller to ensure that this will be followed by an exact number
76 // of bytes written for the field in order to form a valid proto message.
77 //
78 // Args:
79 //   field_number: The field number for the field.
80 //   payload_size: The size of the payload.
81 //   writer: The output writer to write to
82 //
83 //
84 // Returns:
85 // OK - Field key is written successfully
86 //
87 // Errors encountered by the `writer` will be returned as it is.
88 //
89 // Precondition: The field_number must be a ValidFieldNumber.
90 // Precondition: `data_size_bytes` must be smaller than
91 //   std::numeric_limits<uint32_t>::max()
WriteLengthDelimitedKeyAndLengthPrefix(uint32_t field_number,size_t payload_size,stream::Writer & writer)92 inline Status WriteLengthDelimitedKeyAndLengthPrefix(uint32_t field_number,
93                                                      size_t payload_size,
94                                                      stream::Writer& writer) {
95   PW_TRY(WriteVarint(FieldKey(field_number, WireType::kDelimited), writer));
96   return WriteVarint(payload_size, writer);
97 }
98 
99 // Forward declaration. StreamEncoder and MemoryEncoder are very tightly
100 // coupled.
101 class MemoryEncoder;
102 
103 // A protobuf encoder that encodes serialized proto data to a
104 // pw::stream::Writer.
105 class StreamEncoder {
106  public:
107   // The StreamEncoder will serialize proto data to the pw::stream::Writer
108   // provided through the constructor. The scratch buffer provided is for
109   // internal use ONLY and should not be considered valid proto data.
110   //
111   // If a StreamEncoder object will be writing nested proto messages, it must
112   // provide a scratch buffer large enough to hold the largest submessage some
113   // additional overhead incurred by the encoder's implementation. It's a good
114   // idea to be generous when sizing this buffer. MaxScratchBufferSize() can be
115   // helpful in providing an estimated size for this buffer. The scratch buffer
116   // must exist for the lifetime of the StreamEncoder object.
117   //
118   // StreamEncoder objects that do not write nested proto messages can
119   // provide a zero-length scratch buffer.
StreamEncoder(stream::Writer & writer,ByteSpan scratch_buffer)120   constexpr StreamEncoder(stream::Writer& writer, ByteSpan scratch_buffer)
121       : status_(OkStatus()),
122         write_when_empty_(true),
123         parent_(nullptr),
124         nested_field_number_(0),
125         memory_writer_(scratch_buffer),
126         writer_(writer) {}
127 
128   // Precondition: Encoder has no active child encoder.
129   //
130   // Postcondition: If this encoder is a nested one, the parent encoder is
131   //     unlocked and proto encoding may resume on the parent.
~StreamEncoder()132   ~StreamEncoder() { CloseEncoder(); }
133 
134   // Disallow copy/assign to avoid confusion about who owns the buffer.
135   StreamEncoder& operator=(const StreamEncoder& other) = delete;
136   StreamEncoder(const StreamEncoder& other) = delete;
137 
138   // It's not safe to move an encoder as it could cause another encoder's
139   // parent_ pointer to become invalid.
140   StreamEncoder& operator=(StreamEncoder&& other) = delete;
141 
142   // Closes this encoder, finalizing its output.
143   //
144   // This method is called automatically by `StreamEncoder`'s destructor, but
145   // may be invoked manually in order to close an encoder before the end of its
146   // lexical scope.
147   //
148   // Precondition: Encoder has no active child encoder.
149   //
150   // Postcondition: If this encoder is a nested one, the parent encoder is
151   //     unlocked and proto encoding may resume on the parent. No more writes
152   //     to this encoder may be performed.
153   void CloseEncoder();
154 
155   // Forwards the conservative write limit of the underlying
156   // pw::stream::Writer.
157   //
158   // Precondition: Encoder has no active child encoder.
ConservativeWriteLimit()159   size_t ConservativeWriteLimit() const {
160     PW_ASSERT(!nested_encoder_open());
161     return writer_.ConservativeWriteLimit();
162   }
163 
164   enum class EmptyEncoderBehavior { kWriteFieldNumber, kWriteNothing };
165 
166   // Creates a nested encoder with the provided field number. Once this is
167   // called, the parent encoder is locked and not available for use until the
168   // nested encoder is finalized (either explicitly or through destruction).
169   //
170   // Precondition: Encoder has no active child encoder.
171   //
172   // Postcondition: Until the nested child encoder has been destroyed, this
173   //     encoder cannot be used.
174   StreamEncoder GetNestedEncoder(uint32_t field_number,
175                                  EmptyEncoderBehavior empty_encoder_behavior =
176                                      EmptyEncoderBehavior::kWriteFieldNumber) {
177     return GetNestedEncoder(
178         field_number, /*write_when_empty=*/
179         empty_encoder_behavior == EmptyEncoderBehavior::kWriteFieldNumber);
180   }
181 
182   // Returns the current encoder's status.
183   //
184   // Precondition: Encoder has no active child encoder.
status()185   Status status() const {
186     PW_ASSERT(!nested_encoder_open());
187     return status_;
188   }
189 
190   // Writes a proto uint32 key-value pair.
191   //
192   // Precondition: Encoder has no active child encoder.
WriteUint32(uint32_t field_number,uint32_t value)193   Status WriteUint32(uint32_t field_number, uint32_t value) {
194     return WriteUint64(field_number, value);
195   }
196 
197   // Writes a repeated uint32 using packed encoding.
198   //
199   // Precondition: Encoder has no active child encoder.
WritePackedUint32(uint32_t field_number,span<const uint32_t> values)200   Status WritePackedUint32(uint32_t field_number, span<const uint32_t> values) {
201     return WritePackedVarints(
202         field_number, values, internal::VarintType::kNormal);
203   }
204 
205   // Writes a repeated uint32 using packed encoding.
206   //
207   // Precondition: Encoder has no active child encoder.
WriteRepeatedUint32(uint32_t field_number,const pw::Vector<uint32_t> & values)208   Status WriteRepeatedUint32(uint32_t field_number,
209                              const pw::Vector<uint32_t>& values) {
210     return WritePackedVarints(field_number,
211                               span(values.data(), values.size()),
212                               internal::VarintType::kNormal);
213   }
214 
215   // Writes a proto uint64 key-value pair.
216   //
217   // Precondition: Encoder has no active child encoder.
WriteUint64(uint32_t field_number,uint64_t value)218   Status WriteUint64(uint32_t field_number, uint64_t value) {
219     return WriteVarintField(field_number, value);
220   }
221 
222   // Writes a repeated uint64 using packed encoding.
223   //
224   // Precondition: Encoder has no active child encoder.
WritePackedUint64(uint32_t field_number,span<const uint64_t> values)225   Status WritePackedUint64(uint32_t field_number, span<const uint64_t> values) {
226     return WritePackedVarints(
227         field_number, values, internal::VarintType::kNormal);
228   }
229 
230   // Writes a repeated uint64 using packed encoding.
231   //
232   // Precondition: Encoder has no active child encoder.
WriteRepeatedUint64(uint32_t field_number,const pw::Vector<uint64_t> & values)233   Status WriteRepeatedUint64(uint32_t field_number,
234                              const pw::Vector<uint64_t>& values) {
235     return WritePackedVarints(field_number,
236                               span(values.data(), values.size()),
237                               internal::VarintType::kNormal);
238   }
239 
240   // Writes a proto int32 key-value pair.
241   //
242   // Precondition: Encoder has no active child encoder.
WriteInt32(uint32_t field_number,int32_t value)243   Status WriteInt32(uint32_t field_number, int32_t value) {
244     // Signed numbers are sent as 2's complement so this cast is correct.
245     return WriteUint64(field_number, static_cast<uint64_t>(value));
246   }
247 
248   // Writes a repeated int32 using packed encoding.
249   //
250   // Precondition: Encoder has no active child encoder.
WritePackedInt32(uint32_t field_number,span<const int32_t> values)251   Status WritePackedInt32(uint32_t field_number, span<const int32_t> values) {
252     return WritePackedVarints(
253         field_number,
254         span(reinterpret_cast<const uint32_t*>(values.data()), values.size()),
255         internal::VarintType::kNormal);
256   }
257 
258   // Writes a repeated int32 using packed encoding.
259   //
260   // Precondition: Encoder has no active child encoder.
WriteRepeatedInt32(uint32_t field_number,const pw::Vector<int32_t> & values)261   Status WriteRepeatedInt32(uint32_t field_number,
262                             const pw::Vector<int32_t>& values) {
263     return WritePackedVarints(
264         field_number,
265         span(reinterpret_cast<const uint32_t*>(values.data()), values.size()),
266         internal::VarintType::kNormal);
267   }
268 
269   // Writes a proto int64 key-value pair.
270   //
271   // Precondition: Encoder has no active child encoder.
WriteInt64(uint32_t field_number,int64_t value)272   Status WriteInt64(uint32_t field_number, int64_t value) {
273     // Signed numbers are sent as 2's complement so this cast is correct.
274     return WriteUint64(field_number, static_cast<uint64_t>(value));
275   }
276 
277   // Writes a repeated int64 using packed encoding.
278   //
279   // Precondition: Encoder has no active child encoder.
WritePackedInt64(uint32_t field_number,span<const int64_t> values)280   Status WritePackedInt64(uint32_t field_number, span<const int64_t> values) {
281     return WritePackedVarints(
282         field_number,
283         span(reinterpret_cast<const uint64_t*>(values.data()), values.size()),
284         internal::VarintType::kNormal);
285   }
286 
287   // Writes a repeated int64 using packed encoding.
288   //
289   // Precondition: Encoder has no active child encoder.
WriteRepeatedInt64(uint32_t field_number,const pw::Vector<int64_t> & values)290   Status WriteRepeatedInt64(uint32_t field_number,
291                             const pw::Vector<int64_t>& values) {
292     return WritePackedVarints(
293         field_number,
294         span(reinterpret_cast<const uint64_t*>(values.data()), values.size()),
295         internal::VarintType::kNormal);
296   }
297 
298   // Writes a proto sint32 key-value pair.
299   //
300   // Precondition: Encoder has no active child encoder.
WriteSint32(uint32_t field_number,int32_t value)301   Status WriteSint32(uint32_t field_number, int32_t value) {
302     return WriteUint64(field_number, varint::ZigZagEncode(value));
303   }
304 
305   // Writes a repeated sint32 using packed encoding.
306   //
307   // Precondition: Encoder has no active child encoder.
WritePackedSint32(uint32_t field_number,span<const int32_t> values)308   Status WritePackedSint32(uint32_t field_number, span<const int32_t> values) {
309     return WritePackedVarints(
310         field_number,
311         span(reinterpret_cast<const uint32_t*>(values.data()), values.size()),
312         internal::VarintType::kZigZag);
313   }
314 
315   // Writes a repeated sint32 using packed encoding.
316   //
317   // Precondition: Encoder has no active child encoder.
WriteRepeatedSint32(uint32_t field_number,const pw::Vector<int32_t> & values)318   Status WriteRepeatedSint32(uint32_t field_number,
319                              const pw::Vector<int32_t>& values) {
320     return WritePackedVarints(
321         field_number,
322         span(reinterpret_cast<const uint32_t*>(values.data()), values.size()),
323         internal::VarintType::kZigZag);
324   }
325 
326   // Writes a proto sint64 key-value pair.
327   //
328   // Precondition: Encoder has no active child encoder.
WriteSint64(uint32_t field_number,int64_t value)329   Status WriteSint64(uint32_t field_number, int64_t value) {
330     return WriteUint64(field_number, varint::ZigZagEncode(value));
331   }
332 
333   // Writes a repeated sint64 using packed encoding.
334   //
335   // Precondition: Encoder has no active child encoder.
WritePackedSint64(uint32_t field_number,span<const int64_t> values)336   Status WritePackedSint64(uint32_t field_number, span<const int64_t> values) {
337     return WritePackedVarints(
338         field_number,
339         span(reinterpret_cast<const uint64_t*>(values.data()), values.size()),
340         internal::VarintType::kZigZag);
341   }
342 
343   // Writes a repeated sint64 using packed encoding.
344   //
345   // Precondition: Encoder has no active child encoder.
WriteRepeatedSint64(uint32_t field_number,const pw::Vector<int64_t> & values)346   Status WriteRepeatedSint64(uint32_t field_number,
347                              const pw::Vector<int64_t>& values) {
348     return WritePackedVarints(
349         field_number,
350         span(reinterpret_cast<const uint64_t*>(values.data()), values.size()),
351         internal::VarintType::kZigZag);
352   }
353 
354   // Writes a proto bool key-value pair.
355   //
356   // Precondition: Encoder has no active child encoder.
WriteBool(uint32_t field_number,bool value)357   Status WriteBool(uint32_t field_number, bool value) {
358     return WriteUint32(field_number, static_cast<uint32_t>(value));
359   }
360 
361   // Writes a repeated bool using packed encoding.
362   //
363   // Precondition: Encoder has no active child encoder.
WritePackedBool(uint32_t field_number,span<const bool> values)364   Status WritePackedBool(uint32_t field_number, span<const bool> values) {
365     static_assert(sizeof(bool) == sizeof(uint8_t),
366                   "bool must be same size as uint8_t");
367     return WritePackedVarints(
368         field_number,
369         span(reinterpret_cast<const uint8_t*>(values.data()), values.size()),
370         internal::VarintType::kNormal);
371   }
372 
373   // Writes a repeated bool using packed encoding.
374   //
375   // Precondition: Encoder has no active child encoder.
WriteRepeatedBool(uint32_t field_number,const pw::Vector<bool> & values)376   Status WriteRepeatedBool(uint32_t field_number,
377                            const pw::Vector<bool>& values) {
378     static_assert(sizeof(bool) == sizeof(uint8_t),
379                   "bool must be same size as uint8_t");
380 
381     return WritePackedVarints(
382         field_number,
383         span(reinterpret_cast<const uint8_t*>(values.data()), values.size()),
384         internal::VarintType::kNormal);
385   }
386 
387   // Writes a proto fixed32 key-value pair.
388   //
389   // Precondition: Encoder has no active child encoder.
WriteFixed32(uint32_t field_number,uint32_t value)390   Status WriteFixed32(uint32_t field_number, uint32_t value) {
391     std::array<std::byte, sizeof(value)> data =
392         bytes::CopyInOrder(endian::little, value);
393     return WriteFixed(field_number, data);
394   }
395 
396   // Writes a repeated fixed32 field using packed encoding.
397   //
398   // Precondition: Encoder has no active child encoder.
WritePackedFixed32(uint32_t field_number,span<const uint32_t> values)399   Status WritePackedFixed32(uint32_t field_number,
400                             span<const uint32_t> values) {
401     return WritePackedFixed(field_number, as_bytes(values), sizeof(uint32_t));
402   }
403 
404   // Writes a repeated fixed32 field using packed encoding.
405   //
406   // Precondition: Encoder has no active child encoder.
WriteRepeatedFixed32(uint32_t field_number,const pw::Vector<uint32_t> & values)407   Status WriteRepeatedFixed32(uint32_t field_number,
408                               const pw::Vector<uint32_t>& values) {
409     return WritePackedFixed(field_number,
410                             as_bytes(span(values.data(), values.size())),
411                             sizeof(uint32_t));
412   }
413 
414   // Writes a proto fixed64 key-value pair.
415   //
416   // Precondition: Encoder has no active child encoder.
WriteFixed64(uint32_t field_number,uint64_t value)417   Status WriteFixed64(uint32_t field_number, uint64_t value) {
418     std::array<std::byte, sizeof(value)> data =
419         bytes::CopyInOrder(endian::little, value);
420     return WriteFixed(field_number, data);
421   }
422 
423   // Writes a repeated fixed64 field using packed encoding.
424   //
425   // Precondition: Encoder has no active child encoder.
WritePackedFixed64(uint32_t field_number,span<const uint64_t> values)426   Status WritePackedFixed64(uint32_t field_number,
427                             span<const uint64_t> values) {
428     return WritePackedFixed(field_number, as_bytes(values), sizeof(uint64_t));
429   }
430 
431   // Writes a repeated fixed64 field using packed encoding.
432   //
433   // Precondition: Encoder has no active child encoder.
WriteRepeatedFixed64(uint32_t field_number,const pw::Vector<uint64_t> & values)434   Status WriteRepeatedFixed64(uint32_t field_number,
435                               const pw::Vector<uint64_t>& values) {
436     return WritePackedFixed(field_number,
437                             as_bytes(span(values.data(), values.size())),
438                             sizeof(uint64_t));
439   }
440 
441   // Writes a proto sfixed32 key-value pair.
442   //
443   // Precondition: Encoder has no active child encoder.
WriteSfixed32(uint32_t field_number,int32_t value)444   Status WriteSfixed32(uint32_t field_number, int32_t value) {
445     return WriteFixed32(field_number, static_cast<uint32_t>(value));
446   }
447 
448   // Writes a repeated sfixed32 field using packed encoding.
449   //
450   // Precondition: Encoder has no active child encoder.
WritePackedSfixed32(uint32_t field_number,span<const int32_t> values)451   Status WritePackedSfixed32(uint32_t field_number,
452                              span<const int32_t> values) {
453     return WritePackedFixed(field_number, as_bytes(values), sizeof(int32_t));
454   }
455 
456   // Writes a repeated fixed32 field using packed encoding.
457   //
458   // Precondition: Encoder has no active child encoder.
WriteRepeatedSfixed32(uint32_t field_number,const pw::Vector<int32_t> & values)459   Status WriteRepeatedSfixed32(uint32_t field_number,
460                                const pw::Vector<int32_t>& values) {
461     return WritePackedFixed(field_number,
462                             as_bytes(span(values.data(), values.size())),
463                             sizeof(int32_t));
464   }
465 
466   // Writes a proto sfixed64 key-value pair.
467   //
468   // Precondition: Encoder has no active child encoder.
WriteSfixed64(uint32_t field_number,int64_t value)469   Status WriteSfixed64(uint32_t field_number, int64_t value) {
470     return WriteFixed64(field_number, static_cast<uint64_t>(value));
471   }
472 
473   // Writes a repeated sfixed64 field using packed encoding.
474   //
475   // Precondition: Encoder has no active child encoder.
WritePackedSfixed64(uint32_t field_number,span<const int64_t> values)476   Status WritePackedSfixed64(uint32_t field_number,
477                              span<const int64_t> values) {
478     return WritePackedFixed(field_number, as_bytes(values), sizeof(int64_t));
479   }
480 
481   // Writes a repeated fixed64 field using packed encoding.
482   //
483   // Precondition: Encoder has no active child encoder.
WriteRepeatedSfixed64(uint32_t field_number,const pw::Vector<int64_t> & values)484   Status WriteRepeatedSfixed64(uint32_t field_number,
485                                const pw::Vector<int64_t>& values) {
486     return WritePackedFixed(field_number,
487                             as_bytes(span(values.data(), values.size())),
488                             sizeof(int64_t));
489   }
490 
491   // Writes a proto float key-value pair.
492   //
493   // Precondition: Encoder has no active child encoder.
WriteFloat(uint32_t field_number,float value)494   Status WriteFloat(uint32_t field_number, float value) {
495     static_assert(sizeof(float) == sizeof(uint32_t),
496                   "Float and uint32_t are not the same size");
497     uint32_t integral_value;
498     std::memcpy(&integral_value, &value, sizeof(value));
499     std::array<std::byte, sizeof(value)> data =
500         bytes::CopyInOrder(endian::little, integral_value);
501     return WriteFixed(field_number, data);
502   }
503 
504   // Writes a repeated float field using packed encoding.
505   //
506   // Precondition: Encoder has no active child encoder.
WritePackedFloat(uint32_t field_number,span<const float> values)507   Status WritePackedFloat(uint32_t field_number, span<const float> values) {
508     return WritePackedFixed(field_number, as_bytes(values), sizeof(float));
509   }
510 
511   // Writes a repeated float field using packed encoding.
512   //
513   // Precondition: Encoder has no active child encoder.
WriteRepeatedFloat(uint32_t field_number,const pw::Vector<float> & values)514   Status WriteRepeatedFloat(uint32_t field_number,
515                             const pw::Vector<float>& values) {
516     return WritePackedFixed(field_number,
517                             as_bytes(span(values.data(), values.size())),
518                             sizeof(float));
519   }
520 
521   // Writes a proto double key-value pair.
522   //
523   // Precondition: Encoder has no active child encoder.
WriteDouble(uint32_t field_number,double value)524   Status WriteDouble(uint32_t field_number, double value) {
525     static_assert(sizeof(double) == sizeof(uint64_t),
526                   "Double and uint64_t are not the same size");
527     uint64_t integral_value;
528     std::memcpy(&integral_value, &value, sizeof(value));
529     std::array<std::byte, sizeof(value)> data =
530         bytes::CopyInOrder(endian::little, integral_value);
531     return WriteFixed(field_number, data);
532   }
533 
534   // Writes a repeated double field using packed encoding.
535   //
536   // Precondition: Encoder has no active child encoder.
WritePackedDouble(uint32_t field_number,span<const double> values)537   Status WritePackedDouble(uint32_t field_number, span<const double> values) {
538     return WritePackedFixed(field_number, as_bytes(values), sizeof(double));
539   }
540 
541   // Writes a repeated double field using packed encoding.
542   //
543   // Precondition: Encoder has no active child encoder.
WriteRepeatedDouble(uint32_t field_number,const pw::Vector<double> & values)544   Status WriteRepeatedDouble(uint32_t field_number,
545                              const pw::Vector<double>& values) {
546     return WritePackedFixed(field_number,
547                             as_bytes(span(values.data(), values.size())),
548                             sizeof(double));
549   }
550 
551   // Writes a proto `bytes` field as a key-value pair. This can also be used to
552   // write a pre-encoded nested submessage directly without using a nested
553   // encoder.
554   //
555   // Precondition: Encoder has no active child encoder.
WriteBytes(uint32_t field_number,ConstByteSpan value)556   Status WriteBytes(uint32_t field_number, ConstByteSpan value) {
557     return WriteLengthDelimitedField(field_number, value);
558   }
559 
560   // Writes a proto 'bytes' field from the stream bytes_reader.
561   //
562   // The payload for the value is provided through the stream::Reader
563   // `bytes_reader`. The method reads a chunk of the data from the reader using
564   // the `stream_pipe_buffer` and writes it to the encoder.
565   //
566   // Precondition: The stream_pipe_buffer.byte_size() >= 1
567   // Precondition: Encoder has no active child encoder.
568   //
569   // Returns:
570   // OK - Bytes field is written successfully.
571   // RESOURCE_EXHAUSTED - Exceeds write limits.
572   // OUT_OF_RANGE - `bytes_reader` is exhausted before `num_bytes` of
573   //                bytes is read.
574   //
575   // Other errors encountered by the writer will be returned as it is.
WriteBytesFromStream(uint32_t field_number,stream::Reader & bytes_reader,size_t num_bytes,ByteSpan stream_pipe_buffer)576   Status WriteBytesFromStream(uint32_t field_number,
577                               stream::Reader& bytes_reader,
578                               size_t num_bytes,
579                               ByteSpan stream_pipe_buffer) {
580     return WriteLengthDelimitedFieldFromStream(
581         field_number, bytes_reader, num_bytes, stream_pipe_buffer);
582   }
583 
584   // Writes a proto string key-value pair.
585   //
586   // Precondition: Encoder has no active child encoder.
WriteString(uint32_t field_number,std::string_view value)587   Status WriteString(uint32_t field_number, std::string_view value) {
588     return WriteBytes(field_number, as_bytes(span<const char>(value)));
589   }
590 
591   // Writes a proto string key-value pair.
592   //
593   // Precondition: Encoder has no active child encoder.
WriteString(uint32_t field_number,const char * value,size_t len)594   Status WriteString(uint32_t field_number, const char* value, size_t len) {
595     return WriteBytes(field_number, as_bytes(span(value, len)));
596   }
597 
598   // Writes a proto 'string' field from the stream bytes_reader.
599   //
600   // The payload for the value is provided through the stream::Reader
601   // `bytes_reader`. The method reads a chunk of the data from the reader using
602   // the `stream_pipe_buffer` and writes it to the encoder.
603   //
604   // Precondition: The stream_pipe_buffer.byte_size() >= 1
605   // Precondition: Encoder has no active child encoder.
606   //
607   // Returns:
608   // OK - String field is written successfully.
609   // RESOURCE_EXHAUSTED - Exceeds write limits.
610   // OUT_OF_RANGE - `bytes_reader` is exhausted before `num_bytes` of
611   //                bytes is read.
612   //
613   // Other errors encountered by the writer will be returned as it is.
WriteStringFromStream(uint32_t field_number,stream::Reader & bytes_reader,size_t num_bytes,ByteSpan stream_pipe_buffer)614   Status WriteStringFromStream(uint32_t field_number,
615                                stream::Reader& bytes_reader,
616                                size_t num_bytes,
617                                ByteSpan stream_pipe_buffer) {
618     return WriteBytesFromStream(
619         field_number, bytes_reader, num_bytes, stream_pipe_buffer);
620   }
621 
622  protected:
623   // Specialized move constructor used only for codegen.
624   //
625   // Postcondition: The other encoder is invalidated and cannot be used as it
626   //     acts like a parent encoder with an active child encoder.
StreamEncoder(StreamEncoder && other)627   constexpr StreamEncoder(StreamEncoder&& other)
628       : status_(other.status_),
629         write_when_empty_(true),
630         parent_(other.parent_),
631         nested_field_number_(other.nested_field_number_),
632         memory_writer_(std::move(other.memory_writer_)),
633         writer_(&other.writer_ == &other.memory_writer_ ? memory_writer_
634                                                         : other.writer_) {
635     PW_ASSERT(nested_field_number_ == 0);
636     // Make the nested encoder look like it has an open child to block writes
637     // for the remainder of the object's life.
638     other.nested_field_number_ = kFirstReservedNumber;
639     other.parent_ = nullptr;
640   }
641 
642   // Writes proto values to the stream from the structure contained within
643   // message, according to the description of fields in table.
644   //
645   // This is called by codegen subclass Write() functions that accept a typed
646   // struct Message reference, using the appropriate codegen MessageField table
647   // corresponding to that type.
648   Status Write(span<const std::byte> message,
649                span<const internal::MessageField> table);
650 
651   // Protected method to create a nested encoder, specifying whether the field
652   // should be written when no fields were added to the nested encoder. Exposed
653   // using an enum in the public API, for better readability.
654   StreamEncoder GetNestedEncoder(uint32_t field_number, bool write_when_empty);
655 
656  private:
657   friend class MemoryEncoder;
658 
659   constexpr StreamEncoder(StreamEncoder& parent,
660                           ByteSpan scratch_buffer,
661                           bool write_when_empty = true)
status_(OkStatus ())662       : status_(OkStatus()),
663         write_when_empty_(write_when_empty),
664         parent_(&parent),
665         nested_field_number_(0),
666         memory_writer_(scratch_buffer),
667         writer_(memory_writer_) {
668     // If this encoder was spawned from a failed encoder, it should also start
669     // in a failed state.
670     if (&parent != this) {
671       status_.Update(parent.status_);
672     }
673     if (scratch_buffer.empty()) {
674       status_.Update(Status::ResourceExhausted());
675     }
676   }
677 
nested_encoder_open()678   bool nested_encoder_open() const { return nested_field_number_ != 0; }
679 
680   // CloseNestedMessage() is called on the parent encoder as part of the nested
681   // encoder destructor.
682   void CloseNestedMessage(StreamEncoder& nested);
683 
684   // Implementation for encoding all varint field types.
685   Status WriteVarintField(uint32_t field_number, uint64_t value);
686 
687   // Implementation for encoding all length-delimited field types.
688   Status WriteLengthDelimitedField(uint32_t field_number, ConstByteSpan data);
689 
690   // Encoding of length-delimited field where payload comes from `bytes_reader`.
691   Status WriteLengthDelimitedFieldFromStream(uint32_t field_number,
692                                              stream::Reader& bytes_reader,
693                                              size_t num_bytes,
694                                              ByteSpan stream_pipe_buffer);
695 
696   // Implementation for encoding all fixed-length integer types.
697   Status WriteFixed(uint32_t field_number, ConstByteSpan data);
698 
699   // Encodes a base-128 varint to the buffer. This function assumes the caller
700   // has already checked UpdateStatusForWrite() to ensure the writer's
701   // conservative write limit indicates the Writer has sufficient buffer space.
WriteVarint(uint64_t value)702   Status WriteVarint(uint64_t value) {
703     PW_TRY(status_);
704     status_.Update(::pw::protobuf::WriteVarint(value, writer_));
705     return status_;
706   }
707 
WriteZigzagVarint(int64_t value)708   Status WriteZigzagVarint(int64_t value) {
709     return WriteVarint(varint::ZigZagEncode(value));
710   }
711 
712   // Writes a list of varints to the buffer in length-delimited packed encoding.
713   template <typename T, typename = std::enable_if_t<std::is_integral<T>::value>>
WritePackedVarints(uint32_t field_number,span<T> values,internal::VarintType encode_type)714   Status WritePackedVarints(uint32_t field_number,
715                             span<T> values,
716                             internal::VarintType encode_type) {
717     static_assert(std::is_same<T, const uint8_t>::value ||
718                       std::is_same<T, const uint32_t>::value ||
719                       std::is_same<T, const int32_t>::value ||
720                       std::is_same<T, const uint64_t>::value ||
721                       std::is_same<T, const int64_t>::value,
722                   "Packed varints must be of type bool, uint32_t, int32_t, "
723                   "uint64_t, or int64_t");
724 
725     size_t payload_size = 0;
726     for (T val : values) {
727       if (encode_type == internal::VarintType::kZigZag) {
728         int64_t integer =
729             static_cast<int64_t>(static_cast<std::make_signed_t<T>>(val));
730         payload_size += varint::EncodedSize(varint::ZigZagEncode(integer));
731       } else {
732         uint64_t integer = static_cast<uint64_t>(val);
733         payload_size += varint::EncodedSize(integer);
734       }
735     }
736 
737     if (!UpdateStatusForWrite(field_number, WireType::kDelimited, payload_size)
738              .ok()) {
739       return status_;
740     }
741 
742     WriteVarint(FieldKey(field_number, WireType::kDelimited))
743         .IgnoreError();  // TODO: b/242598609 - Handle Status properly
744     WriteVarint(payload_size)
745         .IgnoreError();  // TODO: b/242598609 - Handle Status properly
746     for (T value : values) {
747       if (encode_type == internal::VarintType::kZigZag) {
748         WriteZigzagVarint(static_cast<std::make_signed_t<T>>(value))
749             .IgnoreError();  // TODO: b/242598609 - Handle Status properly
750       } else {
751         WriteVarint(value)
752             .IgnoreError();  // TODO: b/242598609 - Handle Status properly
753       }
754     }
755 
756     return status_;
757   }
758 
759   // Writes a list of fixed-size types to the buffer in length-delimited
760   // packed encoding. Only float, double, uint32_t, int32_t, uint64_t, and
761   // int64_t are permitted
762   Status WritePackedFixed(uint32_t field_number,
763                           span<const std::byte> values,
764                           size_t elem_size);
765 
766   template <typename Container>
WriteStringOrBytes(uint32_t field_number,const std::byte * raw_container)767   Status WriteStringOrBytes(uint32_t field_number,
768                             const std::byte* raw_container) {
769     const auto& container = *reinterpret_cast<Container*>(raw_container);
770     if (container.empty()) {
771       return OkStatus();
772     }
773     return WriteLengthDelimitedField(field_number, as_bytes(span(container)));
774   }
775 
776   // Checks if a write is invalid or will cause the encoder to enter an error
777   // state, and preemptively sets this encoder's status to that error to block
778   // the write. Only the first error encountered is tracked.
779   //
780   // Precondition: Encoder has no active child encoder.
781   //
782   // Returns:
783   //   InvalidArgument: The field number provided was invalid.
784   //   ResourceExhausted: The requested write would have exceeded the
785   //     stream::Writer's conservative write limit.
786   //   Other: If any Write() operations on the stream::Writer caused an error,
787   //     that error will be repeated here.
788   Status UpdateStatusForWrite(uint32_t field_number,
789                               WireType type,
790                               size_t data_size);
791 
792   // Callbacks for oneof fields set a flag to ensure they are only invoked once.
793   // To maintain logical constness of message structs passed to write, this
794   // resets each callback's invoked flag following a write operation.
795   void ResetOneOfCallbacks(ConstByteSpan message,
796                            span<const internal::MessageField> table);
797 
798   // The current encoder status. This status is only updated to reflect the
799   // first error encountered. Any further write operations are blocked when the
800   // encoder enters an error state.
801   Status status_;
802 
803   // Checked by the parent when the nested encoder is closed, and if no bytes
804   // were written, the field is not written.
805   bool write_when_empty_;
806 
807   // If this is a nested encoder, this points to the encoder that created it.
808   // For user-created MemoryEncoders, parent_ points to this object as an
809   // optimization for the MemoryEncoder and nested encoders to use the same
810   // underlying buffer.
811   StreamEncoder* parent_;
812 
813   // If an encoder has a child encoder open, this is the field number of that
814   // submessage. Otherwise, this is 0 to indicate no child encoder is open.
815   uint32_t nested_field_number_;
816 
817   // This memory writer is used for staging proto submessages to the
818   // scratch_buffer.
819   stream::MemoryWriter memory_writer_;
820 
821   // All proto encode operations are directly written to this writer.
822   stream::Writer& writer_;
823 };
824 
825 // A protobuf encoder that writes directly to a provided buffer.
826 //
827 // Example:
828 //
829 //   // Writes a proto response to the provided buffer, returning the encode
830 //   // status and number of bytes written.
831 //   StatusWithSize WriteProtoResponse(ByteSpan response) {
832 //     // All proto writes are directly written to the `response` buffer.
833 //     MemoryEncoder encoder(response);
834 //     encoder.WriteUint32(kMagicNumberField, 0x1a1a2b2b);
835 //     encoder.WriteString(kFavoriteFood, "cookies");
836 //     return StatusWithSize(encoder.status(), encoder.size());
837 //   }
838 //
839 // Note: Avoid using a MemoryEncoder reference as an argument for a function.
840 // The StreamEncoder is more generic.
841 class MemoryEncoder : public StreamEncoder {
842  public:
MemoryEncoder(ByteSpan dest)843   constexpr MemoryEncoder(ByteSpan dest) : StreamEncoder(*this, dest) {}
844 
845   // Precondition: Encoder has no active child encoder.
846   //
847   // Postcondition: If this encoder is a nested one, the parent encoder is
848   //     unlocked and proto encoding may resume on the parent.
849   ~MemoryEncoder() = default;
850 
851   // Disallow copy/assign to avoid confusion about who owns the buffer.
852   MemoryEncoder(const MemoryEncoder& other) = delete;
853   MemoryEncoder& operator=(const MemoryEncoder& other) = delete;
854 
855   // It's not safe to move an encoder as it could cause another encoder's
856   // parent_ pointer to become invalid.
857   MemoryEncoder& operator=(MemoryEncoder&& other) = delete;
858 
data()859   const std::byte* data() const { return memory_writer_.data(); }
size()860   size_t size() const { return memory_writer_.bytes_written(); }
861 
begin()862   const std::byte* begin() const { return data(); }
end()863   const std::byte* end() const { return data() + size(); }
864 
865  protected:
866   // This is needed by codegen.
867   MemoryEncoder(MemoryEncoder&& other) = default;
868 };
869 
870 // pw_protobuf guarantees that all generated StreamEncoder classes can be
871 // converted among each other. It's also safe to convert any MemoryEncoder to
872 // any other StreamEncoder.
873 //
874 // This guarantee exists to facilitate usage of protobuf overlays. Protobuf
875 // overlays are protobuf message definitions that deliberately ensure that
876 // fields defined in one message will not conflict with fields defined in other
877 // messages.
878 //
879 // Example:
880 //
881 //   // The first half of the overlaid message.
882 //   message BaseMessage {
883 //     uint32 length = 1;
884 //     reserved 2;  // Reserved for Overlay
885 //   }
886 //
887 //   // OK: The second half of the overlaid message.
888 //   message Overlay {
889 //     reserved 1;  // Reserved for BaseMessage
890 //     uint32 height = 2;
891 //   }
892 //
893 //   // OK: A message that overlays and bundles both types together.
894 //   message Both {
895 //     uint32 length = 1;  // Defined independently by BaseMessage
896 //     uint32 height = 2;  // Defined independently by Overlay
897 //   }
898 //
899 //   // BAD: Diverges from BaseMessage's definition, and can cause decode
900 //   // errors/corruption.
901 //   message InvalidOverlay {
902 //     fixed32 length = 1;
903 //   }
904 //
905 // While this use case is somewhat uncommon, it's a core supported use case of
906 // pw_protobuf.
907 //
908 // Warning: Using this to convert one stream encoder to another when the
909 // messages themselves do not safely overlay will result in corrupt protos.
910 // Be careful when doing this as there's no compile-time way to detect whether
911 // or not two messages are meant to overlay.
912 template <typename ToStreamEncoder, typename FromStreamEncoder>
StreamEncoderCast(FromStreamEncoder & encoder)913 inline ToStreamEncoder& StreamEncoderCast(FromStreamEncoder& encoder) {
914   static_assert(std::is_base_of<StreamEncoder, FromStreamEncoder>::value,
915                 "Provided argument is not a derived class of "
916                 "pw::protobuf::StreamEncoder");
917   static_assert(std::is_base_of<StreamEncoder, ToStreamEncoder>::value,
918                 "Cannot cast to a type that is not a derived class of "
919                 "pw::protobuf::StreamEncoder");
920   return pw::internal::SiblingCast<ToStreamEncoder&, StreamEncoder>(encoder);
921 }
922 
923 }  // namespace pw::protobuf
924