RustProto: add serialized_len()->usize

PiperOrigin-RevId: 904638364

rust/codegen_traits.rs

@@ -168,6 +168,7 @@ pub(crate) mod read {
 
     pub trait Serialize: SealedInternal {
         fn serialize(&self) -> Result<Vec<u8>, crate::SerializeError>;
+        fn serialized_len(&self) -> usize;
     }
 }
 

rust/cpp_kernel/message.cc

@@ -1,5 +1,6 @@
 #include "google/protobuf/message.h"
 
+#include <cstddef>
 #include <limits>
 
 #include "google/protobuf/message_lite.h"
@@ -37,6 +38,10 @@ bool proto2_rust_Message_serialize(const google::protobuf::MessageLite* m,
   return google::protobuf::rust::SerializeMsg(m, output);
 }
 
+size_t proto2_rust_Message_serialized_len(const google::protobuf::MessageLite* m) {
+  return m->ByteSizeLong();
+}
+
 void proto2_rust_Message_copy_from(google::protobuf::MessageLite* dst,
                                    const google::protobuf::MessageLite& src) {
   dst->Clear();

rust/cpp_kernel/message.rs

@@ -7,6 +7,7 @@ unsafe extern "C" {
     pub fn proto2_rust_Message_parse_dont_enforce_required(m: RawMessage, input: PtrAndLen)
         -> bool;
     pub fn proto2_rust_Message_serialize(m: RawMessage, output: &mut SerializedData) -> bool;
+    pub fn proto2_rust_Message_serialized_len(m: RawMessage) -> usize;
     pub fn proto2_rust_Message_copy_from(dst: RawMessage, src: RawMessage);
     pub fn proto2_rust_Message_merge_from(dst: RawMessage, src: RawMessage);
     pub fn proto2_rust_Message_get_descriptor(m: RawMessage) -> *const std::ffi::c_void;
@@ -293,6 +294,10 @@ impl<T: CppGetRawMessage> Serialize for T {
             Err(SerializeError)
         }
     }
+
+    fn serialized_len(&self) -> usize {
+        unsafe { proto2_rust_Message_serialized_len(self.get_raw_message(Private)) }
+    }
 }
 
 impl<T> TakeFrom for T

rust/test/shared/serialization_test.rs

@@ -36,6 +36,29 @@ macro_rules! generate_parameterized_serialization_test {
                 assert_that!(serialized.len(), eq(0));
             }
 
+            #[gtest]
+            fn [< serialized_len_matches_serialize_ $name_ext >]() {
+                let mut msg = [< $type >]::new();
+                msg.set_optional_int64(42);
+                msg.set_optional_bool(true);
+                msg.set_optional_bytes(b"serialize deserialize test");
+
+                let serialized = msg.serialize().unwrap();
+                assert_that!(msg.serialized_len(), eq(serialized.len()));
+
+                let serialized = msg.as_view().serialize().unwrap();
+                assert_that!(msg.as_view().serialized_len(), eq(serialized.len()));
+
+                let serialized = msg.as_mut().serialize().unwrap();
+                assert_that!(msg.as_mut().serialized_len(), eq(serialized.len()));
+            }
+
+            #[gtest]
+            fn [< serialized_len_empty_message_ $name_ext >]() {
+                let msg = [< $type >]::new();
+                assert_that!(msg.serialized_len(), eq(0));
+            }
+
             #[gtest]
             fn [< serialize_default_view $name_ext>]() {
                 let default = View::<[< $type >]>::default();

rust/upb/sys/BUILD

@@ -58,6 +58,7 @@ cc_library(
         "//upb/mini_descriptor",
         "//upb/mini_table",
         "//upb/text:debug",
+        "//upb/wire:byte_size",
     ],
 )
 

rust/upb/sys/upb_api.c

@@ -21,4 +21,5 @@
 #include "upb/mini_descriptor/decode.h"  // IWYU pragma: keep
 #include "upb/mini_table/message.h"      // IWYU pragma: keep
 #include "upb/text/debug_string.h"       // IWYU pragma: keep
+#include "upb/wire/byte_size.h"          // IWYU pragma: keep
 // go/keep-sorted end

rust/upb/sys/wire/wire.rs

@@ -44,6 +44,7 @@ pub enum DecodeStatus {
 unsafe extern "C" {
     // SAFETY:
     // - `mini_table` is the one associated with `msg`
+    // - `msg` is legal to dereference and read from.
     // - `buf` and `buf_size` are legally writable.
     pub fn upb_Encode(
         msg: RawMessage,
@@ -56,6 +57,7 @@ unsafe extern "C" {
 
     // SAFETY:
     // - `mini_table` is the one associated with `msg`
+    // - `msg` is legal to dereference and read from.
     // - `buf` is legally readable for at least `buf_size` bytes.
     // - `extreg` is either null or points at a valid upb_ExtensionRegistry.
     pub fn upb_Decode(
@@ -67,6 +69,11 @@ unsafe extern "C" {
         options: i32,
         arena: RawArena,
     ) -> DecodeStatus;
+
+    // SAFETY:
+    // - `msg` is legal to dereference and read from.
+    // - `mini_table` is the one associated with `msg`.
+    pub fn upb_ByteSize(msg: RawMessage, mini_table: RawMiniTable) -> usize;
 }
 
 #[cfg(test)]
@@ -79,5 +86,6 @@ mod tests {
         use crate::assert_linked;
         assert_linked!(upb_Encode);
         assert_linked!(upb_Decode);
+        assert_linked!(upb_ByteSize);
     }
 }

rust/upb/wire.rs

@@ -6,7 +6,7 @@
 // https://developers.google.com/open-source/licenses/bsd
 
 use super::sys::mini_table::extension_registry::upb_ExtensionRegistry;
-use super::sys::wire::wire::{upb_Decode, upb_Encode, DecodeStatus, EncodeStatus};
+use super::sys::wire::wire::{upb_ByteSize, upb_Decode, upb_Encode, DecodeStatus, EncodeStatus};
 use super::{Arena, AssociatedMiniTable, MessagePtr};
 
 /// Contains the decode options that can be passed to `decode_with_options`.
@@ -42,6 +42,14 @@ pub fn encode<T: AssociatedMiniTable>(msg: MessagePtr<T>) -> Result<Vec<u8>, Enc
     }
 }
 
+/// Returns the serialized size of the message.
+pub fn byte_size<T: AssociatedMiniTable>(msg: MessagePtr<T>) -> usize {
+    // SAFETY:
+    // - `T::mini_table()` is the one associated with `msg`.
+    // - `msg` is guaranteed live.
+    unsafe { upb_ByteSize(msg.raw(), T::mini_table()) }
+}
+
 /// Decodes into the provided message (merge semantics). If Err, then
 /// DecodeStatus != Ok.
 ///

rust/upb_kernel/message.rs

@@ -392,6 +392,10 @@ where
         //~ of the failure, we should try to keep it instead.
         upb::wire::encode(self.get_ptr(Private)).map_err(|_| SerializeError)
     }
+
+    fn serialized_len(&self) -> usize {
+        upb::wire::byte_size(self.get_ptr(Private))
+    }
 }
 
 impl<T> TakeFrom for T

upb/BUILD

@@ -160,6 +160,7 @@ upb_amalgamation(
         "//upb/text:debug",
         "//upb/text:internal",
         "//upb/wire",
+        "//upb/wire:byte_size",
         "//upb/wire:decoder",
         "//upb/wire:eps_copy_input_stream",
         "//upb/wire:eps_copy_input_stream_internal",