Mutable references returned by custom allocator

I've been playing around with custom allocators lately and came up with the following function. I want to know if my implementation is correct or if it triggers UB. Im kind of torn about this.

pub fn slice_mut_uninit<'a, T: Sized>(&'a self, size: usize) -> &'a mut [MaybeUninit<T>] {
      let layout = alloc::Layout::array::<T>(size).unwrap();
      let ptr = self.allocate(layout).unwrap();
      let s = unsafe { std::slice::from_raw_parts_mut(ptr.cast().as_ptr(), size) };
      s
  }

The idea behind this function is to return mutable slice with it's lifetime tied to that of the allocator itself. Returned slice does not alias Self, but returning &mut from function taking shared reference still feels wrong somehow.

RefCell::borrow_mut does almost exactly the same thing: the lifetime in RefMut<'b, T> is tied to the original &'b RefCell<T>. But there you must go one step further from &'a mut RefMut<'b, T>to &'a mut T.

In your custom allocator is there a way to free the slice without dropping the allocator?

The slice cannot be freed on its own, but since this is an arena allocator, it is possible to reset it to its initial state and free it that way. My reset function takes &mut self so it's not possible to free memory while there are any outstanding references.

This does happen with allocators. It is kind of uncommon, though. See bumpalo::Bump::alloc() and typed_arena::Arena::alloc().

It is more common to return an index or reference-type smart pointer (e.g. RefMut). There's a handy table that compares arena implementations on this Arenas blog post published 2024-08-15.

So long as your type is making use of proper internal mutability constructs, and accounts for multithreading, you should be good; but it's hard to tell from just that snippet. There could also be some provenance concerns, but they're rather unlikely here.