Tracking issue for custom allocators in standard collections

[alexcrichton]

The new Alloc trait defined in RFC 1398 was added in #42313 but the libs teamdecided to hold off on Vec integration just yet to be conservative.

This issue is intended to track the integration of the Alloc trait into standard collections like Vec, BTreeMap, etc. I'm not personally aware of any current blockers, but what we'll probably want to do as part of this issue includes:

• Develop a set of conventions for collections which support custom allocators
• Ensure no regressions with integration of a custom allocator:
  • error messages should look as they do today
  • no performance regressions at runtime
  • ideally minimal compile-time impact

[Ericson2314]

There are some very interesting issues here:

• "Flat collections" (Vec, VecDec) can easily store a &mut to the allocator they were crated with, indirection-heavy ones would be significantly bloated if every Box also had an allocator reference.

• Ideally all collections would be resilient in the face of allocation failure, but that will take some major algorithms design work. It's better to convert quickly blowing up on OOM as they do today.

• I had the idea of doing trait InfallableAllocator = Allocator<Error=!> to signify an allocator that wouldn't fail (and thus must panick on OOM). If, when we en-mass convert collections, we have them take an InfallableAllocator parameter, it will be clearer to the end-user that the collections are not oom-resilient, and there would be less panics strune about.

  • I'm not up-to-date with all the discussion of invalid Layouts, so I'll need to double-check that ! wouldn't instead need to be some strict Unsupported;. Then again, it isn't like the collections will be able to handle an unsupported layout either...

• Not sure what the ETA on default parameters is, so we're probably best off reexporting collections in std with the allocator parameter specialized to the global allcoator. This will work today, right?

[joshlf]

@Ericson2314

I had the idea of doing trait InfallableAllocator = Allocator<Error=!> to signify an allocator that wouldn't fail (and thus must panick on OOM). If, when we en-mass convert collections, we have them take an InfallableAllocator parameter, it will be clearer to the end-user that the collections are not oom-resilient, and there would be less panics strune about.

I believe that the current design doesn't allow you to parametrize the error, so it must return the concrete error type with two arms - Unsupported and Exhausted - and if you want to panic on OOM, it's your (the client's) responsibility to do that by panicking yourself or by calling the oom method.

[Ericson2314]

@joshlf err yes, this is predicated on us adding an associated error type.

[ghost]

To support the allocator trait in the collection types, the constructors will need to be duplicated.
@pnkfelix has used the _in suffix and appended the allocator to the list of arguments:

impl RawVec<T, A>
    fn with_capacity(n: usize) -> RawVec<T, HeapAlloc>;
    fn with_capacity_in(n: usize, a: A) -> RawVec<T, A>;
}

I personally quite like this approach.

[ghost]

@Ericson2314 mentioned this:

"Flat collections" (Vec, VecDec) can easily store a &mut to the allocator they were crated with, indirection-heavy ones would be significantly bloated if every Box also had an allocator reference.

I am currently implementing a non-trivial allocator myself (for a relocatable heap) and noticed that it don't need to keep anything around for deallocation (the type information is still needed).

This still works when multible allocators are used.

I wonder if this could be integrated somehow? Associated types to indicate what needs to be stored?

[joshlf]

@s3bk

I am currently implementing a non-trivial allocator myself (for a relocatable heap) and noticed that it don't need to keep anything around for deallocation (the type information is still needed).

This still works when multible allocators are used.

I wonder if this could be integrated somehow? Associated types to indicate what needs to be stored?

So the idea is that, for allocators where all instances of the type refer to the same global singleton, you wouldn't actually need to store a reference? Presumably that wouldn't work for allocators that could actually have multiple distinct instances, right?

This motivates a follow-up: if you were implementing a global singleton, couldn't you just make the actual Alloc type that you expose to users be a ZST that delegates to some private instance under the hood?

[ghost]

@joshlf

So the idea is that, for allocators where all instances of the type refer to the same global singleton, you wouldn't actually need to store a reference? Presumably that wouldn't work for allocators that could actually have multiple distinct instances, right?

Yes, for allocators with multiple instances the allocation function would benefit from non-zero knowledge.
(Mine needs to know the base address for instance. And working with more than one allocator is perfectly reasonable.)
What I was trying to point at, is that the pointer passed to dealloc is enough information to recover the actual instance used to allocate the memory.

[rkruppe]

