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));
});

[jmillikin]

Minor bike-shedding: would prefer Context vs Ctx for the type name, matching similar idioms in other languages.

Regarding the API surface problems, a second problem with using the existing simple methods (push, append, etc) in the Linux kernel is that they panic on allocation error. Adding try_* variants will unavoidably expand the API, so might as well define them as try_push(&mut self, ctx: A::Context, value: T) -> Result<(), AllocError> to get both features for the same cost.

v.try_push(GFP_KERNEL, 1)?;

The existing try_reserve and try_reserve_exact methods would be stuck with a different error type (unless TryReserveError is aliased to AllocError?) and wouldn't have a parameter for the context. Could define try_reserve_ctx and try_reserve_ctx_exact?

Another option is to have a method shaped like with_alloc_ctx(&mut self, ctx: A::Context) -> VecWithAllocCtx<'_, Self> and having that type implement push() etc with the new signatures. Such a type could also support extend(). The resulting code would look something like this:

v.with_alloc_ctx(GFP_KERNEL).push(1)?;

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

If the methods took their allocator context by reference then wouldn't the call sites in the kernel need an extra & everywhere? It looks a bit awkward in the simple case, and for big context types doesn't provide much benefit over Context = &BigFlags.

// https://rust.docs.kernel.org/kernel/alloc/flags/index.html
pub struct Flags { ... }
pub const GFP_KERNEL: Flags;

// ctx: A::Context where A::Context: Copy
v.try_push(GFP_KERNEL, 1)?;

// ctx: &A::Context
v.try_push(&GFP_KERNEL, 1)?;

[aapoalas]

Seconding @pitaj 's example, a context parameter can be used to kind of turn the entire allocation logic on its head: the allocator parameter on Vec / Box ceases to actually be a way to find the allocator and perform allocation/deallocation; it becomes simply the way to make the allocator invariant (although this isn't strictly necessary; if the allocator is fine with getting called with pointers that it didn't allocate then the allocator can be covariant [I think]) and Context becomes the actual way to pass in that allocator.

This is a very useful feature for bump and slab allocators (or so I believe) since they'd want to avoid storing an allocator reference in each (heap-allocated) Vec / Box. They'd thus prefer to pass the concrete allocator in the context, but then use branding to ensure that the context parameter always matches the allocator (or in the "turn on its head" parlance, that the allocator parameter matches the object's brand).

[aapoalas]

Another option is to have a method shaped like with_alloc_ctx(&mut self, ctx: A::Context) -> VecWithAllocCtx<'_, Self> and having that type implement push() etc with the new signatures. Such a type could also support extend(). The resulting code would look something like this:

This seems like a fruitfulinteresting path to consider. One could think of Vec and Box as being Vec/BoxWithAllocCtx types already; they believe they can call the allocators methods directly on the self.alloc field. A Vec or Box that takes an allocator with Context would be a "lesser" type that would only have eg. try_reserve and try_push like APIs will never allocate. You'd then need to call vec.with_alloc_ctx(ctx) to get a full-fledged Vec which could now actually do allocations.

There's an obvious problem here of course; how do you represent the "lesser" VecWithoutAlloc temporarily becoming a normal Vec, and then going back to being a VecWithoutAlloc once you're done with your allocations? That doesn't seem possible, since the two would need to be at the same memory location for mutations on the Vec to carry over into the VecWithoutAlloc. It seems you'd still need to have that VecWithAllocCtx type which internally is an exclusive(?) reference to the VecWithoutAllocCtx plus the context.