pw_protobuf: Add StreamEncoderCast<>()

Adds a helper template to make casting between StreamEncoder types
easier to read.

Fixes: b/234875243
Change-Id: Ia74ef6259260d389b3a15be93768b43ea56ebee6
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/97909
Reviewed-by: Wyatt Hepler <hepler@google.com>
Commit-Queue: Auto-Submit <auto-submit@pigweed.google.com.iam.gserviceaccount.com>
Pigweed-Auto-Submit: Armando Montanez <amontanez@google.com>
diff --git a/pw_protobuf/codegen_encoder_test.cc b/pw_protobuf/codegen_encoder_test.cc
index 2b64eda..abc9114 100644
--- a/pw_protobuf/codegen_encoder_test.cc
+++ b/pw_protobuf/codegen_encoder_test.cc
@@ -534,5 +534,38 @@
   EXPECT_EQ(packed.status(), OkStatus());
 }
 
+TEST(Codegen, MemoryToStreamConversion) {
+  std::byte encode_buffer[IntegerMetadata::kMaxEncodedSizeBytes];
+  IntegerMetadata::MemoryEncoder metadata(encode_buffer);
+  IntegerMetadata::StreamEncoder& streamed_metadata = metadata;
+  EXPECT_EQ(streamed_metadata.WriteBits(3), OkStatus());
+
+  constexpr uint8_t expected_proto[] = {0x08, 0x03};
+
+  ConstByteSpan result(metadata);
+  ASSERT_EQ(metadata.status(), OkStatus());
+  EXPECT_EQ(result.size(), sizeof(expected_proto));
+  EXPECT_EQ(std::memcmp(result.data(), expected_proto, sizeof(expected_proto)),
+            0);
+}
+
+TEST(Codegen, OverlayConversion) {
+  std::byte encode_buffer[BaseMessage::kMaxEncodedSizeBytes +
+                          Overlay::kMaxEncodedSizeBytes];
+  BaseMessage::MemoryEncoder base(encode_buffer);
+  Overlay::StreamEncoder& overlay =
+      StreamEncoderCast<Overlay::StreamEncoder>(base);
+  EXPECT_EQ(overlay.WriteHeight(15), OkStatus());
+  EXPECT_EQ(base.WriteLength(7), OkStatus());
+
+  constexpr uint8_t expected_proto[] = {0x10, 0x0f, 0x08, 0x07};
+
+  ConstByteSpan result(base);
+  ASSERT_EQ(base.status(), OkStatus());
+  EXPECT_EQ(result.size(), sizeof(expected_proto));
+  EXPECT_EQ(std::memcmp(result.data(), expected_proto, sizeof(expected_proto)),
+            0);
+}
+
 }  // namespace
 }  // namespace pw::protobuf
diff --git a/pw_protobuf/docs.rst b/pw_protobuf/docs.rst
index f7da518..70bfd50 100644
--- a/pw_protobuf/docs.rst
+++ b/pw_protobuf/docs.rst
@@ -253,6 +253,83 @@
 value are final, and can't be referenced or modified as a later step in the
 encode process.
 
+Casting between generated StreamEncoder types
+=============================================
+pw_protobuf guarantees that all generated ``StreamEncoder`` classes can be
+converted among each other. It's also safe to convert any ``MemoryEncoder`` to
+any other ``StreamEncoder``.
+
+This guarantee exists to facilitate usage of protobuf overlays. Protobuf
+overlays are protobuf message definitions that deliberately ensure that
+fields defined in one message will not conflict with fields defined in other
+messages.
+
+For example:
+
+.. code::
+
+  // The first half of the overlaid message.
+  message BaseMessage {
+    uint32 length = 1;
+    reserved 2;  // Reserved for Overlay
+  }
+
+  // OK: The second half of the overlaid message.
+  message Overlay {
+    reserved 1;  // Reserved for BaseMessage
+    uint32 height = 2;
+  }
+
+  // OK: A message that overlays and bundles both types together.
+  message Both {
+    uint32 length = 1;  // Defined independently by BaseMessage
+    uint32 height = 2;  // Defined independently by Overlay
+  }
+
+  // BAD: Diverges from BaseMessage's definition, and can cause decode
+  // errors/corruption.
+  message InvalidOverlay {
+    fixed32 length = 1;
+  }
+
+The ``StreamEncoderCast<>()`` helper template reduces very messy casting into
+a much easier to read syntax:
+
+.. code:: c++
+
+  #include "pw_protobuf/encoder.h"
+  #include "pw_protobuf_test_protos/full_test.pwpb.h"
+
+  Result<ConstByteSpan> EncodeOverlaid(uint32_t height,
+                                       uint32_t length,
+                                       ConstByteSpan encode_buffer) {
+    BaseMessage::MemoryEncoder base(encode_buffer);
+
+    // Without StreamEncoderCast<>(), this line would be:
+    //   Overlay::StreamEncoder& overlay =
+    //       *static_cast<Overlay::StreamEncoder*>(
+    //           static_cast<pw::protobuf::StreamEncoder*>(&base)
+    Overlay::StreamEncoder& overlay =
+        StreamEncoderCast<Overlay::StreamEncoder>(base);
+    if (!overlay.WriteHeight(height).ok()) {
+      return overlay.status();
+    }
+    if (!base.WriteLength(length).ok()) {
+      return base.status();
+    }
+    return ConstByteSpan(base);
+  }
+
+While this use case is somewhat uncommon, it's a core supported use case of
+pw_protobuf.
+
+.. warning::
+
+  Using this to convert one stream encoder to another when the messages
+  themselves do not safely overlay will result in corrupt protos. Be careful
+  when doing this as there's no compile-time way to detect whether or not two
+  messages are meant to overlay.
+
 Decoding
 --------
 For decoding, in addition to the ``Read()`` method that populates a message