@s3bk How common is the case where you will only ever need to deallocate? All the standard collections and smart pointers have methods that allocate new memory from the same allocator (usually clone; Rc and Arc have make_mut).

Hm, now that I think about it, those methods already need to clone the allocator, so that would potentially offer a way out. Although it's unclear to me if we can switch between storing A and A::DeallocInfo in the collection depending on whether A: Clone, and it's also not clear to me if that would be a good design (it seems to work out mostly by chance, assuming it works out at all).

[ghost]

@rkruppe there could be

trait Alloc {
    type DeallocInfo;
    fn dealloc(&self, ptr: *mut T, layout: Layout);
    fn dealloc_with(info: Self::DeallocInfo, ptr: *mut T, layout: Layout);
    // snip
}
struct Rc<T, A: Alloc> {
  dealloc: A::DeallocInfo;
  // snip
}
impl<T> Rc<T, HeapAlloc> {
    fn make_mut(&mut self) -> &mut T { 
         // use HeapAlloc to potentially clone T
    }
}
impl<T, A: Alloc> Rc<T, A> {
    fn make_mut_in(&mut self, a: A) -> &mut T { 
         // use A to potentially clone T
    }
}

This would keep Rc as thin as possible, not break existing code and allow to use make_mut_in for code that is generic over Alloc

edit: The case for only dealloc: Box, and many non-std types that only provide a minimal api.
edit 2: completely wrong.

[rkruppe]

@s3bk make_mut_in makes HeapAlloc a special case for no good reason, granting it better ergonomics than other allocators that can support make_mut just as well.

edit: The case for only dealloc: Box, and many non-std types that only provide a minimal api.

Box::clone allocates. Do you have examples of such non-std types?

[ghost]

@rkruppe It looks like I was completely wrong here. Most cases are covered by Box and the others reallocate.
Not sure how the the ergonomics could be solved, apart from adding another trait.

[joshlf]

@s3bk

Most cases are covered by Box and the others reallocate.

Is that a problem? The key property of DeallocInfo is the ability to, given a pointer allocated from the allocator, locate the allocator itself. Thus, you should be able to have any method that takes an already-allocated pointer have a _with variant - realloc, shrink, etc.

[joshlf]

@s3bk

There's another angle to consider here: how do we obtain a DeallocInfo? Presumably we need to guarantee that an object of that type has the same lifetime as the Alloc it was obtained from, so how do we ensure that? I'm imagining something like fn dealloc_info(&mut self) -> &mut Self::DeallocInfo, but then the lifetimes get pretty hairy. In particular, once somebody has a &mut DeallocInfo, the allocator can't be used again until the reference goes away because it's a mutable reference. This means you can allocate one Box at a time from an allocator...

[ghost]

@joshlf The lifetime problem is indeed tricky…

What if a lifetime would be added to the Alloc trait?

trait Alloc<'a> {
    type DeallocInfo;
    fn dealloc_info(&self) -> Self::DeallocInfo;
    ...
}
impl<'a> Alloc<'a> for &'a MyAllocator {
    type DeallocInfo = MyDeallocInfo<'a>;
    …
}

[joshlf]

@s3bk that's an interesting approach, but doesn't doing impl<'a> Alloc<'a> for &'a MyAllocator require all allocators to be threadsafe? I'm writing an allocator right now that very much benefits from not needing to be threadsafe.

[ghost]

@joshlf this crude example should not be Send or Sync. So it could safely be non-theadsafe .
Or am I missing something?

[joshlf]

@s3bk Ah I think you're right about that. Good call.

[Ericson2314]

Not quite sure where the conversation is now, but we should probably discuss

type SomeCollection<T> = collections::SomeCollection<T, GlobalAlloc>;

As a temporary way to get the ball rolling here without impacting breaking std, needing newtypes or language design work for default params with aliases.

