Implement AllocRef for references to A: AllocRef

[lachlansneff]

For usability purposes, I'm proposing that the following implementation be added:

unsafe impl<A> AllocRef for &'_ mut A where A: AllocRef {
    fn alloc(&mut self, layout: Layout, init: AllocInit) -> Result<MemoryBlock, AllocErr> {
        A::alloc(self, layout, init)
    }

    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
        A::dealloc(self, ptr, layout)
    }

    unsafe fn grow(&mut self, ptr: NonNull<u8>, layout: Layout, new_size: usize, placement: ReallocPlacement, init: AllocInit) -> Result<MemoryBlock, AllocErr> {
        A::grow(self,ptr, layout, new_size, placement, init)
    }

    unsafe fn shrink(&mut self, ptr: NonNull<u8>, layout: Layout, new_size: usize, placement: ReallocPlacement) -> Result<MemoryBlock, AllocErr> {
        A::shrink(self, ptr, layout, new_size, placement)
    }
}

Without this implementation, I believe it's impossible to implement composable allocators, as I have started doing in this gist: https://gist.github.com/lachlansneff/8f59278bccb82ec5b36c359f385b7e3f.

[Amanieu]

I think this is fine to add.

[TimDiekmann]

@lachlansneff Do you want to push this upstream should I take care if it?

[TimDiekmann]

However, I don't see, why composable allocators can't be implemented with the above definition?

Your linked gist has a bug in Segregator: alloc may return a larger size than Threshold. Additionally, you have to overwrite grow and shrink, otherwise your implementation will be unsound. I'm experimenting with composable allocators in alloc-compose and I didn't pushed SegregateAlloc so far because of this. This is my current implementation:

// reused in multiple allocators
unsafe fn grow<A1: AllocRef, A2: AllocRef>(
    a1: &mut A1,
    a2: &mut A2,
    ptr: NonNull<u8>,
    layout: Layout,
    new_size: usize,
    placement: ReallocPlacement,
    init: AllocInit,
) -> Result<MemoryBlock, AllocErr> {
    if placement == ReallocPlacement::MayMove {
        let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
        let new_memory = a2.alloc(new_layout, init)?;
        ptr::copy_nonoverlapping(ptr.as_ptr(), new_memory.ptr.as_ptr(), layout.size());
        a1.dealloc(ptr, layout);
        Ok(new_memory)
    } else {
        Err(AllocErr)
    }
}

// reused in multiple allocators
unsafe fn shrink<A1: AllocRef, A2: AllocRef>(
    a1: &mut A1,
    a2: &mut A2,
    ptr: NonNull<u8>,
    layout: Layout,
    new_size: usize,
    placement: ReallocPlacement,
) -> Result<MemoryBlock, AllocErr> {
    if placement == ReallocPlacement::MayMove {
        let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
        let new_memory = a2.alloc(new_layout, AllocInit::Uninitialized)?;
        ptr::copy_nonoverlapping(ptr.as_ptr(), new_memory.ptr.as_ptr(), new_memory.size);
        a1.dealloc(ptr, layout);
        Ok(new_memory)
    } else {
        Err(AllocErr)
    }
}

#[derive(Debug, Copy, Clone)]
pub struct SegregateAlloc<Small, Large> {
    threshold: usize,
    pub small: Small,
    pub large: Large,
}

impl<Small: AllocRef, Large: AllocRef> SegregateAlloc<Small, Large> {
    fn clamp_memory(&self, memory: MemoryBlock) -> MemoryBlock {
        MemoryBlock {
            ptr: memory.ptr,
            size: cmp::max(memory.size, self.threshold),
        }
    }
}

unsafe impl<Small, Large> AllocRef for SegregateAlloc<Small, Large>
where
    Small: AllocRef,
    Large: AllocRef,
{
    fn alloc(&mut self, layout: Layout, init: AllocInit) -> Result<MemoryBlock, AllocErr> {
        if layout.size() <= self.threshold {
            let memory = self.small.alloc(layout, init)?;
            // memory block must be smaller than threshold
            Ok(self.clamp_memory(memory))
        } else {
            self.large.alloc(layout, init)
        }
    }

    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
        if layout.size() <= self.threshold {
            self.small.dealloc(ptr, layout)
        } else {
            self.large.dealloc(ptr, layout)
        }
    }

    unsafe fn grow(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
        placement: ReallocPlacement,
        init: AllocInit,
    ) -> Result<MemoryBlock, AllocErr> {
        if layout.size() <= self.threshold {
            let memory = if new_size > self.threshold {
                // Move ownership to `self.large`
                grow(
                    &mut self.small,
                    &mut self.large,
                    ptr,
                    layout,
                    new_size,
                    placement,
                    init,
                )?
            } else {
                self.small.grow(ptr, layout, new_size, placement, init)?
            };
            Ok(self.clamp_memory(memory))
        } else {
            self.large.grow(ptr, layout, new_size, placement, init)
        }
    }

    unsafe fn shrink(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
        placement: ReallocPlacement,
    ) -> Result<MemoryBlock, AllocErr> {
        let memory = if layout.size() <= self.threshold {
            let memory = self.small.shrink(ptr, layout, new_size, placement)?;
            Ok(self.clamp_memory(memory))
        } else if new_size <= self.threshold {
            // Move ownership to `self.small`
            let memory = shrink(
                &mut self.large,
                &mut self.small,
                ptr,
                layout,
                new_size,
                placement,
            )?;
            Ok(self.clamp_memory(memory))
        } else {
            self.large.shrink(ptr, layout, new_size, placement)
        }
    }
}

