Move the Rust `Arena` type which wraps a upb_Arena into the upb directory.

PiperOrigin-RevId: 625317910

rust/upb.rs

@@ -13,12 +13,9 @@ use crate::{
     Repeated, RepeatedMut, RepeatedView, SettableValue, View, ViewProxy,
 };
 use core::fmt::Debug;
-use std::alloc;
 use std::alloc::Layout;
-use std::cell::UnsafeCell;
 use std::fmt;
-use std::marker::PhantomData;
-use std::mem::{align_of, size_of, ManuallyDrop, MaybeUninit};
+use std::mem::{size_of, ManuallyDrop, MaybeUninit};
 use std::ops::Deref;
 use std::ptr::{self, NonNull};
 use std::slice;
@@ -43,123 +40,6 @@ impl From<&ProtoStr> for PtrAndLen {
     }
 }
 
-/// See `upb/port/def.inc`.
-const UPB_MALLOC_ALIGN: usize = 8;
-const _CHECK_UPB_MALLOC_ALIGN_AT_LEAST_POINTER_ALIGNED: () =
-    assert!(UPB_MALLOC_ALIGN >= align_of::<*const ()>());
-
-/// 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`.
-#[derive(Debug)]
-pub struct Arena {
-    // Safety invariant: this must always be a valid arena
-    raw: RawArena,
-    _not_sync: PhantomData<UnsafeCell<()>>,
-}
-
-impl Arena {
-    /// Allocates a fresh arena.
-    #[inline]
-    pub fn new() -> Self {
-        #[inline(never)]
-        #[cold]
-        fn arena_new_failed() -> ! {
-            panic!("Could not create a new UPB arena");
-        }
-
-        // SAFETY:
-        // - `upb_Arena_New` is assumed to be implemented correctly and always sound to
-        //   call; if it returned a non-null pointer, it is a valid arena.
-        unsafe {
-            let Some(raw) = upb_Arena_New() else { arena_new_failed() };
-            Self { raw, _not_sync: PhantomData }
-        }
-    }
-
-    /// # Safety
-    /// - The `raw_arena` must point to a valid arena.
-    /// - The caller must ensure that the Arena's destructor does not run.
-    unsafe fn from_raw(raw_arena: RawArena) -> Self {
-        Arena { raw: raw_arena, _not_sync: PhantomData }
-    }
-
-    /// 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);
-        // SAFETY: `self.raw` is a valid UPB arena
-        let ptr = unsafe { upb_Arena_Malloc(self.raw, layout.size()) };
-        if ptr.is_null() {
-            alloc::handle_alloc_error(layout);
-        }
-
-        // SAFETY:
-        // - `upb_Arena_Malloc` promises that if the return pointer is non-null, it is
-        //   dereferencable for `size` bytes and has an alignment of `UPB_MALLOC_ALIGN`
-        //   until the arena is destroyed.
-        // - `[MaybeUninit<u8>]` has no alignment requirement, and `ptr` is aligned to a
-        //   `UPB_MALLOC_ALIGN` boundary.
-        unsafe { slice::from_raw_parts_mut(ptr.cast(), layout.size()) }
-    }
-
-    /// Resizes some memory on the arena.
-    ///
-    /// # Safety
-    ///
-    /// - `ptr` must be the data pointer returned by a previous call to `alloc`
-    ///   or `resize` on `self`.
-    /// - After calling this function, `ptr` is no longer dereferencable - it is
-    ///   zapped.
-    /// - `old` must be the layout `ptr` was allocated with via `alloc` or
-    ///   `realloc`.
-    /// - `new`'s alignment must be less than `UPB_MALLOC_ALIGN`.
-    #[inline]
-    pub unsafe fn resize(&self, ptr: *mut u8, old: Layout, new: Layout) -> &mut [MaybeUninit<u8>] {
-        debug_assert!(new.align() <= UPB_MALLOC_ALIGN);
-        // SAFETY:
-        // - `self.raw` is a valid UPB arena
-        // - `ptr` was allocated by a previous call to `alloc` or `realloc` as promised
-        //   by the caller.
-        let ptr = unsafe { upb_Arena_Realloc(self.raw, ptr, old.size(), new.size()) };
-        if ptr.is_null() {
-            alloc::handle_alloc_error(new);
-        }
-
-        // SAFETY:
-        // - `upb_Arena_Realloc` promises that if the return pointer is non-null, it is
-        //   dereferencable for the new `size` in bytes until the arena is destroyed.
-        // - `[MaybeUninit<u8>]` has no alignment requirement, and `ptr` is aligned to a
-        //   `UPB_MALLOC_ALIGN` boundary.
-        unsafe { slice::from_raw_parts_mut(ptr.cast(), new.size()) }
-    }
-}
-
-impl Drop for Arena {
-    #[inline]
-    fn drop(&mut self) {
-        unsafe {
-            upb_Arena_Free(self.raw);
-        }
-    }
-}
-
 /// The scratch size of 64 KiB matches the maximum supported size that a
 /// upb_Message can possibly be.
 const UPB_SCRATCH_SPACE_BYTES: usize = 65_536;