(Picking up the discussion from #43112 (comment).)

[Ericson2314]

Cross post #32838 (comment)

I've started generalizing the collections myself; the easiest way to demonstrate how the associated return type helps is to just do the code change, I suppose.

[alexreg]

So, as far as I can tell, the real issue here is the one with knowing which allocator to use for further allocations or deallocation using a type. In the case of a single instance of each Alloc type, this can be solved all statically (no memory or runtime overhead), as mentioned. For multiple instances of an allocator type, there needs to be a (mutable) reference in each value to the allocator used. I would thus propose an associated type on Alloc with an AllocRef trait bound. The AllocRef trait would have a constructor method that takes the Alloc instance – for a singleton Alloc type (including the global allocator), the type implementing AllocRef would be zero-sized, and the constructor method wouldn't care about the instance it was passed, whereas it would simple store the reference for multiple-instance allocators.

This would seem to achieve the best of both worlds – no additional cost whatsoever over the current situation when you use the global allocator, and the necessary (but minimum) cost when you use custom allocators.

[SimonSapin]

@alexreg I don’t think there is a need for an AllocRef trait.

What happens at the moment for RawVec is that it has a type parameter A: Alloc and contains a value of type A directly. When the allocator type is zero-size like struct Global; in std::heap, this field has no overhead. For allocators that need to keep per-instance state, the idea is that the trait would be implemented not for the allocator struct, but for references to it: impl<'a> Alloc for &'a MyAllocator, and what you pass to generic collections is such a reference. (In this case with &_ rather than &mut _ so that you can use the same allocator in multiple collections at the same time, the allocator needs to use somthing like Cell or Mutex internally.)

[Ericson2314]

@alexreg I guess I was waiting for some concensus that we would explore this path before working in it further. It wasn't very hard to get that far but these things bit-rot fast. See rust-lang-nursery/portability-wg#1 (comment) for where I propose this as part of a roadmap.

[Ericson2314]

I guess now that there is a global Alloc trait, this takes some of the pressure of regular Alloc. That might help things.

[alexreg]

@Ericson2314 I think there's an implicit go-ahead from people in the lang team, though I could be wrong. Certainly @glandium encouraged me submit a PR for this sort of stuff.

Anyway, this doesn't need to be part of the portability / embedded development roadmap, though I can certainly see why it fits there. But with the interface for Alloc and GlobalAlloc getting pretty settled after some recent PRs, I think you'd be at a good point to submit a PR for your work on this issue, shortly. :-)

[glandium]

Note that I have boxed.rs changes for Box<T, A> on hold pending nightly switching to 1.28 and building with 1.27, which simplifies stage0 greatly (which should happen next week). Also note that as of current nightly, this only realy works with ZSTs (some compiler internals still assume a Box has the size of a pointer). I have some example code I need to publish somewhere.

One thing that I would like to figure out is recursive use. Like, Vec<Box<T, &SomeAlloc>, &SomeAlloc>. It would be good to have a way to not have each of those Boxes having a reference to the allocator that is the same as the Vec they live in, doubling the size of the Vec.

[alexreg]

@glandium Good to know about Box and stage0 compilation. We can wait on that I think.

Any optimisation for recursive usage of values that depend on custom allocators would be tricky, I think... Perhaps the way to go would be for all types supporting a custom allocator to implement a common trait that allows cloning while removing the allocator. Then it would be up to the collection/container to manage deallocation of items correctly using the reference to its own allocator. I reckon this would work without too many hitches, and would probably be an immediate improvement, but I worry about situations like this:

let foo = Vec::new_in(alloc);
foo.push(Box::new_in(123, alloc));

If the push method moved out of its argument and generated a new Box<T, SomeZST> (with no allocator) to store in the collection, would a temporary Box<T, A> still be created? If not, then we are good to go. If this problem exists, then we need to think harder...

[Ericson2314]

@eddy fixed the internals but that PR was closed due to a windows problem. Was there another PR?

@alexreg I'm mainly worried about the associated type which seems like a breaking change for Alloc.

[glandium]

I'm mainly worried about the associated type which seems like a breaking change for Alloc.

What associated type?

[main--]

One thing that I would like to figure out is recursive use. Like, Vec<Box<T, &SomeAlloc>, &SomeAlloc>. It would be good to have a way to not have each of those Boxes having a reference to the allocator that is the same as the Vec they live in, doubling the size of the Vec.

Note that this is not necessarily true: When an allocator is not zero-sized, this means that it carries some kind of contextual information to distinguish between multiple instances of it (most commonly a pointer to some manage information, e.g. for an arena). Considering that not only the Vec's backing storage may be allocated from a different instance than its contents but those may in turn each belong to a different instance of the allocator I can't think of a straightforward way to optimize this.

