Safe aligning of `MaybeUninit` mutable slices

[nickkuk]

Proposal

Problem statement

Add new method for slice:

impl<T> [MaybeUninit<T>] {
    pub fn align_to_uninit_mut<U>(&mut self) -> (&mut Self, &mut [MaybeUninit<U>], &mut Self);
}

Motivating examples or use cases

This method originated from a discussion in users forum: https://users.rust-lang.org/t/typed-scoped-self-referential-arena-in-safe-rust/127155.

It allows to implement scoped multithread typed or untyped arena that allows self-references in safe Rust: gist.

Solution sketch

impl<T> [MaybeUninit<T>] {
    pub fn align_to_uninit_mut<U>(&mut self) -> (&mut Self, &mut [MaybeUninit<U>], &mut Self) {
        unsafe { self.align_to_mut() }
    }
}

Alternatives

I don't know of any so far.

Links and related work

Look at motivating examples.

What happens now?

This issue contains an API change proposal (or ACP) and is part of the libs-api team feature lifecycle. Once this issue is filed, the libs-api team will review open proposals as capability becomes available. Current response times do not have a clear estimate, but may be up to several months.

Possible responses

The libs team may respond in various different ways. First, the team will consider the problem (this doesn't require any concrete solution or alternatives to have been proposed):

• We think this problem seems worth solving, and the standard library might be the right place to solve it.
• We think that this probably doesn't belong in the standard library.

Second, if there's a concrete solution:

• We think this specific solution looks roughly right, approved, you or someone else should implement this. (Further review will still happen on the subsequent implementation PR.)
• We're not sure this is the right solution, and the alternatives or other materials don't give us enough information to be sure about that. Here are some questions we have that aren't answered, or rough ideas about alternatives we'd want to see discussed.

[scottmcm]

Oh, the motivating reason for having this is that it's safe? So you'd use it by MaybeUninit::writeing but never reading from it ever?

[nickkuk]

@scottmcm I apologize for the possibly insufficiently detailed description, but if you look into my gist you can see that MaybeUninit::write into the new method's output is completely enough for using arena's memory inside the scope.

My main motivation is possibility of implementing parallel algorithms like on GPU that use some memory and then flush the entire buffer before reusing it again.

[nickkuk]

Oh, the motivating reason for having this is that it's safe?

Yes, to show people who is implementing arenas that this operations is actually safe. So, the most arenas can be implemented in safe Rust.

So you'd use it by MaybeUninit::writeing but never reading from it ever?

Yes, just holding initialized after MaybeUninit::write mutable references with 'scope lifetime is enough for scoped arena use cases where you don't care about previous memory state during another algorithm run. After a scope ends, it "leaves" the entire buffer in MaybeUninit<u8> state from the point of view of other parts of the program and for subsequent scope.

[hanna-kruppe]

I do think it would be pretty neat to not need any unsafe to implement a basic bump-pointer allocator. I can write a correct one with unsafe and raw pointers, but doing it properly is a huge chore and a risk so it’s rarely a good engineering decision to roll your own (and existing libraries often try to provide more features, which lead to issues like fitzgen/bumpalo#261). I’ve even occasionally had thoughts like “how bad would it be to just use zeroed Box<[u8]> chunks and manually do unaligned loads/stores with $int <-> [u8; N] when necessary” (for simple data consisting mostly or entirely of integers).

[nickkuk]

@hanna-kruppe I'm trying to implement as optimal as possible scoped bump allocator with the following interface:

pub struct BumpAllocator<'scope> { .. }

impl<'scope> BumpAllocator<'scope> {
    pub fn new(memory: &'scope mut [MaybeUninit<u8>]) -> Self { .. }
    pub fn try_alloc_uninit<T>(&mut self) -> Option<&'scope mut MaybeUninit<T>> { .. }
}

Here is godbolt comparison of the best I found with new suggested method vs unsafe implementation from bumpalo. Not as good so far 🙁 Can anyone do better?

It seems we need something like

impl<T> [MaybeUninit<T>] {
    pub fn split_last_aligned<U>(&mut self) -> (&mut MaybeUninit<U>, &mut Self) { .. }
}

to reach such codegen.

[nickkuk]

I made it a bit better by using align_offset: godbolt, but it is still allocation from begin of the buffer. For allocation from end I cannot proof to the compiler that

let end = self.memory.as_ptr_range().end;
end.wrapping_sub((end as usize) & (align_of::<T>() - 1))

is correctly aligned pointer. Maybe we need some operation on pointers like align_offset but for aligning to down?

[hanna-kruppe]

I am not overly worried about bumping up vs down (it’s not like one is always clearly better than the other). Being able to express it at all in safe code, with decent performance, would already be great.

Squeezing out the absolute maximum of performance is a legitimate use case for carefully written unsafe, and the standard library doesn’t need to bend over backwards to remove the need for all such uses. For example, representing a range of memory as (start, end) pointers instead of (start pointer, length) pair is slightly more efficient for some uses, but the standard library doesn’t explicitly provide that alternative representation (slice::Iter* types internally use start/end pointers, but it’s not guaranteed and doesn’t try to to be a general tool for efficiently manipulating start/end pointers pairs in the most efficient way.)

[nickkuk]

Btw, with new split_off_mut method we get quite simple and optimal scoped bump allocator implementation in safe Rust:

pub struct BumpAllocator<'scope> {
    memory: &'scope mut [MaybeUninit<u8>],
}

impl<'scope> BumpAllocator<'scope> {
    pub fn new(memory: &'scope mut [MaybeUninit<u8>]) -> Self {
        Self { memory }
    }
    pub fn try_alloc_uninit<T>(&mut self) -> Option<&'scope mut MaybeUninit<T>> {
        let first_end = self.memory.as_ptr().align_offset(align_of::<T>()) + size_of::<T>();
        let prefix = self.memory.split_off_mut(..first_end)?;
        Some(&mut prefix.align_to_uninit_mut::<T>().1[0])
    }
}

godbolt

[dtolnay]

We discussed this proposal in the most recent standard library API team meeting and are in favor of it, exactly as proposed. Please open a library tracking issue in rust-lang/rust and send the implementation! Thank you.