From 969664ceaa20fcf8fdf950bd77b5067ffcc79c08 Mon Sep 17 00:00:00 2001
From: Thad House <ThadHouse@users.noreply.github.com>
Date: Sat, 16 Nov 2024 15:26:10 +0000
Subject: [PATCH] [wpiutil] Faster nanopb submessage encode (#7374)

Due to how submessages are encoded (with a length prefix), nanopb currently does the encoding twice. It encodes once to get the length to write, then writes the length, then reencodes the entire message a 2nd time.

This results in a requirement that each encode always encodes the same. Generally, this is fine, but it'd be nice to not make this a requirement.

The double encode also requires going through the entire set of fields again, which has the possibility to be slow.

Instead of doing this, write to a temporary SmallVector. Then we can just encode the length of that buffer, and do a memcpy into primary stream. Theoretically in most cases, this should be much faster.
---
 .../src/main/native/cpp/protobuf/Protobuf.cpp | 25 +++++++++++++++++++
 .../native/include/wpi/protobuf/Protobuf.h    |  6 ++++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/wpiutil/src/main/native/cpp/protobuf/Protobuf.cpp b/wpiutil/src/main/native/cpp/protobuf/Protobuf.cpp
index a99cb7c164..b822776b37 100644
--- a/wpiutil/src/main/native/cpp/protobuf/Protobuf.cpp
+++ b/wpiutil/src/main/native/cpp/protobuf/Protobuf.cpp
@@ -48,3 +48,28 @@ bool detail::WriteFromStdVector(pb_ostream_t* stream, const pb_byte_t* buf,
   vec->insert(vec->end(), buf, buf + count);
   return true;
 }
+
+bool detail::WriteSubmessage(pb_ostream_t* stream, const pb_msgdesc_t* desc,
+                             const void* msg) {
+  // Write the submessage to a separate buffer
+  wpi::SmallVector<uint8_t, 64> buf;
+  pb_ostream_t subStream{
+      .callback = WriteFromSmallVector,
+      .state = &buf,
+      .max_size = SIZE_MAX,
+      .bytes_written = 0,
+      .errmsg = nullptr,
+  };
+  if (!pb_encode(&subStream, desc, msg)) {
+    return false;
+  }
+
+  uint64_t size = static_cast<uint64_t>(buf.size());
+
+  // Write the size to the original stream.
+  if (!pb_encode_varint(stream, size)) {
+    return false;
+  }
+
+  return pb_write(stream, reinterpret_cast<const pb_byte_t*>(buf.data()), size);
+}
diff --git a/wpiutil/src/main/native/include/wpi/protobuf/Protobuf.h b/wpiutil/src/main/native/include/wpi/protobuf/Protobuf.h
index 34ad3b3aca..381ae698a2 100644
--- a/wpiutil/src/main/native/include/wpi/protobuf/Protobuf.h
+++ b/wpiutil/src/main/native/include/wpi/protobuf/Protobuf.h
@@ -43,6 +43,9 @@ bool WriteFromSmallVector(pb_ostream_t* stream, const pb_byte_t* buf,
 
 bool WriteFromStdVector(pb_ostream_t* stream, const pb_byte_t* buf,
                         size_t count);
+
+bool WriteSubmessage(pb_ostream_t* stream, const pb_msgdesc_t* desc,
+                     const void* msg);
 }  // namespace detail
 
 /**
@@ -208,7 +211,8 @@ class ProtoOutputStream {
   bool Encode(
       const typename Protobuf<std::remove_cvref_t<T>>::MessageStruct& msg) {
     if (m_streamMsg) {
-      return pb_encode_submessage(m_streamMsg, m_msgDesc, &msg);
+      return detail::WriteSubmessage(m_streamMsg, m_msgDesc, &msg);
+      // return pb_encode_submessage(m_streamMsg, m_msgDesc, &msg);
     }
     return pb_encode(&m_streamLocal, m_msgDesc, &msg);
   }