The Vec would have to specialize on heap-allocated contents (Box, Vec, etc), then allocators would need to support an equality check as well as cloning so the allocators can be elided if they're the same for every element. But even then I don't see this working without some compiler magic that makes allocators behave kind-of like fat pointers so the Vec can actually return valid references to its allocatorless elements.

[alexreg]

The Vec would have to specialize on heap-allocated contents (Box, Vec, etc), then allocators would need to support an equality check as well as cloning so the allocators can be elided if they're the same for every element. But even then I don't see this working without some compiler magic that makes allocators behave kind-of like fat pointers so the Vec can actually return valid references to its allocatorless elements.

Yes, implementing Eq or at least PartialEq would be necessary for allocators, but this is pretty trivial. As for the solution, do you see any reason my proposal wouldn't work?

Also, do you have any thoughts on my suggestion, @glandium?

[Ericson2314]

@glandium QuiltOS@0d29c18 It's crucial to how I share code between infallible and fallible stuff.

[main--]

@alexreg I don't see a solution for returning references to objects that had their allocator stripped in your proposal. Of course, the entire thing also falls flat on its face the moment you wrap a heap-allocated type into something else before storing it in the Vec as the specialization can no longer trigger.

[glandium]

@Ericson2314: I have a proposal for fallible allocations (and a replacement for *_excess too) that I'll write up and post either as a pre-RFC or a RFC next week.

[alexreg]

@main-- For references, I was imagining that any containers/collections would only return "augmented references" (with a reference to the allocator, at least if it's not single-instance) – these would have Deref impls of course. There would then be a specialisation of the trait CustomAlloced for the type &AugmentedReference.

As for the second issue, this would best be solved by making CustomAlloced (for lack of a better name right now) an auto-trait, i.e. compositional.

[Ericson2314]

@glandium good look; hope what I linked is useful.

Vec<Box<T, &SomeAlloc>, &SomeAlloc>

I've been thinking (albeit just occasionally) about this for years, and the only solution I know of is linear types. I wouldn't worry about it for now.

[alexreg]

@Ericson2314 Any progress lately? Also, thoughts on the above regarding nested containers?

[joshlf]

I wrote up a proposal of a general-purpose execution context that might solve the problems with having to store Allocs in order to deallocate: pre-pre-RFC: Execution Context

[Ericson2314]

Here's a few PRs on this. Thanks @glandium!

Compiler

1. #50097 preparatory work in rustc to allow extra type params on Box. Take 2; it's actually merged, but only supports ZST.

2. Get non-0 sized types working for Box too. Some of #47043 (upon which the previous PR was based) might be useful. Thankfully no library effort is blocked on this.

Library

1. #50882 convert box to use Alloc trait.
2. #52420 convert other collections to use Alloc trait, and allow fallibility polymorphism

[gnzlbg]

The current API of shrink_in_place is broken.

It returns a Result which might indicate an error, but it cannot actually ever error: the memory block that ptr might point to on success must fit the new size, but this is always the case. First, the pointer does not change (the API is "in place"), so the alignment does not change. Second, the original size is larger than the new desired smaller size, so the original memory block, and every possible block in between, all fit the desired smaller size because they are all larger than it, and have the same alignment.

Note that this means that an implementation of shrink_in_place that actually does nothing beyond just simply returning success is ok. Such an implementation would be very fast, but not very useful. The problem is that not doing anything is actually the right call in some cases (e.g. shrinking a 1Gb allocation by 1 byte), this means that if an implementation were to error to indicate that an allocation cannot be shrunk at all, the user would then need to differentiate between whether this is because shrink_in_place is not supported by the allocator, or because for the size requested not shrinking the allocation is the right call.

So how do we fix this. Well, first I thought about passing a size range, that is, a desired size, and a larger size, so that if the allocation cannot be shrunk in place to be in that range, we can error.

But this is too complicated: shrink_in_place should always succeed and return the Excess. This allows the user to check, whether the new size of the allocation is small enough, or not, and do something else.

Also, while the API of grow_in_place is fine, if we fix this by going the "return Excess" route, I find the arguments for splitting realloc_in_place into grow/shrink_in_place might be worth revisiting. Note that for grow_in_place we can return a meaningful error if we cannot grow the allocation enough, but we can also grow the allocation beyond the desired size, and currently the user would then need to use usable_size to obtain that size, and potentially, we would need to add a grow_in_place_excess method.

At this point, I think I would just prefer a realloc_in_place method that always returns Excess, so that the user can shrink or grow their allocation, and just check the Excess to reuse the extra space, or to see if the allocation is small or big enough. This basically trims 4 methods (shrink/grow_in_place/_excess), two of which their API is currently broken, into 1 method with a functioning API.

EDIT: we could also make this method return a Result, but if the user is shrinking the allocation that will be meaningless. The alternative would be just two methods:

• shrink_in_place(...) -> Excess
• grow_in_place(...) -> Result<Excess, ...>

Most ABIs support returning at least two scalar values in registers, so depending on how many registers it takes to return Result<Excess,...>, this might be another pro for realloc_in_place(...) -> Excess.

[Ericson2314]

I just rebased my old branch making a WIP PR #52420 for this.

[fitzgen]

Hello!

I've been working on a bump allocator and my use cases have required me to fork a couple (and eventually I'll probably end up forking more) std collection types to get them to work with my bump allocator. My hope is that once this issue is resolved and stabilized, I can delete my forked versions and switch back to the std ones. I have a few notes on my experience of porting these types to use my bump allocator that I think also generalize pretty well to the problem of adding support for generic allocators to the std collections.