diff --git a/pw_protobuf/public/pw_protobuf/encoder.h b/pw_protobuf/public/pw_protobuf/encoder.h
index fc724fb..8920511 100644
--- a/pw_protobuf/public/pw_protobuf/encoder.h
+++ b/pw_protobuf/public/pw_protobuf/encoder.h
@@ -849,4 +849,57 @@
   MemoryEncoder(MemoryEncoder&& other) = default;
 };
 
+// pw_protobuf guarantees that all generated StreamEncoder classes can be
+// converted among each other. It's also safe to convert any MemoryEncoder to
+// any other StreamEncoder.
+//
+// This guarantee exists to facilitate usage of protobuf overlays. Protobuf
+// overlays are protobuf message definitions that deliberately ensure that
+// fields defined in one message will not conflict with fields defined in other
+// messages.
+//
+// Example:
+//
+//   // The first half of the overlaid message.
+//   message BaseMessage {
+//     uint32 length = 1;
+//     reserved 2;  // Reserved for Overlay
+//   }
+//
+//   // OK: The second half of the overlaid message.
+//   message Overlay {
+//     reserved 1;  // Reserved for BaseMessage
+//     uint32 height = 2;
+//   }
+//
+//   // OK: A message that overlays and bundles both types together.
+//   message Both {
+//     uint32 length = 1;  // Defined independently by BaseMessage
+//     uint32 height = 2;  // Defined independently by Overlay
+//   }
+//
+//   // BAD: Diverges from BaseMessage's definition, and can cause decode
+//   // errors/corruption.
+//   message InvalidOverlay {
+//     fixed32 length = 1;
+//   }
+//
+// While this use case is somewhat uncommon, it's a core supported use case of
+// pw_protobuf.
+//
+// Warning: Using this to convert one stream encoder to another when the
+// messages themselves do not safely overlay will result in corrupt protos.
+// Be careful when doing this as there's no compile-time way to detect whether
+// or not two messages are meant to overlay.
+template <typename ToStreamEncoder, typename FromStreamEncoder>
+inline ToStreamEncoder& StreamEncoderCast(FromStreamEncoder& encoder) {
+  static_assert(std::is_base_of<StreamEncoder, FromStreamEncoder>::value,
+                "Provided argument is not a derived class of "
+                "pw::protobuf::StreamEncoder");
+  static_assert(std::is_base_of<StreamEncoder, ToStreamEncoder>::value,
+                "Cannot cast to a type that is not a derived class of "
+                "pw::protobuf::StreamEncoder");
+  return static_cast<ToStreamEncoder&>(static_cast<StreamEncoder&>(encoder));
+}
+
 }  // namespace pw::protobuf
diff --git a/pw_protobuf/pw_protobuf_test_protos/full_test.proto b/pw_protobuf/pw_protobuf_test_protos/full_test.proto
index 1309a6c..1aa8b85 100644
--- a/pw_protobuf/pw_protobuf_test_protos/full_test.proto
+++ b/pw_protobuf/pw_protobuf_test_protos/full_test.proto
@@ -143,6 +143,17 @@
   bool signed = 2;  // `signed` should become `signed_` in the C++ code.
 }
 
+// Two messages that should properly overlay without collisions.
+message BaseMessage {
+  uint32 length = 1;
+  reserved 2;
+}
+
+message Overlay {
+  reserved 1;
+  uint32 height = 2;
+}
+
 // Ensure that messages named `Message` don't cause namespace-resolution issues
 // when the codegen internally references the generated type `Message::Message`.
 message Message {