This needs more testing before I can push this.

[lachlansneff]

@TimDiekmann Well, from my understanding, though I may have missed something, just having two generics with AllocRef bounds means you can't just build up single type with all the allocators in it.

For example, a stack allocator must stay in a specific location, so the part that implements AllocRef would be &StackAlloc<N> or &mut StackAlloc<N>, not StackAlloc<N>. Therefore, you couldn't have this: type Alloc = Fallback<StackAlloc<N>, Global>.

Implementing AllocRef for &mut AllocRef (or &AllocRef might be better, not sure) lets your allocator implementations take something that is either an AllocRef itself (like Global) or something who's reference implements AllocRef (like StackAlloc<N>). There are some complicated trait bounds to get that to work, but my example has them.

As for the unsoundness, I'm not surprised. I wrote it in about an hour and wasn't paying too much attention to details.

[TimDiekmann]

You are right, that you must not implement AllocRef for StackAlloc<N> directly, but for &mut StackAlloc<N> or &StackAlloc<N>. But it's possible to define type Alloc<'a> = FallbackAlloc<&'a StackAlloc<N>, Global>;.

Implementing your proposal makes sense anyway as &mut A may be smaller than A in size. I'd also add AllocRef::by_ref just like in Write::by_ref for a convenient construction.

[lachlansneff]

Yes, you can define an Allocator that way, but it's certainly nice to be able to store everything inline if you want.

let allocator = FallBack<StackAlloc<N>, Global>::default();

There may also be cache-locality benefits from making the allocator state very compact and inline.

[TimDiekmann]

That would require &mut StackAlloc<N>: Default which is not possible, as you need to dereference to an actual object.

There may also be cache-locality benefits from making the allocator state very compact and inline.

I don't see how this would be more efficient than

let mut stack_alloc = StackAlloc::new();
let alloc = FallBack {
    primary: &mut stack_alloc,
    fallback: Global
};

Actually, StackAlloc is a bad example for AllocRef for &mut A, as StackAlloc does not even implement AllocRef.

[lachlansneff]

Well, I actually implemented it in my example, so it's definitely possible. Default is just only implemented when the generic parameters have Default, so it works when the parameter allocators are inline.

I don't see how this would be more efficient than...

Well, for stack allocator it's probably not. But with more complex allocators keeping all the state in single struct is definitely more efficient both because pointer chasing can add overhead and also because being able to fit the entire allocator in a few cache lines as possible can really speed things up.

Actually, StackAlloc is a bad example for AllocRef for &mut A, as StackAlloc does not even implement AllocRef.

Sure, but if you want to be able to take both Global and StackAlloc (which doesn't implement AllocRef, as you noted), you have to assume the references to them will have AllocRef.

[TimDiekmann]

Default is just only implemented when the generic parameters have Default

Default is implemented on the StackAlloc, not on &mut StackAlloc. You cannot implement Default for a struct containing mutable references:

// Either derive Default
#[derive(Default)]
struct TakesRef<'a, T> {
    t: &'a mut T,
//     ^^^^^^^^^ the trait `std::default::Default` is not implemented for `&mut T`
}
    
// Or implement it manually
impl<T: Default> Default for TakesRef<'_, T> {
    fn default() -> Self {
        Self {
            t: &mut T::default()
//                  ------------ temporary value created here
        }
    }
}
}

Well, for stack allocator it's probably not. But with more complex allocators keeping all the state in single struct is definitely more efficient both because pointer chasing can add overhead and also because being able to fit the entire allocator in a few cache lines as possible can really speed things up.

I still don't see, why you want to force the user to use &mut A instead of A. Let the user chose, if they want to use A or A::by_ref.

Sure, but if you want to be able to take both Global and StackAlloc (which doesn't implement AllocRef, as you noted), you have to assume the references to them will have AllocRef.

Then you just pass &(mut) StackAlloc and you wil be fine, as &(mut) StackAlloc implements AllocRef. The reference is part of the type you pass as generic parameter.

[lachlansneff]

My whole thing is about giving the user choice here.

Here's the Fallback implementation:

#[derive(Default)]
pub struct Fallback<Primary, Secondary> {
    primary: Primary,
    secondary: Secondary,
}

unsafe impl<Primary, Secondary> AllocRef for &'_ mut Fallback<Primary, Secondary>
where
    for<'a> &'a mut Primary: AllocRef,
    for<'a> &'a Primary: AllocOwner,
    for<'a> &'a mut Secondary: AllocRef,
{
    // ...
}

If you require that Primary and Secondary implement AllocRef, the user cannot pass in a allocator type that stores the data within itself.

You can use this either way!

let alloc = Fallback::<StackAlloc<N>, Global>::default()

or

let mut stack = StackAlloc<N>::default();
let alloc: Fallback<&mut StackAlloc<N>, Global> = Fallback::new(&mut stack, Global);

This gives the user choice to do either.

[TimDiekmann]

Then I still don't see, why &mut A is required in your usecase, has your implementation any advantages over this?

[TimDiekmann]

I think this is pretty offtopic here, we may move the discussion to your gist or to alloc-compose instead 🙂