@@ -934,12 +814,6 @@ mod tests {
     use super::*;
     use googletest::prelude::*;
 
-    #[test]
-    fn test_arena_new_and_free() {
-        let arena = Arena::new();
-        drop(arena);
-    }
-
     #[test]
     fn test_serialized_data_roundtrip() {
         let arena = Arena::new();

rust/upb/BUILD

@@ -12,7 +12,11 @@ load("@rules_rust//rust:defs.bzl", "rust_library", "rust_test")
 
 rust_library(
     name = "upb",
-    srcs = ["shared.rs"],
+    srcs = [
+        "arena.rs",
+        "lib.rs",
+        "opaque_pointee.rs",
+    ],
     visibility = [
         "//rust:__subpackages__",
         "//src/google/protobuf:__subpackages__",

rust/upb/arena.rs

@@ -0,0 +1,138 @@
+use crate::opaque_pointee::opaque_pointee;
+use std::alloc::{self, Layout};
+use std::cell::UnsafeCell;
+use std::marker::PhantomData;
+use std::mem::{align_of, MaybeUninit};
+use std::ptr::NonNull;
+use std::slice;
+
+opaque_pointee!(upb_Arena);
+pub type RawArena = NonNull<upb_Arena>;
+
+/// See `upb/port/def.inc`.
+const UPB_MALLOC_ALIGN: usize = 8;
+const _CHECK_UPB_MALLOC_ALIGN_AT_LEAST_POINTER_ALIGNED: () =
+    assert!(UPB_MALLOC_ALIGN >= align_of::<*const ()>());
+
+/// 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`. The upb_Arena C object
+/// could be understood as being Sync (at least vacuously, as there are no
+/// const-ptr functions on upb_Arena's API), but the Rust Arena is
+/// necessarily expressed as interior mutability (&self rather than &mut self
+/// receivers) See https://doc.rust-lang.org/nomicon/lifetime-mismatch.html and
+/// https://blog.reverberate.org/2021/12/19/arenas-and-rust.html, and the
+/// 'known problems' section of https://rust-lang.github.io/rust-clippy/master/index.html#/mut_from_ref.
+#[derive(Debug)]
+pub struct Arena {
+    // Safety invariant: this must always be a valid arena
+    raw: RawArena,
+    _not_sync: PhantomData<UnsafeCell<()>>,
+}
+
+impl Arena {
+    /// Allocates a fresh arena.
+    #[inline]
+    pub fn new() -> Self {
+        #[inline(never)]
+        #[cold]
+        fn arena_new_failed() -> ! {
+            panic!("Could not create a new UPB arena");
+        }
+
+        // SAFETY:
+        // - `upb_Arena_New` is assumed to be implemented correctly and always sound to
+        //   call; if it returned a non-null pointer, it is a valid arena.
+        unsafe {
+            let Some(raw) = upb_Arena_New() else { arena_new_failed() };
+            Self { raw, _not_sync: PhantomData }
+        }
+    }
+
+    /// # Safety
+    /// - The `raw_arena` must point to a valid arena.
+    /// - The caller must ensure that the Arena's destructor does not run.
+    pub unsafe fn from_raw(raw_arena: RawArena) -> Self {
+        Arena { raw: raw_arena, _not_sync: PhantomData }
+    }
+
+    /// 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`.
+    #[allow(clippy::mut_from_ref)]
+    #[inline]
+    pub unsafe fn alloc(&self, layout: Layout) -> &mut [MaybeUninit<u8>] {
+        debug_assert!(layout.align() <= UPB_MALLOC_ALIGN);
+        // SAFETY: `self.raw` is a valid UPB arena
+        let ptr = unsafe { upb_Arena_Malloc(self.raw, layout.size()) };
+        if ptr.is_null() {
+            alloc::handle_alloc_error(layout);
+        }
+
+        // SAFETY:
+        // - `upb_Arena_Malloc` promises that if the return pointer is non-null, it is
+        //   dereferencable for `size` bytes and has an alignment of `UPB_MALLOC_ALIGN`
+        //   until the arena is destroyed.
+        // - `[MaybeUninit<u8>]` has no alignment requirement, and `ptr` is aligned to a
+        //   `UPB_MALLOC_ALIGN` boundary.
+        unsafe { slice::from_raw_parts_mut(ptr.cast(), layout.size()) }
+    }
+}
+
+impl Default for Arena {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Drop for Arena {
+    #[inline]
+    fn drop(&mut self) {
+        unsafe {
+            upb_Arena_Free(self.raw);
+        }
+    }
+}
+
+extern "C" {
+    // `Option<NonNull<T: Sized>>` is ABI-compatible with `*mut T`
+    fn upb_Arena_New() -> Option<RawArena>;
+    fn upb_Arena_Free(arena: RawArena);
+    fn upb_Arena_Malloc(arena: RawArena, size: usize) -> *mut u8;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn raw_ffi_test() {
+        // SAFETY: FFI unit test uses C API under expected patterns.
+        unsafe {
+            let arena = upb_Arena_New().unwrap();
+            let bytes = upb_Arena_Malloc(arena, 3);
+            *bytes.add(2) = 7;
+            upb_Arena_Free(arena);
+        }
+    }
+
+    #[test]
+    fn test_arena_new_and_free() {
+        let arena = Arena::new();
+        drop(arena);
+    }
+}

rust/upb/lib.rs

@@ -1,23 +1,11 @@
 use std::ptr::NonNull;
 use std::slice;
 
-// Macro to create structs that will act as opaque pointees. These structs are
-// never intended to be dereferenced in Rust.
-// This 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)>,
-        }
-    };
-}
+mod opaque_pointee;
+use opaque_pointee::opaque_pointee;
 
-opaque_pointee!(upb_Arena);
-pub type RawArena = NonNull<upb_Arena>;
+mod arena;
+pub use arena::{upb_Arena, Arena, RawArena};
 
 opaque_pointee!(upb_Message);
 pub type RawMessage = NonNull<upb_Message>;
@@ -34,14 +22,6 @@ pub type RawMap = NonNull<upb_Map>;
 opaque_pointee!(upb_Array);
 pub type RawArray = NonNull<upb_Array>;
 
-extern "C" {
-    // `Option<NonNull<T: Sized>>` is ABI-compatible with `*mut T`
-    pub fn upb_Arena_New() -> Option<RawArena>;
-    pub fn upb_Arena_Free(arena: RawArena);
-    pub fn upb_Arena_Malloc(arena: RawArena, size: usize) -> *mut u8;
-    pub fn upb_Arena_Realloc(arena: RawArena, ptr: *mut u8, old: usize, new: usize) -> *mut u8;
-}
-
 extern "C" {
     pub fn upb_Message_DeepCopy(
         dst: RawMessage,
@@ -253,35 +233,22 @@ extern "C" {
 mod tests {
     use super::*;
 
-    #[test]
-    fn arena_ffi_test() {
-        // SAFETY: FFI unit test uses C API under expected patterns.
-        unsafe {
-            let arena = upb_Arena_New().unwrap();
-            let bytes = upb_Arena_Malloc(arena, 3);
-            let bytes = upb_Arena_Realloc(arena, bytes, 3, 512);
-            *bytes.add(511) = 7;
-            upb_Arena_Free(arena);
-        }
-    }
-
     #[test]
     fn array_ffi_test() {
         // SAFETY: FFI unit test uses C API under expected patterns.
         unsafe {
-            let arena = upb_Arena_New().unwrap();
-            let array = upb_Array_New(arena, CType::Float);
+            let arena = Arena::new();
+            let raw_arena = arena.raw();
+            let array = upb_Array_New(raw_arena, CType::Float);
 
-            assert!(upb_Array_Append(array, upb_MessageValue { float_val: 7.0 }, arena));
-            assert!(upb_Array_Append(array, upb_MessageValue { float_val: 42.0 }, arena));
+            assert!(upb_Array_Append(array, upb_MessageValue { float_val: 7.0 }, raw_arena));
+            assert!(upb_Array_Append(array, upb_MessageValue { float_val: 42.0 }, raw_arena));
             assert_eq!(upb_Array_Size(array), 2);
             assert!(matches!(upb_Array_Get(array, 1), upb_MessageValue { float_val: 42.0 }));
 
-            assert!(upb_Array_Resize(array, 3, arena));
+            assert!(upb_Array_Resize(array, 3, raw_arena));
             assert_eq!(upb_Array_Size(array), 3);
             assert!(matches!(upb_Array_Get(array, 2), upb_MessageValue { float_val: 0.0 }));
-
-            upb_Arena_Free(arena);
         }
     }
 
@@ -289,15 +256,16 @@ mod tests {
     fn map_ffi_test() {
         // SAFETY: FFI unit test uses C API under expected patterns.
         unsafe {
-            let arena = upb_Arena_New().unwrap();
-            let map = upb_Map_New(arena, CType::Bool, CType::Double);
+            let arena = Arena::new();
+            let raw_arena = arena.raw();
+            let map = upb_Map_New(raw_arena, CType::Bool, CType::Double);
             assert_eq!(upb_Map_Size(map), 0);
             assert_eq!(
                 upb_Map_Insert(
                     map,
                     upb_MessageValue { bool_val: true },
                     upb_MessageValue { double_val: 2.0 },
-                    arena
+                    raw_arena
                 ),
                 MapInsertStatus::Inserted
             );
@@ -306,7 +274,7 @@ mod tests {
                     map,
                     upb_MessageValue { bool_val: true },
                     upb_MessageValue { double_val: 3.0 },
-                    arena,
+                    raw_arena,
                 ),
                 MapInsertStatus::Replaced,
             );
@@ -318,7 +286,7 @@ mod tests {
                     map,
                     upb_MessageValue { bool_val: true },
                     upb_MessageValue { double_val: 4.0 },
-                    arena
+                    raw_arena
                 ),
                 MapInsertStatus::Inserted
             );
@@ -326,8 +294,6 @@ mod tests {
             let mut out = upb_MessageValue::zeroed();
             assert!(upb_Map_Get(map, upb_MessageValue { bool_val: true }, &mut out));
             assert!(matches!(out, upb_MessageValue { double_val: 4.0 }));
-
-            upb_Arena_Free(arena);
         }
     }
 }

rust/upb/opaque_pointee.rs

@@ -0,0 +1,16 @@
+// Macro to create structs that will act as opaque pointees. These structs are
+// never intended to be dereferenced in Rust.
+// This 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)>,
+        }
+    };
+}
+
+pub(crate) use opaque_pointee;