Skip to content

Instantly share code, notes, and snippets.

@thomcc

thomcc/multialloc.rs

Created November 9, 2024 14:10
Show Gist options
  • Save thomcc/b04b5d8514e5b2d7bbbf032de49a0654 to your computer and use it in GitHub Desktop.
Save thomcc/b04b5d8514e5b2d7bbbf032de49a0654 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
multialloc.rs
use core::{alloc::Layout, marker::PhantomData, ptr::NonNull};
/// Utility for performing several conceptually-separate allocations in a single
/// shot. Usage:
/// ```ignore
/// let (foos, numbers, bars): (NonNull<Foo>, NonNull<u32>, NonNull<Bar>) =
/// MultiAllocBuilder::<4>::new()
/// .push::<Foo>(Layout::array::<Foo>(num_foos).unwrap())
/// .push::<u32>(Layout::array::<u32>(num_numbers).unwrap())
/// .push::<Bar>(Layout::array::<Bar>(num_bars).unwrap())
/// // ...up to 16 times...
/// .alloc();
/// crate::mem::multi::dealloc_multi_alloc(foos);
/// ```
pub struct MultiAllocBuilder<const PLANNED_ALLOCS: usize, Tup: TupExtend = ()> {
layouts: [Layout; PLANNED_ALLOCS],
so_far: usize,
_boo: PhantomData<Tup>,
}
impl<const PLANNED_ALLOCS: usize> MultiAllocBuilder<PLANNED_ALLOCS, ()> {
#[inline]
pub const fn new() -> Self {
const {
assert!(
PLANNED_ALLOCS != 0,
"Don't use this if you don't plan on doing any allocations",
)
};
Self {
layouts: [Layout::new::<()>(); PLANNED_ALLOCS],
so_far: 0,
_boo: PhantomData,
}
}
}
impl<const PLANNED_ALLOCS: usize, Tup: TupExtend> MultiAllocBuilder<PLANNED_ALLOCS, Tup> {
#[inline]
pub const fn push<T>(
mut self,
layout: Layout,
) -> MultiAllocBuilder<PLANNED_ALLOCS, <Tup as TupExtend>::ExtendTup<T>>
where
<Tup as TupExtend>::ExtendTup<T>: TupExtend,
{
const {
assert!(
<Tup as TupExtend>::SIZE < PLANNED_ALLOCS,
"Too many `push` calls on MultiAllocBuilder for the \
provided `PLANNED_ALLOCS` value."
);
assert!(
core::mem::size_of::<T>() != 0,
"No ZST's allowed in the MultiAllocBuilder's house",
);
}
self.layouts[self.so_far] = layout;
MultiAllocBuilder {
layouts: self.layouts,
so_far: self.so_far + 1,
_boo: PhantomData,
}
}
#[inline]
pub fn alloc(self) -> <Tup as TupExtend>::TupOfNonNulls {
const {
assert!(
<Tup as TupExtend>::SIZE == PLANNED_ALLOCS,
"`MultiAllocBuilder::alloc` called before the planned number of `alloc` calls.",
);
};
let mut result = [NonNull::<u8>::dangling(); PLANNED_ALLOCS];
let mut offsets = [0; PLANNED_ALLOCS];
let mut current = Layout::new::<(Layout, *mut u8)>();
for (&layout, offset_out) in self.layouts.iter().zip(&mut offsets) {
let Ok((next_layout, next_offset)) = current.extend(layout) else {
// Not quite right to call it with this but whatever.
std::alloc::handle_alloc_error(layout);
};
*offset_out = next_offset;
current = next_layout;
}
let mut pointers = [NonNull::<u8>::dangling(); PLANNED_ALLOCS];
unsafe {
let alloc = std::alloc::alloc(current);
if alloc.is_null() {
std::alloc::handle_alloc_error(current);
}
for (&offset, pointer_out) in offsets.iter().zip(&mut pointers) {
*pointer_out = NonNull::new_unchecked(alloc.byte_add(offset));
}
// Write the data that we need to properly deallocate directly before the first pointer.
pointers[0]
.cast::<(Layout, *mut u8)>()
.sub(1)
.write_unaligned((current, alloc));
}
Tup::make_result_tuple(&mut pointers)
}
}
#[inline]
pub unsafe fn dealloc_multi_alloc<T>(first_pointer: NonNull<T>) {
let (alloced_layout, to_free) = first_pointer
.as_ptr()
.cast::<(Layout, *mut u8)>()
.sub(1)
.read_unaligned();
std::alloc::dealloc(to_free.cast::<u8>(), alloced_layout);
}
mod sealed {
pub struct SealMe;
}
use sealed::SealMe;
pub trait TupExtend: Sized {
type ExtendTup<T>;
type TupOfNonNulls;
// type SameSizeArray<T>;
const SIZE: usize;
fn extend_tup<T>(self, next: T) -> Self::ExtendTup<T>;
#[doc(hidden)]
fn __sealer__(_: &SealMe);
fn make_result_tuple(
data_as_slice_because_of_const_restrictions: &mut [NonNull<u8>],
) -> Self::TupOfNonNulls;
}
impl TupExtend for () {
// type SameSizeArray<T> = [T; 0];
const SIZE: usize = 0;
type ExtendTup<Next> = (Next,);
type TupOfNonNulls = ();
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(next,)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(_: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
()
}
}
impl<A> TupExtend for (A,) {
// type SameSizeArray<T> = [T; 1];
const SIZE: usize = 1;
type ExtendTup<Next> = (A, Next);
type TupOfNonNulls = (NonNull<A>,);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(self.0, next)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(data[0].cast(),)
}
}
impl<A, B> TupExtend for (A, B) {
// type SameSizeArray<T> = [T; 2];
const SIZE: usize = 2;
type ExtendTup<Next> = (A, B, Next);
type TupOfNonNulls = (NonNull<A>, NonNull<B>);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(self.0, self.1, next)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(data[0].cast(), data[1].cast())
}
}
impl<A, B, C> TupExtend for (A, B, C) {
// type SameSizeArray<T> = [T; 3];
const SIZE: usize = 3;
type ExtendTup<Next> = (A, B, C, Next);
type TupOfNonNulls = (NonNull<A>, NonNull<B>, NonNull<C>);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(self.0, self.1, self.2, next)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(data[0].cast(), data[1].cast(), data[2].cast())
}
}
impl<A, B, C, D> TupExtend for (A, B, C, D) {
// type SameSizeArray<T> = [T; 4];
const SIZE: usize = 4;
type ExtendTup<Next> = (A, B, C, D, Next);
type TupOfNonNulls = (NonNull<A>, NonNull<B>, NonNull<C>, NonNull<D>);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(self.0, self.1, self.2, self.3, next)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(
data[0].cast(),
data[1].cast(),
data[2].cast(),
data[3].cast(),
)
}
}
impl<A, B, C, D, E> TupExtend for (A, B, C, D, E) {
// type SameSizeArray<T> = [T; 5];
const SIZE: usize = 5;
type ExtendTup<Next> = (A, B, C, D, E, Next);
type TupOfNonNulls = (NonNull<A>, NonNull<B>, NonNull<C>, NonNull<D>, NonNull<E>);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(self.0, self.1, self.2, self.3, self.4, next)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(
data[0].cast(),
data[1].cast(),
data[2].cast(),
data[3].cast(),
data[4].cast(),
)
}
}
impl<A, B, C, D, E, F> TupExtend for (A, B, C, D, E, F) {
// type SameSizeArray<T> = [T; 6];
const SIZE: usize = 6;
type ExtendTup<Next> = (A, B, C, D, E, F, Next);
type TupOfNonNulls = (
NonNull<A>,
NonNull<B>,
NonNull<C>,
NonNull<D>,
NonNull<E>,
NonNull<F>,
);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(self.0, self.1, self.2, self.3, self.4, self.5, next)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(
data[0].cast(),
data[1].cast(),
data[2].cast(),
data[3].cast(),
data[4].cast(),
data[5].cast(),
)
}
}
impl<A, B, C, D, E, F, G> TupExtend for (A, B, C, D, E, F, G) {
// type SameSizeArray<T> = [T; 7];
const SIZE: usize = 7;
type ExtendTup<Next> = (A, B, C, D, E, F, G, Next);
type TupOfNonNulls = (
NonNull<A>,
NonNull<B>,
NonNull<C>,
NonNull<D>,
NonNull<E>,
NonNull<F>,
NonNull<G>,
);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(self.0, self.1, self.2, self.3, self.4, self.5, self.6, next)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(
data[0].cast(),
data[1].cast(),
data[2].cast(),
data[3].cast(),
data[4].cast(),
data[5].cast(),
data[6].cast(),
)
}
}
impl<A, B, C, D, E, F, G, H> TupExtend for (A, B, C, D, E, F, G, H) {
// type SameSizeArray<T> = [T; 8];
const SIZE: usize = 8;
type ExtendTup<Next> = (A, B, C, D, E, F, G, H, Next);
type TupOfNonNulls = (
NonNull<A>,
NonNull<B>,
NonNull<C>,
NonNull<D>,
NonNull<E>,
NonNull<F>,
NonNull<G>,
NonNull<H>,
);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(
self.0, self.1, self.2, self.3, self.4, self.5, self.6, self.7, next,
)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(
data[0].cast(),
data[1].cast(),
data[2].cast(),
data[3].cast(),
data[4].cast(),
data[5].cast(),
data[6].cast(),
data[7].cast(),
)
}
}
impl<A, B, C, D, E, F, G, H, I> TupExtend for (A, B, C, D, E, F, G, H, I) {
// type SameSizeArray<T> = [T; 9];
const SIZE: usize = 9;
type ExtendTup<Next> = (A, B, C, D, E, F, G, H, I, Next);
type TupOfNonNulls = (
NonNull<A>,
NonNull<B>,
NonNull<C>,
NonNull<D>,
NonNull<E>,
NonNull<F>,
NonNull<G>,
NonNull<H>,
NonNull<I>,
);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(
self.0, self.1, self.2, self.3, self.4, self.5, self.6, self.7, self.8, next,
)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(
data[0].cast(),
data[1].cast(),
data[2].cast(),
data[3].cast(),
data[4].cast(),
data[5].cast(),
data[6].cast(),
data[7].cast(),
data[8].cast(),
)
}
}
impl<A, B, C, D, E, F, G, H, I, J> TupExtend for (A, B, C, D, E, F, G, H, I, J) {
// type SameSizeArray<T> = [T; 10];
const SIZE: usize = 10;
type ExtendTup<Next> = (A, B, C, D, E, F, G, H, I, J, Next);
type TupOfNonNulls = (
NonNull<A>,
NonNull<B>,
NonNull<C>,
NonNull<D>,
NonNull<E>,
NonNull<F>,
NonNull<G>,
NonNull<H>,
NonNull<I>,
NonNull<J>,
);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(
self.0, self.1, self.2, self.3, self.4, self.5, self.6, self.7, self.8, self.9, next,
)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(
data[0].cast(),
data[1].cast(),
data[2].cast(),
data[3].cast(),
data[4].cast(),
data[5].cast(),
data[6].cast(),
data[7].cast(),
data[8].cast(),
data[9].cast(),
)
}
}
impl<A, B, C, D, E, F, G, H, I, J, K> TupExtend for (A, B, C, D, E, F, G, H, I, J, K) {
// type SameSizeArray<T> = [T; 11];
const SIZE: usize = 11;
type ExtendTup<Next> = (A, B, C, D, E, F, G, H, I, J, K, Next);
type TupOfNonNulls = (
NonNull<A>,
NonNull<B>,
NonNull<C>,
NonNull<D>,
NonNull<E>,
NonNull<F>,
NonNull<G>,
NonNull<H>,
NonNull<I>,
NonNull<J>,
NonNull<K>,
);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(
self.0, self.1, self.2, self.3, self.4, self.5, self.6, self.7, self.8, self.9,
self.10, next,
)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(
data[0].cast(),
data[1].cast(),
data[2].cast(),
data[3].cast(),
data[4].cast(),
data[5].cast(),
data[6].cast(),
data[7].cast(),
data[8].cast(),
data[9].cast(),
data[10].cast(),
)
}
}
impl<A, B, C, D, E, F, G, H, I, J, K, L> TupExtend for (A, B, C, D, E, F, G, H, I, J, K, L) {
// type SameSizeArray<T> = [T; 12];
const SIZE: usize = 12;
type ExtendTup<Next> = (A, B, C, D, E, F, G, H, I, J, K, L, Next);
type TupOfNonNulls = (
NonNull<A>,
NonNull<B>,
NonNull<C>,
NonNull<D>,
NonNull<E>,
NonNull<F>,
NonNull<G>,
NonNull<H>,
NonNull<I>,
NonNull<J>,
NonNull<K>,
NonNull<L>,
);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(
self.0, self.1, self.2, self.3, self.4, self.5, self.6, self.7, self.8, self.9,
self.10, self.11, next,
)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(
data[0].cast(),
data[1].cast(),
data[2].cast(),
data[3].cast(),
data[4].cast(),
data[5].cast(),
data[6].cast(),
data[7].cast(),
data[8].cast(),
data[9].cast(),
data[10].cast(),
data[11].cast(),
)
}
}
impl<A, B, C, D, E, F, G, H, I, J, K, L, M> TupExtend for (A, B, C, D, E, F, G, H, I, J, K, L, M) {
// type SameSizeArray<T> = [T; 13];
const SIZE: usize = 13;
type ExtendTup<Next> = (A, B, C, D, E, F, G, H, I, J, K, L, M, Next);
type TupOfNonNulls = (
NonNull<A>,
NonNull<B>,
NonNull<C>,
NonNull<D>,
NonNull<E>,
NonNull<F>,
NonNull<G>,
NonNull<H>,
NonNull<I>,
NonNull<J>,
NonNull<K>,
NonNull<L>,
NonNull<M>,
);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(
self.0, self.1, self.2, self.3, self.4, self.5, self.6, self.7, self.8, self.9,
self.10, self.11, self.12, next,
)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(
data[0].cast(),
data[1].cast(),
data[2].cast(),
data[3].cast(),
data[4].cast(),
data[5].cast(),
data[6].cast(),
data[7].cast(),
data[8].cast(),
data[9].cast(),
data[10].cast(),
data[11].cast(),
data[12].cast(),
)
}
}
impl<A, B, C, D, E, F, G, H, I, J, K, L, M, N> TupExtend
for (A, B, C, D, E, F, G, H, I, J, K, L, M, N)
{
// type SameSizeArray<T> = [T; 14];
const SIZE: usize = 14;
type ExtendTup<Next> = (A, B, C, D, E, F, G, H, I, J, K, L, M, N, Next);
type TupOfNonNulls = (
NonNull<A>,
NonNull<B>,
NonNull<C>,
NonNull<D>,
NonNull<E>,
NonNull<F>,
NonNull<G>,
NonNull<H>,
NonNull<I>,
NonNull<J>,
NonNull<K>,
NonNull<L>,
NonNull<M>,
NonNull<N>,
);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(
self.0, self.1, self.2, self.3, self.4, self.5, self.6, self.7, self.8, self.9,
self.10, self.11, self.12, self.13, next,
)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(
data[0].cast(),
data[1].cast(),
data[2].cast(),
data[3].cast(),
data[4].cast(),
data[5].cast(),
data[6].cast(),
data[7].cast(),
data[8].cast(),
data[9].cast(),
data[10].cast(),
data[11].cast(),
data[12].cast(),
data[13].cast(),
)
}
}
impl<A, B, C, D, E, F, G, H, I, J, K, L, M, N, O> TupExtend
for (A, B, C, D, E, F, G, H, I, J, K, L, M, N, O)
{
// type SameSizeArray<T> = [T; 15];
const SIZE: usize = 15;
type ExtendTup<Next> = (A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, Next);
type TupOfNonNulls = (
NonNull<A>,
NonNull<B>,
NonNull<C>,
NonNull<D>,
NonNull<E>,
NonNull<F>,
NonNull<G>,
NonNull<H>,
NonNull<I>,
NonNull<J>,
NonNull<K>,
NonNull<L>,
NonNull<M>,
NonNull<N>,
NonNull<O>,
);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(
self.0, self.1, self.2, self.3, self.4, self.5, self.6, self.7, self.8, self.9,
self.10, self.11, self.12, self.13, self.14, next,
)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(
data[0].cast(),
data[1].cast(),
data[2].cast(),
data[3].cast(),
data[4].cast(),
data[5].cast(),
data[6].cast(),
data[7].cast(),
data[8].cast(),
data[9].cast(),
data[10].cast(),
data[11].cast(),
data[12].cast(),
data[13].cast(),
data[14].cast(),
)
}
}
impl<A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P> TupExtend
for (A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P)
{
// type SameSizeArray<T> = [T; 16];
const SIZE: usize = 16;
type ExtendTup<Next> = (A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P, Next);
type TupOfNonNulls = (
NonNull<A>,
NonNull<B>,
NonNull<C>,
NonNull<D>,
NonNull<E>,
NonNull<F>,
NonNull<G>,
NonNull<H>,
NonNull<I>,
NonNull<J>,
NonNull<K>,
NonNull<L>,
NonNull<M>,
NonNull<N>,
NonNull<O>,
NonNull<P>,
);
#[inline(always)]
fn extend_tup<Next>(self, next: Next) -> Self::ExtendTup<Next> {
(
self.0, self.1, self.2, self.3, self.4, self.5, self.6, self.7, self.8, self.9,
self.10, self.11, self.12, self.13, self.14, self.15, next,
)
}
fn __sealer__(_: &SealMe) {}
#[inline(always)]
fn make_result_tuple(data: &mut [NonNull<u8>]) -> Self::TupOfNonNulls {
(
data[0].cast(),
data[1].cast(),
data[2].cast(),
data[3].cast(),
data[4].cast(),
data[5].cast(),
data[6].cast(),
data[7].cast(),
data[8].cast(),
data[9].cast(),
data[10].cast(),
data[11].cast(),
data[12].cast(),
data[13].cast(),
data[14].cast(),
data[15].cast(),
)
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment