Change Rust-upb to use upb_Encode and upb_Decode directly, instead of going through the upb C gencode.

PiperOrigin-RevId: 623474125

rust/upb.rs

@@ -332,19 +332,28 @@ fn copy_bytes_in_arena<'msg>(arena: &'msg Arena, val: &'msg [u8]) -> &'msg [u8]
     }
 }
 
-/// Opaque struct containing a upb_MiniTable.
-///
-/// This wrapper is a workaround until stabilization of [`extern type`].
-/// TODO: convert to extern type once stabilized.
-/// [`extern type`]: https://github.com/rust-lang/rust/issues/43467
-#[repr(C)]
-pub struct OpaqueMiniTable {
-    // TODO: consider importing a minitable struct declared in
-    // google3/third_party/upb/bits.
-    _data: [u8; 0],
-    _marker: std::marker::PhantomData<(*mut u8, ::std::marker::PhantomPinned)>,
+// A macro for opaque pointees.
+// This wrapper is a workaround until stabilization of [`extern type`].
+// TODO: convert to extern type once stabilized.
+// [`extern type`]: https://github.com/rust-lang/rust/issues/43467
+macro_rules! opaque_pointee {
+    ($name:ident) => {
+        #[repr(C)]
+        pub struct $name {
+            _data: [u8; 0],
+            _marker: std::marker::PhantomData<(*mut u8, ::std::marker::PhantomPinned)>,
+        }
+    };
 }
 
+// Opaque pointee for upb_MiniTable.
+// TODO: consider importing a minitable struct declared in
+// google3/third_party/upb/bits.
+opaque_pointee!(OpaqueMiniTable);
+
+// Opaque pointee for upb_ExtensionRegistry
+opaque_pointee!(OpaqueExtensionRegistry);
+
 extern "C" {
     pub fn upb_Message_DeepCopy(
         dst: RawMessage,
@@ -352,6 +361,7 @@ extern "C" {
         mini_table: *const OpaqueMiniTable,
         arena: RawArena,
     );
+
     pub fn upb_Message_DeepClone(
         m: RawMessage,
         mini_table: *const OpaqueMiniTable,
@@ -359,6 +369,52 @@ extern "C" {
     ) -> Option<RawMessage>;
 }
 
+// LINT.IfChange(encode_status)
+#[repr(C)]
+#[derive(PartialEq, Eq, Copy, Clone)]
+pub enum EncodeStatus {
+    Ok = 0,
+    OutOfMemory = 1,
+    MaxDepthExceeded = 2,
+    MissingRequired = 3,
+}
+// LINT.ThenChange()
+
+// LINT.IfChange(decode_status)
+#[repr(C)]
+#[derive(PartialEq, Eq, Copy, Clone)]
+pub enum DecodeStatus {
+    Ok = 0,
+    Malformed = 1,
+    OutOfMemory = 2,
+    BadUtf8 = 3,
+    MaxDepthExceeded = 4,
+    MissingRequired = 5,
+    UnlinkedSubMessage = 6,
+}
+// LINT.ThenChange()
+
+extern "C" {
+    pub fn upb_Encode(
+        msg: RawMessage,
+        mini_table: *const OpaqueMiniTable,
+        options: i32,
+        arena: RawArena,
+        buf: *mut *mut u8,
+        buf_size: *mut usize,
+    ) -> EncodeStatus;
+
+    pub fn upb_Decode(
+        buf: *const u8,
+        buf_size: usize,
+        msg: RawMessage,
+        mini_table: *const OpaqueMiniTable,
+        extreg: *const OpaqueExtensionRegistry,
+        options: i32,
+        arena: RawArena,
+    ) -> DecodeStatus;
+}
+
 /// The raw type-erased version of an owned `Repeated`.
 #[derive(Debug)]
 pub struct InnerRepeated {

src/google/protobuf/compiler/rust/message.cc

@@ -66,11 +66,35 @@ void MessageSerialize(Context& ctx, const Descriptor& msg) {
       return;
 
     case Kernel::kUpb:
-      ctx.Emit({{"serialize_thunk", ThunkName(ctx, msg, "serialize")}}, R"rs(
+      ctx.Emit({{"minitable", UpbMinitableName(msg)}},
+               R"rs(
         let arena = $pbr$::Arena::new();
+        // SAFETY: $minitable$ is a static of a const object.
+        let mini_table = unsafe { $std$::ptr::addr_of!($minitable$) };
+        let options = 0;
+        let mut buf: *mut u8 = std::ptr::null_mut();
         let mut len = 0;
+
+        // SAFETY: `mini_table` is the corresponding one that was used to
+        // construct `self.raw_msg()`.
+        let status = unsafe {
+          $pbr$::upb_Encode(self.raw_msg(), mini_table, options, arena.raw(),
+              &mut buf, &mut len)
+        };
+
+        //~ TODO: Currently serialize() on the Rust API is an
+        //~ infallible fn, so if upb signals an error here we can only panic.
+        assert!(status == $pbr$::EncodeStatus::Ok);
+        let data = if len == 0 {
+          std::ptr::NonNull::dangling()
+        } else {
+          std::ptr::NonNull::new(buf).unwrap()
+        };
+
+        // SAFETY:
+        // - `arena` allocated `data`.
+        // - `data` is valid for reads up to `len` and will not be mutated.
         unsafe {
-          let data = $serialize_thunk$(self.raw_msg(), arena.raw(), &mut len);
           $pbr$::SerializedData::from_raw_parts(arena, data, len)
         }
       )rs");
@@ -89,6 +113,7 @@ void MessageClearAndParse(Context& ctx, const Descriptor& msg) {
           },
           R"rs(
           let success = unsafe {
+            // SAFETY: `data.as_ptr()` is valid to read for `data.len()`.
             let data = $pbr$::SerializedData::from_raw_parts(
               $NonNull$::new(data.as_ptr() as *mut _).unwrap(),
               data.len(),
@@ -101,21 +126,32 @@ void MessageClearAndParse(Context& ctx, const Descriptor& msg) {
       return;
 
     case Kernel::kUpb:
-      ctx.Emit({{"parse_thunk", ThunkName(ctx, msg, "parse")}}, R"rs(
-        let arena = $pbr$::Arena::new();
-        let msg = unsafe {
-          $parse_thunk$(data.as_ptr(), data.len(), arena.raw())
-        };
+      ctx.Emit({{"minitable", UpbMinitableName(msg)}},
+               R"rs(
+        let mut msg = Self::new();
+        // SAFETY: $minitable$ is a static of a const object.
+        let mini_table = unsafe { $std$::ptr::addr_of!($minitable$) };
+        let ext_reg = std::ptr::null();
+        let options = 0;
 
-        match msg {
-          None => Err($pb$::ParseError),
-          Some(msg) => {
-            //~ This assignment causes self.arena to be dropped and to deallocate
-            //~ any previous message pointed/owned to by self.inner.msg.
-            self.inner.arena = arena;
-            self.inner.msg = msg;
+        // SAFETY:
+        // - `data.as_ptr()` is valid to read for `data.len()`
+        // - `mini_table` is the one used to construct `msg.raw_msg()`
+        // - `msg.arena().raw()` is held for the same lifetime as `msg`.
+        let status = unsafe {
+          $pbr$::upb_Decode(
+              data.as_ptr(), data.len(), msg.raw_msg(),
+              mini_table, ext_reg, options, msg.arena().raw())
+        };
+        match status {
+          $pbr$::DecodeStatus::Ok => {
+            //~ This swap causes the old self.inner.arena to be moved into `msg`
+            //~ which we immediately drop, which will release any previous
+            //~ message that was held here.
+            std::mem::swap(self, &mut msg);
             Ok(())
           }
+          _ => Err($pb$::ParseError)
         }
       )rs");
       return;
@@ -184,14 +220,10 @@ void MessageExterns(Context& ctx, const Descriptor& msg) {
       ctx.Emit(
           {
               {"new_thunk", ThunkName(ctx, msg, "new")},
-              {"serialize_thunk", ThunkName(ctx, msg, "serialize")},
-              {"parse_thunk", ThunkName(ctx, msg, "parse")},
               {"minitable", UpbMinitableName(msg)},
           },
           R"rs(
           fn $new_thunk$(arena: $pbi$::RawArena) -> $pbi$::RawMessage;
-          fn $serialize_thunk$(msg: $pbi$::RawMessage, arena: $pbi$::RawArena, len: &mut usize) -> $NonNull$<u8>;
-          fn $parse_thunk$(data: *const u8, size: usize, arena: $pbi$::RawArena) -> Option<$pbi$::RawMessage>;
           /// Opaque wrapper for this message's MiniTable. The only valid way to
           /// reference this static is with `std::ptr::addr_of!(..)`.
           static $minitable$: $pbr$::OpaqueMiniTable;

upb/wire/decode.h

@@ -110,6 +110,7 @@ UPB_INLINE int upb_Decode_LimitDepth(uint32_t decode_options, uint32_t limit) {
   return upb_DecodeOptions_MaxDepth(max_depth) | (decode_options & 0xffff);
 }
 
+// LINT.IfChange
 typedef enum {
   kUpb_DecodeStatus_Ok = 0,
   kUpb_DecodeStatus_Malformed = 1,    // Wire format was corrupt
@@ -127,6 +128,7 @@ typedef enum {
   // of options.
   kUpb_DecodeStatus_UnlinkedSubMessage = 6,
 } upb_DecodeStatus;
+// LINT.ThenChange(//depot/google3/third_party/protobuf/rust/upb.rs:decode_status)
 
 UPB_API upb_DecodeStatus upb_Decode(const char* buf, size_t size,
                                     upb_Message* msg, const upb_MiniTable* mt,

upb/wire/encode.h

@@ -40,6 +40,7 @@ enum {
   kUpb_EncodeOption_CheckRequired = 4,
 };
 
+// LINT.IfChange
 typedef enum {
   kUpb_EncodeStatus_Ok = 0,
   kUpb_EncodeStatus_OutOfMemory = 1,  // Arena alloc failed
@@ -48,6 +49,7 @@ typedef enum {
   // kUpb_EncodeOption_CheckRequired failed but the parse otherwise succeeded.
   kUpb_EncodeStatus_MissingRequired = 3,
 } upb_EncodeStatus;
+// LINT.ThenChange(//depot/google3/third_party/protobuf/rust/upb.rs:encode_status)
 
 UPB_INLINE uint32_t upb_EncodeOptions_MaxDepth(uint16_t depth) {
   return (uint32_t)depth << 16;