pw_protobuf: Expand functions for calculating encoded proto size

- Add SizeOf* functions for each basic protobuf field type.
- Add tests for serialized_size.h.

Change-Id: Icb6fb674d1904bdbdd73cd66f792e8047734baf2
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/81800
Pigweed-Auto-Submit: Wyatt Hepler <hepler@google.com>
Reviewed-by: Armando Montanez <amontanez@google.com>
Commit-Queue: Wyatt Hepler <hepler@google.com>

7 files changed, 242 insertions, 13 deletions

diff --git a/pw_protobuf/BUILD.bazel b/pw_protobuf/BUILD.bazel
index e9dc16900..711ac97d6 100644
--- a/pw_protobuf/BUILD.bazel
+++ b/pw_protobuf/BUILD.bazel
@@ -119,6 +119,15 @@ pw_cc_test(
 )
 
 pw_cc_test(
+    name = "serialized_size_test",
+    srcs = ["serialized_size_test.cc"],
+    deps = [
+        ":pw_protobuf",
+        "//pw_unit_test",
+    ],
+)
+
+pw_cc_test(
     name = "stream_decoder_test",
     srcs = ["stream_decoder_test.cc"],
     deps = [
diff --git a/pw_protobuf/BUILD.gn b/pw_protobuf/BUILD.gn
index 9d9bea55d..e5b06cd78 100644
--- a/pw_protobuf/BUILD.gn
+++ b/pw_protobuf/BUILD.gn
@@ -104,6 +104,7 @@ pw_test_group("tests") {
     ":find_test",
     ":map_utils_test",
     ":message_test",
+    ":serialized_size_test",
     ":stream_decoder_test",
     ":varint_size_test",
   ]
@@ -129,6 +130,11 @@ pw_test("codegen_test") {
   sources = [ "codegen_test.cc" ]
 }
 
+pw_test("serialized_size_test") {
+  deps = [ ":pw_protobuf" ]
+  sources = [ "serialized_size_test.cc" ]
+}
+
 pw_test("stream_decoder_test") {
   deps = [ ":pw_protobuf" ]
   sources = [ "stream_decoder_test.cc" ]
diff --git a/pw_protobuf/CMakeLists.txt b/pw_protobuf/CMakeLists.txt
index e9d8591c8..271c06a75 100644
--- a/pw_protobuf/CMakeLists.txt
+++ b/pw_protobuf/CMakeLists.txt
@@ -113,6 +113,16 @@ pw_add_test(pw_protobuf.codegen_test
     pw_protobuf
 )
 
+pw_add_test(pw_protobuf.serialized_size_test
+  SOURCES
+    serialized_size_test.cc
+  DEPS
+    pw_protobuf
+  GROUPS
+    modules
+    pw_protobuf
+)
+
 pw_add_test(pw_protobuf.stream_decoder_test
   SOURCES
     stream_decoder_test.cc
diff --git a/pw_protobuf/docs.rst b/pw_protobuf/docs.rst
index d605eade7..e9a970588 100644
--- a/pw_protobuf/docs.rst
+++ b/pw_protobuf/docs.rst
@@ -592,6 +592,16 @@ fields in a message.
 
 .. include:: size_report/decoder_incremental
 
+---------------------------
+Serialized size calculation
+---------------------------
+``pw_protobuf/serialized_size.h`` provides a set of functions for calculating
+how much memory serialized protocol buffer fields require. The
+``kMaxSizeBytes*`` variables provide the maximum encoded sizes of each field
+type. The ``SizeOfField*()`` functions calculate the encoded size of a field of
+the specified type, given a particular key and, for variable length fields
+(varint or delimited), a value.
+
 --------------------------
 Available protobuf modules
 --------------------------
diff --git a/pw_protobuf/public/pw_protobuf/encoder.h b/pw_protobuf/public/pw_protobuf/encoder.h
index 537492de5..c3d68ad6b 100644
--- a/pw_protobuf/public/pw_protobuf/encoder.h
+++ b/pw_protobuf/public/pw_protobuf/encoder.h
@@ -46,8 +46,8 @@ namespace pw::protobuf {
 //     size of the largest one.
 //  max_nested_depth: The max number of nested submessage encoders that are
 //     expected to be open simultaneously to encode this proto message.
-inline constexpr size_t MaxScratchBufferSize(size_t max_message_size,
-                                             size_t max_nested_depth) {
+constexpr size_t MaxScratchBufferSize(size_t max_message_size,
+                                      size_t max_nested_depth) {
   return max_message_size + max_nested_depth * config::kMaxVarintSize;
 }
 
diff --git a/pw_protobuf/public/pw_protobuf/serialized_size.h b/pw_protobuf/public/pw_protobuf/serialized_size.h
index a20ba04e0..7fe5518c4 100644
--- a/pw_protobuf/public/pw_protobuf/serialized_size.h
+++ b/pw_protobuf/public/pw_protobuf/serialized_size.h
@@ -40,7 +40,9 @@ inline constexpr size_t kMaxSizeBytesInt64 = varint::kMaxVarint64SizeBytes;
 // The bool field type is backed by a varint, but has a limited value range.
 inline constexpr size_t kMaxSizeBytesBool = 1;
 
-inline constexpr size_t kMaxSizeOfFieldKey = varint::kMaxVarint32SizeBytes;
+inline constexpr size_t kMaxSizeBytesEnum = kMaxSizeBytesInt32;
+
+inline constexpr size_t kMaxSizeOfFieldNumber = varint::kMaxVarint32SizeBytes;
 
 inline constexpr size_t kMaxSizeOfLength = varint::kMaxVarint32SizeBytes;
 
@@ -55,15 +57,33 @@ inline constexpr size_t kMaxSizeOfLength = varint::kMaxVarint32SizeBytes;
 //   The size of the field key.
 //
 // Precondition: The field_number must be a ValidFieldNumber.
-// Precondition: The field_number must be a ValidFieldNumber.
-constexpr size_t SizeOfFieldKey(uint32_t field_number) {
+template <typename T>
+constexpr size_t FieldNumberSizeBytes(T field_number) {
+  static_assert((std::is_enum<T>() || std::is_integral<T>()) &&
+                    sizeof(T) <= sizeof(uint32_t),
+                "Field numbers must be 32-bit enums or integers");
   // The wiretype does not impact the serialized size, so use kVarint (0), which
   // will be optimized out by the compiler.
-  return varint::EncodedSize(FieldKey(field_number, WireType::kVarint));
+  return varint::EncodedSize(
+      FieldKey(static_cast<uint32_t>(field_number), WireType::kVarint));
+}
+
+// Calculates the size of a varint field (uint32/64, int32/64, sint32/64, enum).
+template <typename T, typename U>
+constexpr size_t SizeOfVarintField(T field_number, U value) {
+  return FieldNumberSizeBytes(field_number) + varint::EncodedSize(value);
+}
+
+// Calculates the size of a delimited field (string, bytes, nested message,
+// packed repeated).
+template <typename T>
+constexpr size_t SizeOfDelimitedField(T field_number, size_t length_bytes) {
+  return FieldNumberSizeBytes(field_number) +
+         varint::EncodedSize(length_bytes) + length_bytes;
 }
 
-// Calculate the size of a proto field in wire format. This is the size of a
-// final serialized protobuf entry, including the key (field number + wire
+// Calculates the size of a proto field in the wire format. This is the size of
+// a final serialized protobuf entry, including the key (field number + wire
 // type), encoded payload size (for length-delimited types), and data.
 //
 // Args:
@@ -77,15 +97,89 @@ constexpr size_t SizeOfFieldKey(uint32_t field_number) {
 // Precondition: The field_number must be a ValidFieldNumber.
 // Precondition: `data_size_bytes` must be smaller than
 //   std::numeric_limits<uint32_t>::max()
-constexpr size_t SizeOfField(uint32_t field_number,
+template <typename T>
+constexpr size_t SizeOfField(T field_number,
                              WireType type,
                              size_t data_size_bytes) {
-  size_t size = SizeOfFieldKey(field_number);
   if (type == WireType::kDelimited) {
-    size += varint::EncodedSize(data_size_bytes);
+    return SizeOfDelimitedField(field_number, data_size_bytes);
   }
-  size += data_size_bytes;
-  return size;
+  return FieldNumberSizeBytes(field_number) + data_size_bytes;
+}
+
+// Functions for calculating the size of each type of protobuf field. Varint
+// fields (int32, uint64, etc.) accept a value argument that defaults to the
+// largest-to-encode value for the type.
+template <typename T>
+constexpr size_t SizeOfFieldFloat(T field_number) {
+  return FieldNumberSizeBytes(field_number) + sizeof(float);
+}
+template <typename T>
+constexpr size_t SizeOfFieldDouble(T field_number) {
+  return FieldNumberSizeBytes(field_number) + sizeof(double);
+}
+template <typename T>
+constexpr size_t SizeOfFieldInt32(T field_number, int32_t value = -1) {
+  return SizeOfVarintField(field_number, value);
+}
+template <typename T>
+constexpr size_t SizeOfFieldInt64(T field_number, int64_t value = -1) {
+  return SizeOfVarintField(field_number, value);
+}
+template <typename T>
+constexpr size_t SizeOfFieldSint32(
+    T field_number, int32_t value = std::numeric_limits<int32_t>::min()) {
+  return SizeOfVarintField(field_number, varint::ZigZagEncode(value));
+}
+template <typename T>
+constexpr size_t SizeOfFieldSint64(
+    T field_number, int64_t value = std::numeric_limits<int64_t>::min()) {
+  return SizeOfVarintField(field_number, varint::ZigZagEncode(value));
+}
+template <typename T>
+constexpr size_t SizeOfFieldUint32(
+    T field_number, uint32_t value = std::numeric_limits<uint32_t>::max()) {
+  return SizeOfVarintField(field_number, value);
+}
+template <typename T>
+constexpr size_t SizeOfFieldUint64(
+    T field_number, uint64_t value = std::numeric_limits<uint64_t>::max()) {
+  return SizeOfVarintField(field_number, value);
+}
+template <typename T>
+constexpr size_t SizeOfFieldFixed32(T field_number) {
+  return FieldNumberSizeBytes(field_number) + sizeof(uint32_t);
+}
+template <typename T>
+constexpr size_t SizeOfFieldFixed64(T field_number) {
+  return FieldNumberSizeBytes(field_number) + sizeof(uint64_t);
+}
+template <typename T>
+constexpr size_t SizeOfFieldSfixed32(T field_number) {
+  return FieldNumberSizeBytes(field_number) + sizeof(uint32_t);
+}
+template <typename T>
+constexpr size_t SizeOfFieldSfixed64(T field_number) {
+  return FieldNumberSizeBytes(field_number) + sizeof(uint64_t);
+}
+template <typename T>
+constexpr size_t SizeOfFieldBool(T field_number) {
+  return FieldNumberSizeBytes(field_number) + 1;
+}
+template <typename T>
+constexpr size_t SizeOfFieldString(T field_number, size_t length_bytes) {
+  return SizeOfDelimitedField(field_number, length_bytes);
+}
+template <typename T>
+constexpr size_t SizeOfFieldBytes(T field_number, size_t length_bytes) {
+  return SizeOfDelimitedField(field_number, length_bytes);
+}
+template <typename T, typename U = int32_t>
+constexpr size_t SizeOfFieldEnum(T field_number, U value = static_cast<U>(-1)) {
+  static_assert((std::is_enum<U>() || std::is_integral<U>()) &&
+                    sizeof(U) <= sizeof(uint32_t),
+                "Enum values must be 32-bit enums or integers");
+  return SizeOfFieldInt32(field_number, static_cast<int32_t>(value));
 }
 
 }  // namespace pw::protobuf
diff --git a/pw_protobuf/serialized_size_test.cc b/pw_protobuf/serialized_size_test.cc
new file mode 100644
index 000000000..3327c66fe
--- /dev/null
+++ b/pw_protobuf/serialized_size_test.cc
@@ -0,0 +1,100 @@
+// Copyright 2022 The Pigweed Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not
+// use this file except in compliance with the License. You may obtain a copy of
+// the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations under
+// the License.
+#include "pw_protobuf/serialized_size.h"
+
+#include <cinttypes>
+
+#include "gtest/gtest.h"
+
+namespace pw::protobuf {
+namespace {
+
+#define TEST_VARINT(type)                                           \
+  TEST(SerializedSize, type) {                                      \
+    static_assert(SizeOfField##type(1, 0) == 1 + 1,                 \
+                  #type " minimum encoded size, key 1");            \
+    static_assert(SizeOfField##type(1) == 1 + kMaxSizeBytes##type,  \
+                  #type " maximum encoded size, key 1");            \
+    static_assert(SizeOfField##type(16, 0) == 2 + 1,                \
+                  #type " minimum encoded size, key 16");           \
+    static_assert(SizeOfField##type(16) == 2 + kMaxSizeBytes##type, \
+                  #type " maximum encoded size, key 16");           \
+  }                                                                 \
+  static_assert(true, "require semicolons")
+
+#define TEST_FIXED(type)                                            \
+  TEST(SerializedSize, SizeOf##type##Field) {                       \
+    static_assert(SizeOfField##type(1) == 1 + kMaxSizeBytes##type,  \
+                  #type " key 1");                                  \
+    static_assert(SizeOfField##type(15) == 1 + kMaxSizeBytes##type, \
+                  #type " key 15");                                 \
+    static_assert(SizeOfField##type(16) == 2 + kMaxSizeBytes##type, \
+                  #type " key 16");                                 \
+    static_assert(SizeOfField##type(17) == 2 + kMaxSizeBytes##type, \
+                  #type " key 17");                                 \
+  }                                                                 \
+  static_assert(true, "require semicolons")
+
+#define TEST_DELIMITED(function)                                            \
+  TEST(SerializedSize, function) {                                          \
+    static_assert(function(1, 0) == 1 + 1 + 0, #function " key 1");         \
+    static_assert(function(1, 1) == 1 + 1 + 1, #function " key 1");         \
+    static_assert(function(1, 128) == 1 + 2 + 128, #function " key 1");     \
+    static_assert(function(1, 1000) == 1 + 2 + 1000, #function " key 1");   \
+    static_assert(function(16, 0) == 2 + 1 + 0, #function " key 16");       \
+    static_assert(function(16, 1) == 2 + 1 + 1, #function " key 16");       \
+    static_assert(function(16, 128) == 2 + 2 + 128, #function " key 16");   \
+    static_assert(function(16, 1000) == 2 + 2 + 1000, #function " key 16"); \
+  }                                                                         \
+  static_assert(true, "require semicolons")
+
+TEST(SerializedSize, SizeOfVarintField) {
+  static_assert(SizeOfVarintField(1, 0) == 1 + 1);
+  static_assert(SizeOfVarintField(1, 127) == 1 + 1);
+
+  static_assert(SizeOfVarintField(1, INT32_C(-1)) == 1 + 10);
+  static_assert(SizeOfVarintField(1, INT64_C(-1)) == 1 + 10);
+
+  static_assert(SizeOfVarintField(16, 0) == 2 + 1);
+  static_assert(SizeOfVarintField(16, 127) == 2 + 1);
+
+  static_assert(SizeOfVarintField(16, INT32_C(-1)) == 2 + 10);
+  static_assert(SizeOfVarintField(16, INT64_C(-1)) == 2 + 10);
+}
+
+TEST_DELIMITED(SizeOfDelimitedField);
+
+TEST_FIXED(Float);
+TEST_FIXED(Double);
+
+TEST_VARINT(Int32);
+TEST_VARINT(Int64);
+TEST_VARINT(Sint32);
+TEST_VARINT(Sint64);
+TEST_VARINT(Uint32);
+TEST_VARINT(Uint64);
+
+TEST_FIXED(Fixed32);
+TEST_FIXED(Fixed64);
+TEST_FIXED(Sfixed32);
+TEST_FIXED(Sfixed64);
+TEST_FIXED(Bool);
+
+TEST_DELIMITED(SizeOfFieldString);
+TEST_DELIMITED(SizeOfFieldBytes);
+
+TEST_VARINT(Enum);
+
+}  // namespace
+}  // namespace pw::protobuf