149 lines
4.2 KiB
Rust
149 lines
4.2 KiB
Rust
use core::alloc::{GlobalAlloc, Layout};
|
|
use core::ptr::null_mut;
|
|
use core::sync::atomic::{AtomicBool, Ordering};
|
|
use core::ops::{Deref, DerefMut};
|
|
|
|
pub struct Locked<A> {
|
|
inner: core::cell::UnsafeCell<A>,
|
|
lock: AtomicBool,
|
|
}
|
|
|
|
pub struct LockedGuard<'a, A> {
|
|
lock: &'a AtomicBool,
|
|
data: &'a mut A,
|
|
}
|
|
|
|
unsafe impl<A> Sync for Locked<A> {}
|
|
|
|
impl<A> Locked<A> {
|
|
pub const fn new(inner: A) -> Self {
|
|
Self {
|
|
inner: core::cell::UnsafeCell::new(inner),
|
|
lock: AtomicBool::new(false),
|
|
}
|
|
}
|
|
|
|
pub fn lock(&self) -> LockedGuard<'_, A> {
|
|
while self.lock.compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed).is_err() {
|
|
core::hint::spin_loop();
|
|
}
|
|
LockedGuard {
|
|
lock: &self.lock,
|
|
data: unsafe { &mut *self.inner.get() },
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<A> Drop for LockedGuard<'_, A> {
|
|
fn drop(&mut self) {
|
|
self.lock.store(false, Ordering::Release);
|
|
}
|
|
}
|
|
|
|
impl<A> Deref for LockedGuard<'_, A> {
|
|
type Target = A;
|
|
fn deref(&self) -> &Self::Target { self.data }
|
|
}
|
|
|
|
impl<A> DerefMut for LockedGuard<'_, A> {
|
|
fn deref_mut(&mut self) -> &mut Self::Target { self.data }
|
|
}
|
|
|
|
struct ListNode {
|
|
next: Option<&'static mut ListNode>,
|
|
}
|
|
|
|
const BLOCK_SIZES: &[usize] = &[8, 16, 32, 64, 128, 256, 512, 1024, 2048];
|
|
|
|
pub struct SlabAllocator {
|
|
list_heads: [Option<&'static mut ListNode>; BLOCK_SIZES.len()],
|
|
heap_start: usize,
|
|
heap_end: usize,
|
|
next_bump: usize,
|
|
}
|
|
|
|
impl SlabAllocator {
|
|
pub const fn new() -> Self {
|
|
Self {
|
|
list_heads: [None, None, None, None, None, None, None, None, None],
|
|
heap_start: 0,
|
|
heap_end: 0,
|
|
next_bump: 0,
|
|
}
|
|
}
|
|
|
|
pub fn init(&mut self, start: usize, size: usize) {
|
|
self.heap_start = start;
|
|
self.next_bump = start;
|
|
self.heap_end = start + size;
|
|
}
|
|
|
|
fn list_index(layout: &Layout) -> Option<usize> {
|
|
let required_block_size = layout.size().max(layout.align());
|
|
BLOCK_SIZES.iter().position(|&s| s >= required_block_size)
|
|
}
|
|
|
|
fn fallback_alloc(&mut self, layout: Layout) -> *mut u8 {
|
|
let alloc_start = (self.next_bump + layout.align() - 1) & !(layout.align() - 1);
|
|
let alloc_end = alloc_start.checked_add(layout.size()).unwrap_or(self.heap_end + 1);
|
|
|
|
if alloc_end > self.heap_end {
|
|
null_mut() // Out of memory
|
|
} else {
|
|
self.next_bump = alloc_end;
|
|
alloc_start as *mut u8
|
|
}
|
|
}
|
|
}
|
|
|
|
unsafe impl GlobalAlloc for Locked<SlabAllocator> {
|
|
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
|
|
let mut allocator = self.lock();
|
|
|
|
match SlabAllocator::list_index(&layout) {
|
|
Some(index) => {
|
|
match allocator.list_heads[index].take() {
|
|
Some(node) => {
|
|
allocator.list_heads[index] = node.next.take();
|
|
node as *mut ListNode as *mut u8
|
|
}
|
|
None => {
|
|
let block_size = BLOCK_SIZES[index];
|
|
let block_align = block_size;
|
|
let layout = Layout::from_size_align(block_size, block_align).unwrap();
|
|
allocator.fallback_alloc(layout)
|
|
}
|
|
}
|
|
}
|
|
None => allocator.fallback_alloc(layout)
|
|
}
|
|
}
|
|
|
|
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
|
|
let mut allocator = self.lock();
|
|
|
|
match SlabAllocator::list_index(&layout) {
|
|
Some(index) => {
|
|
let new_node = ListNode {
|
|
next: allocator.list_heads[index].take(),
|
|
};
|
|
|
|
assert!(layout.size() >= core::mem::size_of::<ListNode>());
|
|
|
|
let new_node_ptr = ptr as *mut ListNode;
|
|
|
|
// Production-fix для Rust 2024: явная изоляция unsafe-операций
|
|
unsafe {
|
|
new_node_ptr.write(new_node);
|
|
allocator.list_heads[index] = Some(&mut *new_node_ptr);
|
|
}
|
|
}
|
|
None => {
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[global_allocator]
|
|
pub static ALLOCATOR: Locked<SlabAllocator> = Locked::new(SlabAllocator::new());
|