Allocator API taking context

[pitaj]

The rust-for-linux project have their own allocator API including a custom allocator trait. I'm working on some modifications to the Allocator trait that would enable their use-cases. Essentially, what they need is a way to pass some flags at each (possible) allocation. For instance:

v.push(1, GFP_KERNEL)?;

Context as associated type on Allocator

Add an associated type for additional context passed into (re)allocating functions.

pub unsafe trait Allocator {
    type Ctx: Copy;

    fn allocate(&self, context: Self::Ctx, layout: Layout) -> Result<NonNull<[u8]>, AllocError>;

    fn allocate_zeroed(&self, context: Self::Ctx, layout: Layout) -> Result<NonNull<[u8]>, AllocError> { ... }

    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout);

    unsafe fn grow(
        &self,
        context: Self::Ctx,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8]>, AllocError> { ... }
    unsafe fn grow_zeroed(
        &self,
        context: Self::Ctx,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8]>, AllocError> { ... }
    unsafe fn shrink(
        &self,
        context: Self::Ctx,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8]>, AllocError> { ... }

    fn by_ref(&self) -> &Self where Self: Sized { self }
}

// Support storing the context alongside the allocator
unsafe impl<A: Allocator> Allocator for (A, A::Ctx) {
    type Ctx = ();

    fn allocate(&self, _context: (), layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        self.0.allocate(self.1, layout)
    }

    ...
}

Then for each heap container, we would have a copy of each (re)allocating function that takes a generic context:

pub struct Vec<T, A = Global> { ... }

// Normal `push` without context
impl<T, A: Allocator<Ctx = ()>> Vec<T, A> {
    pub fn push(&mut self, elem: T) { ... }
}
// `cpush` with context
impl<T, A: Allocator> Vec<T, A> {
    pub fn cpush(&mut self, elem: T, context: A::Ctx) { ... }
}

rust-for-linux folks would have to use cpush

Could even use #![feature(default_associated_types)] to make Ctx default to (). Though since the allocator API is unstable, this isn't necessary.

Variant: Ctx as a trait generic

You could make Ctx a type parameter of the trait instead:

pub unsafe trait Allocator<Ctx> { ... }

But I fail to see a case where an allocator would want to support multiple types of context.

This would also require that every heap type wanting to generically support allocators taking context would need to add a PhantomData<Ctx> to their struct.

Questions

• Should Ctx be passed by reference?
  Then we could remove the Copy bound, but people could also just set Ctx = &Something.

Problems

1. Duplicated API surface: push, reserve, etc would all need to be duplicated with a context-taking variant

The only way to avoid this is some new language feature, such as making final () arguments optional or calculating disjointness based on associated types.

2. Restricted API surface: Extend, etc can only be implemented for containers where Ctx = ()

This could possibly be mitigated by providing ways to bundle or split up the allocator and context, for instance:

impl<T, A, Ctx> Vec<T, A>
where
    A: Allocator<Ctx = Ctx>,
    (A, Ctx): Allocator<Ctx = ()>,
{
    pub fn to_bundled_context(self, context: Ctx) -> Vec<T, (A, Ctx)>;
}
impl<T, A, Ctx> Vec<T, (A, Ctx)>
where
    A: Allocator<Ctx = Ctx>,
    (A, Ctx): Allocator<Ctx = ()>,
{
    pub fn to_split_context(self) -> (Vec<T, A>, Ctx);
}

// Usage
fn takes_extend<E: Extend<u32>>(e: &mut E);

let mut v = v.to_bundled_context(context);
takes_extend(&mut v);
let (mut v, context) = v.to_split_context();

[pitaj]

There are other uses for this context parameter. For instance, here is an arena allocator that uses lifetime tagging to prevent increasing the size of containers:

// Cell<T> is invariant in T; so Cell<&'id _> makes `id` invariant.
// This means that the inference engine is not allowed to shrink or
// grow 'id to solve the borrow system. 
type Id<'id> = PhantomData<::std::cell::Cell<&'id mut ()>>;

#[derive(Clone, Copy)]
pub struct Child<'id, 'p> {
    _id: Id<'id>,
    _parent: PhantomData<&'p Parent>,
}

pub struct Parent {
    mem: Bump,
}

impl Parent {
    pub fn new() -> Self {
        Self {
            mem: Bump::new()
        }
    }
}

#[derive(Clone, Copy)]
pub struct Context<'id, 'p> {
    _id: Id<'id>,
    parent: &'p Parent,
}

unsafe impl Allocator for Parent {
    type Ctx = ();

    fn allocate(&self, _context: (), layout: std::alloc::Layout) -> Result<std::ptr::NonNull<[u8]>, std::alloc::AllocError> {
        match self.mem.try_alloc_layout(layout) {
            Ok(ptr) => Ok(NonNull::slice_from_raw_parts(ptr, layout.size())),
            Err(_) => Err(AllocError),
        }
    }

    unsafe fn deallocate(&self, _ptr: std::ptr::NonNull<u8>, _layout: std::alloc::Layout) {
        // Deallocation handled by drop of `Parent`
    }
}

unsafe impl<'id, 'p> Allocator for Child<'id, 'p> {
    type Ctx = Context<'id, 'p>;

    fn allocate(&self, context: Context<'id, 'p>, layout: std::alloc::Layout) -> Result<std::ptr::NonNull<[u8]>, std::alloc::AllocError> {
        context.parent.allocate((), layout)
    }

    unsafe fn deallocate(&self, _ptr: std::ptr::NonNull<u8>, _layout: std::alloc::Layout) {
        // Deallocation handled by drop of `Parent`
    }
}

impl Parent {
    pub fn scope<'p, F>(&'p self, f: F)
    where
        F: for<'id> FnOnce(Child<'id, 'p>, Context<'id, 'p>)
    {
        let child = Child { _id: PhantomData, _parent: PhantomData };
        let context = Context { _id: PhantomData, parent: self };
        f(child, context)
    }
}

Usage:

let arena = qalloc::Parent::new();
arena.scope(|child, context| {
    let mut v = Vec::new_in(child);
    v.cpush(123, context);

    let arena2 = qalloc::Parent::new();
    arena2.scope(|child2, context2| {
        v.cpush(456, context);

        let mut v2 = Vec::new_in(child2);
        v2.cpush("abc", context2);

        // v2.push("def", context); // Wrong context -> lifetime error
        // v.push(789, context2); // Wrong context -> lifetime error
    });

    let mut v3 = Vec::new_in(Global);
    // This box is only one pointer wide, since `Child` is a ZST
    // So `v3` uses less heap memory
    v3.push(Box::cnew_in("derp", child, context));
});