Collections I have ported thus far:

• Vec: https://github.com/fitzgen/bumpalo/blob/master/src/collections/vec.rs
• String: https://github.com/fitzgen/bumpalo/blob/master/src/collections/string.rs

General notes:

• A bunch of trait implementations aren't feasible for custom allocators, like Default and FromIterator, because the trait's signature doesn't provide access to the allocator. In these cases, it seems fine to just have the appropriate bounds on the impl block to exclude custom allocators. It does raise the question of whether we want to consider introducing additional traits in the future that also provide an allocator.

• Box::leak will be interesting when the Box's allocator is constrained by lifetimes. Eg, if my custom allocator is &'bump Bump, I had better not be able to leak the box to a &'static reference, only &'bump at most. I haven't ported Box yet, and instead turned String::into_boxed_str into String::into_bump_str where it returns &'bump str. (This was just an oversight on my part, where I thought that this was the "equivalent" for bump allocators, but as long as Box::leak does the right thing, then we don't need into_bump_str. Basically I've just been trying to port collections quickly and need to go back and clean this up...)

• My bump allocator doesn't run any Drop implementations, so I don't really want to run any collection's Drops either (deallocation is a no-op anyways, but I have not verified whether LLVM sees through everything). I realize this isn't applicable to the general custom-allocators case though, so I'm not proposing anything in particular, just musing out loud here.

• Mostly, this porting has actually been really easy, and I expect adding custom allocators to std types themselves should be pretty easy too! The hardest parts have been trying to remove unstable usage, and trying to reuse as much of std as possible -- neither of which is an issue for adding custom allocator support to std itself!

• Finally, I just want to express my desire and excitement for a future where std collections are parameterized by custom allocators! It will be awesome!

I hope these notes are helpful in some small way!

[SimonSapin]

Regarding Default and FromIterator: those can exist iff the allocator type implements Default.

Regarding Box::leak, this is an interesting case. I think the today’s type system might not be able to express this in its full generality as that would require something like "associated lifetimes" in traits (if that even makes sense). The simplest for now would be to only implement it for Box<T, std::alloc::Global>.

Even if "deallocation" of the memory directly owned by Vec is a no-op, you should still run drop_in_place on the (slice of) elements. You could for example have Vec<Box<Foo>, &'bump Bump> and those inner boxes do need dropping to avoid memory leaks.

[Ericson2314]

Yes that is what I did for Default. We can add additional methods like new_in and from_iterator_in; new classes can be made later.

[glandium]

@fitzgen See the allocator_api crate, which has the Alloc trait, Box and RawVec.

[fitzgen]

@SimonSapin

Regarding Box::leak, this is an interesting case. I think the today’s type system might not be able to express this in its full generality as that would require something like "associated lifetimes" in traits (if that even makes sense). The simplest for now would be to only implement it for Box<T, std::alloc::Global>.

I think this signature would work with Rust today (although I have not tried it yet):

impl<T, A: Alloc> Box<T, A> {
    pub fn leak<'a>(self) -> &'a T
    where
        A: 'a,
    { ... }
}

[fitzgen]

@fitzgen See the allocator_api crate, which has the Alloc trait, Box and RawVec.

Filed fitzgen/bumpalo#2