#![deny(missing_docs)]
#![deny(clippy::indexing_slicing)]
#![deny(clippy::unwrap_used)]
#![deny(clippy::expect_used)]
#![deny(clippy::panic)]
#![cfg_attr(
test,
allow(
clippy::expect_used,
clippy::indexing_slicing,
clippy::panic,
clippy::unwrap_used
)
)]
use std::alloc::{GlobalAlloc, Layout, System};
use std::cell::Cell;
pub const LARGEST: usize = 4_096;
const GRAIN: usize = 16;
const CLASSES: usize = LARGEST / GRAIN + 1;
const PER_CLASS_BYTES: usize = 64 << 10;
const PER_CLASS_BLOCKS: usize = 1_024;
const CAPS: [usize; CLASSES] = caps();
#[allow(clippy::indexing_slicing)]
const fn caps() -> [usize; CLASSES] {
let mut caps = [1usize; CLASSES];
let mut class = 0usize;
while class < CLASSES {
let size = if class * GRAIN > GRAIN {
class * GRAIN
} else {
GRAIN
};
let mut cap = PER_CLASS_BYTES / size;
if cap > PER_CLASS_BLOCKS {
cap = PER_CLASS_BLOCKS;
}
if cap < 1 {
cap = 1;
}
caps[class] = cap;
class += 1;
}
caps
}
#[inline]
fn per_class(class: usize) -> usize {
CAPS.get(class).copied().unwrap_or(1)
}
thread_local! {
static HEADS: [Cell<*mut u8>; CLASSES] =
const { [const { Cell::new(std::ptr::null_mut()) }; CLASSES] };
static HELD: [Cell<usize>; CLASSES] = const { [const { Cell::new(0) }; CLASSES] };
}
#[inline]
fn class_of(layout: Layout) -> Option<usize> {
if layout.size() > LARGEST || layout.align() > GRAIN || layout.size() == 0 {
return None;
}
Some(layout.size().div_ceil(GRAIN))
}
#[inline]
fn layout_of(class: usize) -> Layout {
Layout::from_size_align(class.saturating_mul(GRAIN).max(GRAIN), GRAIN)
.unwrap_or_else(|_| Layout::new::<u128>())
}
const BIG_SMALLEST: usize = 8 << 10;
const BIG_CLASSES: usize = 4;
const BIG_PER_CLASS_BYTES: usize = 512 << 10;
#[inline]
fn big_class_of(layout: Layout) -> Option<usize> {
let largest = BIG_SMALLEST << (BIG_CLASSES - 1);
if layout.size() <= LARGEST || layout.size() > largest || layout.align() > GRAIN {
return None;
}
let rounded = layout.size().next_power_of_two().max(BIG_SMALLEST);
Some((rounded.trailing_zeros() - BIG_SMALLEST.trailing_zeros()) as usize)
}
#[inline]
fn big_layout_of(class: usize) -> Layout {
Layout::from_size_align(BIG_SMALLEST << class.min(BIG_CLASSES - 1), GRAIN)
.unwrap_or_else(|_| Layout::new::<u128>())
}
static BIG_HEADS: [std::sync::atomic::AtomicPtr<u8>; BIG_CLASSES] =
[const { std::sync::atomic::AtomicPtr::new(std::ptr::null_mut()) }; BIG_CLASSES];
static BIG_HELD: [std::sync::atomic::AtomicUsize; BIG_CLASSES] =
[const { std::sync::atomic::AtomicUsize::new(0) }; BIG_CLASSES];
static BIG_LOCKS: [std::sync::atomic::AtomicBool; BIG_CLASSES] =
[const { std::sync::atomic::AtomicBool::new(false) }; BIG_CLASSES];
#[inline]
fn with_big_lock<R>(class: usize, work: impl FnOnce() -> R) -> Option<R> {
use std::sync::atomic::Ordering;
let lock = BIG_LOCKS.get(class)?;
while lock
.compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_err()
{
std::hint::spin_loop();
}
let result = work();
lock.store(false, Ordering::Release);
Some(result)
}
#[inline]
unsafe fn big_alloc(layout: Layout) -> *mut u8 {
use std::sync::atomic::Ordering;
let Some(class) = big_class_of(layout) else {
return unsafe { System.alloc(layout) };
};
if let (Some(head), Some(held)) = (BIG_HEADS.get(class), BIG_HELD.get(class)) {
if !head.load(Ordering::Relaxed).is_null() {
let taken = with_big_lock(class, || {
let block = head.load(Ordering::Relaxed);
if block.is_null() {
return block;
}
let next = unsafe { block.cast::<*mut u8>().read() };
head.store(next, Ordering::Relaxed);
held.store(
held.load(Ordering::Relaxed).saturating_sub(1),
Ordering::Relaxed,
);
block
})
.unwrap_or(std::ptr::null_mut());
if !taken.is_null() {
return taken;
}
}
}
unsafe { System.alloc(big_layout_of(class)) }
}
#[inline]
unsafe fn big_dealloc(pointer: *mut u8, layout: Layout) {
use std::sync::atomic::Ordering;
let Some(class) = big_class_of(layout) else {
return unsafe { System.dealloc(pointer, layout) };
};
let cap = BIG_PER_CLASS_BYTES / big_layout_of(class).size();
let kept = match (BIG_HEADS.get(class), BIG_HELD.get(class)) {
(Some(head), Some(held)) => with_big_lock(class, || {
let count = held.load(Ordering::Relaxed);
if count >= cap {
return false;
}
unsafe {
pointer
.cast::<*mut u8>()
.write(head.load(Ordering::Relaxed))
};
head.store(pointer, Ordering::Relaxed);
held.store(count.saturating_add(1), Ordering::Relaxed);
true
})
.unwrap_or(false),
_ => false,
};
if !kept {
unsafe { System.dealloc(pointer, big_layout_of(class)) };
}
}
#[inline]
fn same_block(layout: Layout, new_size: usize) -> bool {
let Ok(wanted) = Layout::from_size_align(new_size, layout.align()) else {
return false;
};
match (class_of(layout), class_of(wanted)) {
(Some(old), Some(new)) => old == new,
(None, None) => {
let old = big_class_of(layout);
old.is_some() && old == big_class_of(wanted)
}
_ => false,
}
}
#[inline]
fn neither_kept(layout: Layout, new_size: usize) -> bool {
let Ok(wanted) = Layout::from_size_align(new_size, layout.align()) else {
return false;
};
class_of(layout).is_none()
&& class_of(wanted).is_none()
&& big_class_of(layout).is_none()
&& big_class_of(wanted).is_none()
}
pub struct Pooled;
unsafe impl GlobalAlloc for Pooled {
#[inline]
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
let Some(class) = class_of(layout) else {
return unsafe { big_alloc(layout) };
};
let taken = HEADS
.try_with(|heads| {
let Some(head) = heads.get(class) else {
return std::ptr::null_mut();
};
let block = head.get();
if block.is_null() {
return std::ptr::null_mut();
}
let next = unsafe { block.cast::<*mut u8>().read() };
head.set(next);
let _ = HELD.try_with(|held| {
if let Some(count) = held.get(class) {
count.set(count.get().saturating_sub(1));
}
});
block
})
.unwrap_or(std::ptr::null_mut());
if !taken.is_null() {
return taken;
}
unsafe { System.alloc(layout_of(class)) }
}
#[inline]
unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
let Some(class) = class_of(layout) else {
return unsafe { big_dealloc(pointer, layout) };
};
let kept = HELD
.try_with(|held| {
let Some(count) = held.get(class) else {
return false;
};
if count.get() >= per_class(class) {
return false;
}
HEADS
.try_with(|heads| {
let Some(head) = heads.get(class) else {
return false;
};
unsafe { pointer.cast::<*mut u8>().write(head.get()) };
head.set(pointer);
count.set(count.get().saturating_add(1));
true
})
.unwrap_or(false)
})
.unwrap_or(false);
if kept {
return;
}
unsafe { System.dealloc(pointer, layout_of(class)) };
}
#[inline]
unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
if same_block(layout, new_size) {
return pointer;
}
unsafe {
let Ok(wanted) = Layout::from_size_align(new_size, layout.align()) else {
return std::ptr::null_mut();
};
let fresh = self.alloc(wanted);
if !fresh.is_null() {
std::ptr::copy_nonoverlapping(pointer, fresh, layout.size().min(new_size));
self.dealloc(pointer, layout);
}
fresh
}
}
}
const CHUNK_BYTES: usize = 64 << 10;
thread_local! {
static CARVED_HEADS: [Cell<*mut u8>; CLASSES] =
const { [const { Cell::new(std::ptr::null_mut()) }; CLASSES] };
static CARVED_SPANS: [Cell<(usize, usize)>; CLASSES] =
const { [const { Cell::new((0, 0)) }; CLASSES] };
}
pub struct Carved;
impl Carved {
#[inline]
fn carve(class: usize) -> *mut u8 {
let size = layout_of(class).size();
CARVED_SPANS
.try_with(|spans| {
let Some(span) = spans.get(class) else {
return std::ptr::null_mut();
};
let (mut next, mut end) = span.get();
if next.saturating_add(size) > end || next == 0 {
let Ok(chunk) = Layout::from_size_align(CHUNK_BYTES.max(size), GRAIN) else {
return std::ptr::null_mut();
};
let base = unsafe { System.alloc(chunk) };
if base.is_null() {
return std::ptr::null_mut();
}
next = base as usize;
end = next.saturating_add(chunk.size());
}
span.set((next.saturating_add(size), end));
next as *mut u8
})
.unwrap_or(std::ptr::null_mut())
}
}
unsafe impl GlobalAlloc for Carved {
#[inline]
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
let Some(class) = class_of(layout) else {
return unsafe { big_alloc(layout) };
};
let taken = CARVED_HEADS
.try_with(|heads| {
let Some(head) = heads.get(class) else {
return std::ptr::null_mut();
};
let block = head.get();
if block.is_null() {
return std::ptr::null_mut();
}
head.set(unsafe { block.cast::<*mut u8>().read() });
block
})
.unwrap_or(std::ptr::null_mut());
if !taken.is_null() {
return taken;
}
Carved::carve(class)
}
#[inline]
unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
let Some(class) = class_of(layout) else {
return unsafe { big_dealloc(pointer, layout) };
};
let _ = CARVED_HEADS.try_with(|heads| {
if let Some(head) = heads.get(class) {
unsafe { pointer.cast::<*mut u8>().write(head.get()) };
head.set(pointer);
}
});
}
#[inline]
unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
if same_block(layout, new_size) {
return pointer;
}
unsafe {
let Ok(wanted) = Layout::from_size_align(new_size, layout.align()) else {
return std::ptr::null_mut();
};
let fresh = self.alloc(wanted);
if !fresh.is_null() {
std::ptr::copy_nonoverlapping(pointer, fresh, layout.size().min(new_size));
self.dealloc(pointer, layout);
}
fresh
}
}
}
pub struct Shared;
static SHARED_HEADS: [std::sync::atomic::AtomicPtr<u8>; CLASSES] =
[const { std::sync::atomic::AtomicPtr::new(std::ptr::null_mut()) }; CLASSES];
static SHARED_HELD: [std::sync::atomic::AtomicUsize; CLASSES] =
[const { std::sync::atomic::AtomicUsize::new(0) }; CLASSES];
static SHARED_LOCKS: [std::sync::atomic::AtomicBool; CLASSES] =
[const { std::sync::atomic::AtomicBool::new(false) }; CLASSES];
#[inline]
fn with_class_lock<R>(class: usize, work: impl FnOnce() -> R) -> Option<R> {
use std::sync::atomic::Ordering;
let lock = SHARED_LOCKS.get(class)?;
while lock
.compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_err()
{
std::hint::spin_loop();
}
let result = work();
lock.store(false, Ordering::Release);
Some(result)
}
unsafe impl GlobalAlloc for Shared {
#[inline]
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
use std::sync::atomic::Ordering;
let Some(class) = class_of(layout) else {
return unsafe { big_alloc(layout) };
};
let (Some(head), Some(held)) = (SHARED_HEADS.get(class), SHARED_HELD.get(class)) else {
return unsafe { System.alloc(layout_of(class)) };
};
if !head.load(Ordering::Relaxed).is_null() {
let taken = with_class_lock(class, || {
let block = head.load(Ordering::Relaxed);
if block.is_null() {
return block;
}
let next = unsafe { block.cast::<*mut u8>().read() };
head.store(next, Ordering::Relaxed);
held.store(
held.load(Ordering::Relaxed).saturating_sub(1),
Ordering::Relaxed,
);
block
})
.unwrap_or(std::ptr::null_mut());
if !taken.is_null() {
return taken;
}
}
unsafe { System.alloc(layout_of(class)) }
}
#[inline]
unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
use std::sync::atomic::Ordering;
let Some(class) = class_of(layout) else {
return unsafe { big_dealloc(pointer, layout) };
};
let (Some(head), Some(held)) = (SHARED_HEADS.get(class), SHARED_HELD.get(class)) else {
return unsafe { System.dealloc(pointer, layout_of(class)) };
};
let kept = with_class_lock(class, || {
let count = held.load(Ordering::Relaxed);
if count >= per_class(class) {
return false;
}
unsafe {
pointer
.cast::<*mut u8>()
.write(head.load(Ordering::Relaxed))
};
head.store(pointer, Ordering::Relaxed);
held.store(count.saturating_add(1), Ordering::Relaxed);
true
})
.unwrap_or(false);
if !kept {
unsafe { System.dealloc(pointer, layout_of(class)) };
}
}
#[inline]
unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
if same_block(layout, new_size) {
return pointer;
}
if neither_kept(layout, new_size) {
return unsafe { System.realloc(pointer, layout, new_size) };
}
unsafe {
let Ok(wanted) = Layout::from_size_align(new_size, layout.align()) else {
return std::ptr::null_mut();
};
let fresh = self.alloc(wanted);
if !fresh.is_null() {
std::ptr::copy_nonoverlapping(pointer, fresh, layout.size().min(new_size));
self.dealloc(pointer, layout);
}
fresh
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_classes_are_the_ones_documented() {
assert_eq!(class_of(Layout::from_size_align(8, 8).unwrap()), Some(1));
assert_eq!(class_of(Layout::from_size_align(16, 16).unwrap()), Some(1));
assert_eq!(class_of(Layout::from_size_align(17, 8).unwrap()), Some(2));
assert_eq!(
class_of(Layout::from_size_align(LARGEST, 16).unwrap()),
Some(CLASSES - 1)
);
assert_eq!(class_of(Layout::from_size_align(4, 4).unwrap()), Some(1));
assert_eq!(class_of(Layout::from_size_align(1, 1).unwrap()), Some(1));
assert_eq!(
class_of(Layout::from_size_align(LARGEST + 1, 8).unwrap()),
None
);
assert_eq!(class_of(Layout::from_size_align(32, 32).unwrap()), None);
assert_eq!(class_of(Layout::from_size_align(0, 1).unwrap()), None);
}
#[test]
fn a_class_block_covers_every_request_in_it() {
for size in 1..=LARGEST {
let Some(layout) = Layout::from_size_align(size, 8).ok() else {
continue;
};
let Some(class) = class_of(layout) else {
continue;
};
let block = layout_of(class);
assert!(block.size() >= layout.size(), "size {size}");
assert!(block.align() >= layout.align(), "size {size}");
}
}
#[test]
fn a_freed_block_is_the_one_handed_back() {
let layout = Layout::from_size_align(64, 8).expect("a layout");
unsafe {
let first = Pooled.alloc(layout);
assert!(!first.is_null());
Pooled.dealloc(first, layout);
let second = Pooled.alloc(layout);
assert_eq!(first, second, "the freed block was not recycled");
Pooled.dealloc(second, layout);
}
}
#[test]
fn recycling_stays_inside_the_block() {
let layout = Layout::from_size_align(16, 8).expect("a layout");
unsafe {
let guard = Pooled.alloc(layout);
let block = Pooled.alloc(layout);
std::ptr::write_bytes(guard, 0xAB, layout.size());
Pooled.dealloc(block, layout);
let again = Pooled.alloc(layout);
assert_eq!(block, again);
for at in 0..layout.size() {
assert_eq!(guard.add(at).read(), 0xAB, "the link ran past its block");
}
Pooled.dealloc(again, layout);
Pooled.dealloc(guard, layout);
}
}
#[test]
fn a_shared_block_freed_on_one_thread_serves_another() {
let layout = Layout::from_size_align(3_000, 8).expect("a layout");
let freed = std::thread::spawn(move || unsafe {
let block = Shared.alloc(layout);
assert!(!block.is_null());
Shared.dealloc(block, layout);
block as usize
})
.join()
.expect("the thread ran");
unsafe {
let again = Shared.alloc(layout);
assert_eq!(
again as usize, freed,
"the shared list did not recycle the block"
);
let grown = Shared.realloc(again, layout, 3_001);
assert_eq!(grown, again, "a realloc inside one class moved the block");
Shared.dealloc(grown, Layout::from_size_align(3_001, 8).expect("a layout"));
}
}
#[test]
fn carved_blocks_are_distinct_and_recycled() {
std::thread::spawn(|| {
let layout = Layout::from_size_align(1_000, 8).expect("a layout");
let block = layout_of(class_of(layout).expect("a class")).size();
let per_chunk = CHUNK_BYTES / block;
unsafe {
let mut held = Vec::new();
for nth in 0..per_chunk.saturating_add(3) {
let pointer = Carved.alloc(layout);
assert!(!pointer.is_null());
std::ptr::write_bytes(pointer, (nth % 251) as u8, layout.size());
held.push(pointer);
}
for (nth, pointer) in held.iter().enumerate() {
for at in [0, layout.size() - 1] {
assert_eq!(pointer.add(at).read(), (nth % 251) as u8, "block {nth}");
}
}
let last = held.pop().expect("a block");
Carved.dealloc(last, layout);
let again = Carved.alloc(layout);
assert_eq!(again, last, "the freed block was not recycled");
held.push(again);
for pointer in held {
Carved.dealloc(pointer, layout);
}
}
})
.join()
.expect("the thread ran");
}
#[test]
fn page_sized_requests_fall_in_power_of_two_classes() {
let class = |size| big_class_of(Layout::from_size_align(size, 8).unwrap());
assert_eq!(class(LARGEST), None, "the small classes hold this");
assert_eq!(class(LARGEST + 1), Some(0));
assert_eq!(class(8 << 10), Some(0));
assert_eq!(class((8 << 10) + 1), Some(1));
assert_eq!(class(32 << 10), Some(2));
assert_eq!(class(64 << 10), Some(3));
assert_eq!(
class((64 << 10) + 1),
None,
"the system allocator holds this"
);
assert_eq!(
big_class_of(Layout::from_size_align(32 << 10, 64).unwrap()),
None,
"more aligned than the grain"
);
}
#[test]
fn page_sized_blocks_are_recycled() {
let small = Layout::from_size_align(20_000, 8).unwrap();
unsafe {
for allocator in [&Shared as &dyn GlobalAlloc, &Carved, &Pooled] {
let pointer = allocator.alloc(small);
assert!(!pointer.is_null());
std::ptr::write_bytes(pointer, 9, small.size());
let grown = allocator.realloc(pointer, small, 30_000);
assert_eq!(grown, pointer, "30,000 bytes is still the 32 KiB class");
let moved =
allocator.realloc(grown, Layout::from_size_align(30_000, 8).unwrap(), 40_000);
assert!(!moved.is_null());
assert_eq!(moved.add(19_999).read(), 9, "the bytes were not kept");
let at_forty = Layout::from_size_align(40_000, 8).unwrap();
allocator.dealloc(moved, at_forty);
let again = allocator.alloc(Layout::from_size_align(64 << 10, 8).unwrap());
assert_eq!(
again, moved,
"the freed 64 KiB block was not handed out again"
);
allocator.dealloc(again, Layout::from_size_align(64 << 10, 8).unwrap());
}
let held = BIG_HELD
.iter()
.map(|count| count.load(std::sync::atomic::Ordering::Relaxed))
.sum::<usize>();
let huge = Layout::from_size_align(128 << 10, 8).unwrap();
let pointer = Shared.alloc(huge);
assert!(!pointer.is_null());
Shared.dealloc(pointer, huge);
let after = BIG_HELD
.iter()
.map(|count| count.load(std::sync::atomic::Ordering::Relaxed))
.sum::<usize>();
assert_eq!(after, held, "a block above 64 KiB was kept");
}
}
#[test]
fn realloc_keeps_the_bytes() {
unsafe {
let small = Layout::from_size_align(16, 8).expect("a layout");
let block = Pooled.alloc(small);
std::ptr::write_bytes(block, 0x5A, small.size());
let inside = Pooled.realloc(block, small, 12);
assert_eq!(inside, block, "a realloc inside one class moved the block");
let across = Pooled.realloc(inside, small, 200);
assert!(!across.is_null());
for at in 0..small.size() {
assert_eq!(across.add(at).read(), 0x5A, "realloc lost a byte");
}
Pooled.dealloc(across, Layout::from_size_align(200, 8).expect("a layout"));
}
}
}