rust/upb/arena.rs

   1 use crate::opaque_pointee::opaque_pointee;
   2 use std::alloc::{self, Layout};
   3 use std::cell::UnsafeCell;
   4 use std::marker::PhantomData;
   5 use std::mem::{align_of, MaybeUninit};
   6 use std::ptr::{self, NonNull};
   7 use std::slice;
   8
   9 opaque_pointee!(upb_Arena);
  10 pub type RawArena = NonNull<upb_Arena>;
  11
  12 /// See `upb/port/def.inc`.
  13 const UPB_MALLOC_ALIGN: usize = 8;
  14 const _CHECK_UPB_MALLOC_ALIGN_AT_LEAST_POINTER_ALIGNED: () =
  15     assert!(UPB_MALLOC_ALIGN >= align_of::<*const ()>());
  16
  17 /// A wrapper over a `upb_Arena`.
  18 ///
  19 /// This is not a safe wrapper per se, because the allocation functions still
  20 /// have sharp edges (see their safety docs for more info).
  21 ///
  22 /// This is an owning type and will automatically free the arena when
  23 /// dropped.
  24 ///
  25 /// Note that this type is not `Sync` as it implements unsynchronized interior
  26 /// mutability. The upb_Arena C object could be understood as being Sync (at
  27 /// least vacuously under current API since there are not any const upb_Arena*
  28 /// API functions), but the Rust Arena is necessarily expressed as interior
  29 /// mutability (&self rather than &mut self receivers) See https://doc.rust-lang.org/nomicon/lifetime-mismatch.html and
  30 /// https://blog.reverberate.org/2021/12/19/arenas-and-rust.html, and the
  31 /// 'known problems' section of https://rust-lang.github.io/rust-clippy/master/index.html#/mut_from_ref.
  32 #[derive(Debug)]
  33 pub struct Arena {
  34     // Safety invariant: this must always be a valid arena
  35     raw: RawArena,
  36     _not_sync: PhantomData<UnsafeCell<()>>,
  37 }
  38
  39 // SAFETY: `Arena` uniquely holds the underlying RawArena and has no
  40 // thread-local data.
  41 unsafe impl Send for Arena {}
  42
  43 impl Arena {
  44     /// Allocates a fresh arena.
  45     #[inline]
  46     pub fn new() -> Self {
  47         #[inline(never)]
  48         #[cold]
  49         fn arena_new_failed() -> ! {
  50             panic!("Could not create a new UPB arena");
  51         }
  52
  53         // SAFETY:
  54         // - `upb_Arena_New` is assumed to be implemented correctly and always sound to
  55         //   call; if it returned a non-null pointer, it is a valid arena.
  56         unsafe {
  57             let Some(raw) = upb_Arena_New() else { arena_new_failed() };
  58             Self { raw, _not_sync: PhantomData }
  59         }
  60     }
  61
  62     /// # Safety
  63     /// - The `raw_arena` must point to a valid arena.
  64     /// - The caller must ensure that the Arena's destructor does not run.
  65     pub unsafe fn from_raw(raw_arena: RawArena) -> Self {
  66         Arena { raw: raw_arena, _not_sync: PhantomData }
  67     }
  68
  69     /// Returns the raw, UPB-managed pointer to the arena.
  70     #[inline]
  71     pub fn raw(&self) -> RawArena {
  72         self.raw
  73     }
  74
  75     /// Allocates some memory on the arena.
  76     ///
  77     /// # Safety
  78     ///
  79     /// - `layout`'s alignment must be less than `UPB_MALLOC_ALIGN`.
  80     #[allow(clippy::mut_from_ref)]
  81     #[inline]
  82     pub unsafe fn alloc(&self, layout: Layout) -> &mut [MaybeUninit<u8>] {
  83         debug_assert!(layout.align() <= UPB_MALLOC_ALIGN);
  84         // SAFETY: `self.raw` is a valid UPB arena
  85         let ptr = unsafe { upb_Arena_Malloc(self.raw, layout.size()) };
  86         if ptr.is_null() {
  87             alloc::handle_alloc_error(layout);
  88         }
  89
  90         // SAFETY:
  91         // - `upb_Arena_Malloc` promises that if the return pointer is non-null, it is
  92         //   dereferencable for `size` bytes and has an alignment of `UPB_MALLOC_ALIGN`
  93         //   until the arena is destroyed.
  94         // - `[MaybeUninit<u8>]` has no alignment requirement, and `ptr` is aligned to a
  95         //   `UPB_MALLOC_ALIGN` boundary.
  96         unsafe { slice::from_raw_parts_mut(ptr.cast(), layout.size()) }
  97     }
  98
  99     /// Same as alloc() but panics if `layout.align() > UPB_MALLOC_ALIGN`.
 100     #[allow(clippy::mut_from_ref)]
 101     #[inline]
 102     pub fn checked_alloc(&self, layout: Layout) -> &mut [MaybeUninit<u8>] {
 103         assert!(layout.align() <= UPB_MALLOC_ALIGN);
 104         // SAFETY: layout.align() <= UPB_MALLOC_ALIGN asserted.
 105         unsafe { self.alloc(layout) }
 106     }
 107
 108     /// Copies the T into this arena and returns a pointer to the T data inside
 109     /// the arena.
 110     pub fn copy_in<'a, T: Copy>(&'a self, data: &T) -> &'a T {
 111         let layout = Layout::for_value(data);
 112         let alloc = self.checked_alloc(layout);
 113
 114         // SAFETY:
 115         // - alloc is valid for `layout.len()` bytes and is the uninit bytes are written
 116         //   to not read from until written.
 117         // - T is copy so copying the bytes of the value is sound.
 118         unsafe {
 119             let alloc = alloc.as_mut_ptr().cast::<MaybeUninit<T>>();
 120             // let data = (data as *const T).cast::<MaybeUninit<T>>();
 121             (*alloc).write(*data)
 122         }
 123     }
 124
 125     pub fn copy_str_in<'a>(&'a self, s: &str) -> &'a str {
 126         let copied_bytes = self.copy_slice_in(s.as_bytes());
 127         // SAFETY: `copied_bytes` has same contents as `s` and so must meet &str
 128         // criteria.
 129         unsafe { std::str::from_utf8_unchecked(copied_bytes) }
 130     }
 131
 132     pub fn copy_slice_in<'a, T: Copy>(&'a self, data: &[T]) -> &'a [T] {
 133         let layout = Layout::for_value(data);
 134         let alloc: *mut T = self.checked_alloc(layout).as_mut_ptr().cast();
 135
 136         // SAFETY:
 137         // - uninit_alloc is valid for `layout.len()` bytes and is the uninit bytes are
 138         //   written to not read from until written.
 139         // - T is copy so copying the bytes of the values is sound.
 140         unsafe {
 141             ptr::copy_nonoverlapping(data.as_ptr(), alloc, data.len());
 142             slice::from_raw_parts_mut(alloc, data.len())
 143         }
 144     }
 145 }
 146
 147 impl Default for Arena {
 148     fn default() -> Self {
 149         Self::new()
 150     }
 151 }
 152
 153 impl Drop for Arena {
 154     #[inline]
 155     fn drop(&mut self) {
 156         unsafe {
 157             upb_Arena_Free(self.raw);
 158         }
 159     }
 160 }
 161
 162 extern "C" {
 163     // `Option<NonNull<T: Sized>>` is ABI-compatible with `*mut T`
 164     fn upb_Arena_New() -> Option<RawArena>;
 165     fn upb_Arena_Free(arena: RawArena);
 166     fn upb_Arena_Malloc(arena: RawArena, size: usize) -> *mut u8;
 167 }
 168
 169 #[cfg(test)]
 170 mod tests {
 171     use super::*;
 172
 173     #[test]
 174     fn raw_ffi_test() {
 175         // SAFETY: FFI unit test uses C API under expected patterns.
 176         unsafe {
 177             let arena = upb_Arena_New().unwrap();
 178             let bytes = upb_Arena_Malloc(arena, 3);
 179             *bytes.add(2) = 7;
 180             upb_Arena_Free(arena);
 181         }
 182     }
 183
 184     #[test]
 185     fn test_arena_new_and_free() {
 186         let arena = Arena::new();
 187         drop(arena);
 188     }
 189 }