Make protobuf_upb::Arena into a proper Rust RAII type.

PiperOrigin-RevId: 525775419

rust/cpp_kernel/cpp.rs

@@ -30,27 +30,73 @@
 
 // Rust Protobuf runtime using the C++ kernel.
 
-use std::alloc::{dealloc, Layout};
+use std::alloc;
+use std::alloc::Layout;
 use std::boxed::Box;
+use std::cell::UnsafeCell;
+use std::fmt;
+use std::marker::PhantomData;
+use std::mem::MaybeUninit;
 use std::ops::Deref;
 use std::ptr::NonNull;
 use std::slice;
 
+/// A wrapper over a `proto2::Arena`.
+///
+/// This is not a safe wrapper per se, because the allocation functions still
+/// have sharp edges (see their safety docs for more info).
+///
+/// This is an owning type and will automatically free the arena when
+/// dropped.
+///
+/// Note that this type is neither `Sync` nor `Send`.
+///
 /// TODO(b/272728844): Replace this placeholder code with a real implementation.
-#[repr(C)]
 pub struct Arena {
-    _data: [u8; 0],
+    ptr: NonNull<u8>,
+    _not_sync: PhantomData<UnsafeCell<()>>,
 }
 
 impl Arena {
-    pub unsafe fn new() -> *mut Self {
-        let arena = Box::new(Arena { _data: [] });
-        Box::leak(arena) as *mut _
+    /// Allocates a fresh arena.
+    #[inline]
+    pub fn new() -> Self {
+        Self { ptr: NonNull::dangling(), _not_sync: PhantomData }
+    }
+
+    /// Returns the raw, C++-managed pointer to the arena.
+    #[inline]
+    pub fn raw(&self) -> ! {
+        unimplemented!()
+    }
+
+    /// Allocates some memory on the arena.
+    ///
+    /// # Safety
+    ///
+    /// `layout`'s alignment must be less than `UPB_MALLOC_ALIGN`.
+    #[inline]
+    pub unsafe fn alloc(&self, layout: Layout) -> &mut [MaybeUninit<u8>] {
+        unimplemented!()
+    }
+
+    /// Resizes some memory on the arena.
+    ///
+    /// # Safety
+    ///
+    /// After calling this function, `ptr` is essentially zapped. `old` must
+    /// be the layout `ptr` was allocated with via [`Arena::alloc()`]. `new`'s
+    /// alignment must be less than `UPB_MALLOC_ALIGN`.
+    #[inline]
+    pub unsafe fn resize(&self, ptr: *mut u8, old: Layout, new: Layout) -> &[MaybeUninit<u8>] {
+        unimplemented!()
+    }
 }
 
-    pub unsafe fn free(arena: *mut Self) {
-        let arena = Box::from_raw(arena);
-        std::mem::drop(arena);
+impl Drop for Arena {
+    #[inline]
+    fn drop(&mut self) {
+        // unimplemented
     }
 }
 
@@ -63,7 +109,6 @@ impl Arena {
 // LINT.IfChange
 // copybara:strip_end
 #[repr(C)]
-#[derive(Debug)]
 pub struct SerializedData {
     /// Owns the memory.
     data: NonNull<u8>,
@@ -89,11 +134,17 @@ impl Deref for SerializedData {
 impl Drop for SerializedData {
     fn drop(&mut self) {
         unsafe {
-            dealloc(self.data.as_ptr(), Layout::array::<u8>(self.len).unwrap());
+            alloc::dealloc(self.data.as_ptr(), Layout::array::<u8>(self.len).unwrap());
         };
     }
 }
 
+impl fmt::Debug for SerializedData {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        fmt::Debug::fmt(self.deref(), f)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

rust/protobuf.rs

@@ -35,12 +35,7 @@
 //! this crate exists is to be able to use `protobuf` as a crate name for both
 //! cpp and upb kernels from user code.
 
-#[cfg(cpp_kernel)]
-pub use protobuf_cpp::__runtime;
 #[cfg(cpp_kernel)]
 pub use protobuf_cpp::*;
-
-#[cfg(upb_kernel)]
-pub use protobuf_upb::__runtime;
 #[cfg(upb_kernel)]
 pub use protobuf_upb::*;

rust/shared.rs

@@ -38,7 +38,6 @@ pub extern crate cpp as __runtime;
 #[cfg(upb_kernel)]
 pub extern crate upb as __runtime;
 
-pub use __runtime::Arena;
 pub use __runtime::SerializedData;
 
 use std::fmt;

rust/upb_kernel/upb.rs

@@ -28,46 +28,122 @@
 // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-// Rust Protobuf runtime using the UPB kernel.
+//! UPB FFI wrapper code for use by Rust Protobuf.
 
-/// Represents UPB's upb_Arena.
+use std::alloc;
+use std::alloc::Layout;
+use std::cell::UnsafeCell;
+use std::fmt;
+use std::marker::PhantomData;
+use std::mem::MaybeUninit;
 use std::ops::Deref;
 use std::ptr::NonNull;
 use std::slice;
 
+/// See `upb/port/def.inc`.
+const UPB_MALLOC_ALIGN: usize = 8;
+
+/// A UPB-managed pointer to a raw arena.
+pub type RawArena = NonNull<RawArenaData>;
+
+/// The data behind a [`RawArena`]. Do not use this type.
 #[repr(C)]
-pub struct Arena {
+pub struct RawArenaData {
     _data: [u8; 0],
-    _marker: core::marker::PhantomData<(*mut u8, core::marker::PhantomPinned)>,
+}
+
+/// A wrapper over a `upb_Arena`.
+///
+/// This is not a safe wrapper per se, because the allocation functions still
+/// have sharp edges (see their safety docs for more info).
+///
+/// This is an owning type and will automatically free the arena when
+/// dropped.
+///
+/// Note that this type is neither `Sync` nor `Send`.
+pub struct Arena {
+    raw: RawArena,
+    _not_sync: PhantomData<UnsafeCell<()>>,
+}
+
+extern "C" {
+    fn upb_Arena_New() -> RawArena;
+    fn upb_Arena_Free(arena: RawArena);
+    fn upb_Arena_Malloc(arena: RawArena, size: usize) -> *mut u8;
+    fn upb_Arena_Realloc(arena: RawArena, ptr: *mut u8, old: usize, new: usize) -> *mut u8;
 }
 
 impl Arena {
-    pub unsafe fn new() -> *mut Self {
-        upb_Arena_New()
+    /// Allocates a fresh arena.
+    #[inline]
+    pub fn new() -> Self {
+        Self { raw: unsafe { upb_Arena_New() }, _not_sync: PhantomData }
     }
 
-    pub unsafe fn free(arena: *mut Self) {
-        upb_Arena_Free(arena)
+    /// Returns the raw, UPB-managed pointer to the arena.
+    #[inline]
+    pub fn raw(&self) -> RawArena {
+        self.raw
     }
+
+    /// Allocates some memory on the arena.
+    ///
+    /// # Safety
+    ///
+    /// `layout`'s alignment must be less than `UPB_MALLOC_ALIGN`.
+    #[inline]
+    pub unsafe fn alloc(&self, layout: Layout) -> &mut [MaybeUninit<u8>] {
+        debug_assert!(layout.align() <= UPB_MALLOC_ALIGN);
+        let ptr = upb_Arena_Malloc(self.raw, layout.size());
+        if ptr.is_null() {
+            alloc::handle_alloc_error(layout);
         }
 
-extern "C" {
-    pub fn upb_Arena_New() -> *mut Arena;
-    pub fn upb_Arena_Free(arena: *mut Arena);
+        slice::from_raw_parts_mut(ptr.cast(), layout.size())
     }
 
-/// Represents serialized Protobuf wire format data. It's typically produced by
-/// `<Message>.serialize()`.
-#[derive(Debug)]
+    /// Resizes some memory on the arena.
+    ///
+    /// # Safety
+    ///
+    /// After calling this function, `ptr` is essentially zapped. `old` must
+    /// be the layout `ptr` was allocated with via [`Arena::alloc()`]. `new`'s
+    /// alignment must be less than `UPB_MALLOC_ALIGN`.
+    #[inline]
+    pub unsafe fn resize(&self, ptr: *mut u8, old: Layout, new: Layout) -> &[MaybeUninit<u8>] {
+        debug_assert!(new.align() <= UPB_MALLOC_ALIGN);
+        let ptr = upb_Arena_Realloc(self.raw, ptr, old.size(), new.size());
+        if ptr.is_null() {
+            alloc::handle_alloc_error(new);
+        }
+
+        slice::from_raw_parts_mut(ptr.cast(), new.size())
+    }
+}
+
+impl Drop for Arena {
+    #[inline]
+    fn drop(&mut self) {
+        unsafe {
+            upb_Arena_Free(self.raw);
+        }
+    }
+}
+
+/// Represents serialized Protobuf wire format data.
+///
+/// It's typically produced by `<Message>::serialize()`.
 pub struct SerializedData {
     data: NonNull<u8>,
     len: usize,
-    arena: *mut Arena,
+
+    // The arena that owns `data`.
+    _arena: Arena,
 }
 
 impl SerializedData {
-    pub unsafe fn from_raw_parts(arena: *mut Arena, data: NonNull<u8>, len: usize) -> Self {
-        SerializedData { arena, data, len }
+    pub unsafe fn from_raw_parts(arena: Arena, data: NonNull<u8>, len: usize) -> Self {
+        SerializedData { _arena: arena, data, len }
     }
 }
 
@@ -78,9 +154,9 @@ impl Deref for SerializedData {
     }
 }
 
-impl Drop for SerializedData {
-    fn drop(&mut self) {
-        unsafe { Arena::free(self.arena) };
+impl fmt::Debug for SerializedData {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        fmt::Debug::fmt(self.deref(), f)
     }
 }
 
@@ -90,13 +166,13 @@ mod tests {
 
     #[test]
     fn test_arena_new_and_free() {
-        let arena = unsafe { Arena::new() };
-        unsafe { Arena::free(arena) };
+        let arena = Arena::new();
+        drop(arena);
     }
 
     #[test]
     fn test_serialized_data_roundtrip() {
-        let arena = unsafe { Arena::new() };
+        let arena = Arena::new();
         let original_data = b"Hello world";
         let len = original_data.len();
 

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

@@ -97,6 +97,7 @@ bool RustGenerator::Generate(const FileDescriptor* file_desc,
   auto v = file.printer().WithVars({
       {"std", "::__std"},
       {"pb", "::__pb"},
+      {"pbi", "::__pb::__runtime"},
       {"NonNull", "::__std::ptr::NonNull"},
   });
 

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

@@ -56,7 +56,10 @@ void MessageStructFields(Context<Descriptor> msg) {
     case Kernel::kUpb:
       msg.Emit(R"rs(
         msg: $NonNull$<u8>,
-        arena: *mut $pb$::Arena,
+        //~ rustc incorrectly thinks this field is never read, even though
+        //~ it has a destructor!
+        #[allow(dead_code)]
+        arena: $pbi$::Arena,
       )rs");
       return;
   }
@@ -74,9 +77,11 @@ void MessageNew(Context<Descriptor> msg) {
 
     case Kernel::kUpb:
       msg.Emit({{"new_thunk", Thunk(msg, "new")}}, R"rs(
-        let arena = unsafe { $pb$::Arena::new() };
-        let msg = unsafe { $new_thunk$(arena) };
-        $Msg$ { msg, arena }
+        let arena = unsafe { $pbi$::Arena::new() };
+        Self {
+          msg: unsafe { $new_thunk$(arena.raw()) },
+          arena,
+        }
       )rs");
       return;
   }
@@ -94,10 +99,10 @@ void MessageSerialize(Context<Descriptor> msg) {
 
     case Kernel::kUpb:
       msg.Emit({{"serialize_thunk", Thunk(msg, "serialize")}}, R"rs(
-        let arena = unsafe { $pb$::__runtime::upb_Arena_New() };
+        let arena = $pbi$::Arena::new();
         let mut len = 0;
         unsafe {
-          let data = $serialize_thunk$(self.msg, arena, &mut len);
+          let data = $serialize_thunk$(self.msg, arena.raw(), &mut len);
           $pb$::SerializedData::from_raw_parts(arena, data, len)
         }
       )rs");
@@ -163,12 +168,8 @@ void MessageExterns(Context<Descriptor> msg) {
               {"serialize_thunk", Thunk(msg, "serialize")},
           },
           R"rs(
-          fn $new_thunk$(arena: *mut $pb$::Arena) -> $NonNull$<u8>;
-          fn $serialize_thunk$(
-            msg: $NonNull$<u8>,
-            arena: *mut $pb$::Arena,
-            len: &mut usize,
-          ) -> $NonNull$<u8>;
+          fn $new_thunk$(arena: $pbi$::RawArena) -> $NonNull$<u8>;
+          fn $serialize_thunk$(msg: $NonNull$<u8>, arena: $pbi$::RawArena, len: &mut usize) -> $NonNull$<u8>;
       )rs");
       return;
   }
@@ -347,9 +348,9 @@ void MessageGenerator::GenerateThunksCc(Context<Descriptor> msg) {
            }},
       },
       R"cc(
-        // $abi$ is a workaround for a syntax highlight bug in VSCode. However,
-        // that confuses clang-format (it refuses to keep the newline after
-        // `$abi${`). Disabling clang-format for the block.
+        //~ $abi$ is a workaround for a syntax highlight bug in VSCode. However,
+        //~ that confuses clang-format (it refuses to keep the newline after
+        //~ `$abi${`). Disabling clang-format for the block.
         // clang-format off
         extern $abi$ {
         void* $new_thunk$(){return new $QualifiedMsg$(); }