use std::alloc::{alloc, dealloc, Layout};
use std::marker::PhantomData;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use crate::sync::mutex::Mutex;
type Generation = u32;
#[derive(Debug)]
pub struct Handle<T> {
index: u32,
generation: Generation,
_marker: PhantomData<T>,
}
impl<T> Copy for Handle<T> {}
impl<T> Clone for Handle<T> {
fn clone(&self) -> Self {
*self
}
}
impl<T> PartialEq for Handle<T> {
fn eq(&self, other: &Self) -> bool {
self.index == other.index && self.generation == other.generation
}
}
impl<T> Eq for Handle<T> {}
impl<T> std::hash::Hash for Handle<T> {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.index.hash(state);
self.generation.hash(state);
}
}
impl<T> Handle<T> {
pub const fn dangling() -> Self {
Self {
index: u32::MAX,
generation: 0,
_marker: PhantomData,
}
}
pub fn is_dangling(&self) -> bool {
self.index == u32::MAX
}
pub fn raw_index(&self) -> u32 {
self.index
}
pub fn raw_generation(&self) -> u32 {
self.generation
}
}
impl<T> Default for Handle<T> {
fn default() -> Self {
Self::dangling()
}
}
struct Slot {
ptr: *mut u8,
size: usize,
generation: Generation,
in_use: bool,
relocatable: bool,
on_relocate: Option<Box<dyn Fn(*mut u8, *mut u8) + Send + Sync>>,
}
impl Slot {
#[allow(dead_code)]
fn empty() -> Self {
Self {
ptr: std::ptr::null_mut(),
size: 0,
generation: 0,
in_use: false,
relocatable: true,
on_relocate: None,
}
}
}
pub struct HandleAllocator {
slots: Mutex<Vec<Slot>>,
free_list: Mutex<Vec<u32>>,
total_allocated: AtomicU64,
active_count: AtomicU32,
relocation_count: AtomicU64,
}
impl HandleAllocator {
pub fn new() -> Self {
Self {
slots: Mutex::new(Vec::with_capacity(1024)),
free_list: Mutex::new(Vec::new()),
total_allocated: AtomicU64::new(0),
active_count: AtomicU32::new(0),
relocation_count: AtomicU64::new(0),
}
}
pub fn alloc<T>(&self) -> Option<Handle<T>> {
self.alloc_with_options::<T>(true, None)
}
pub fn alloc_with_options<T>(
&self,
relocatable: bool,
on_relocate: Option<Box<dyn Fn(*mut u8, *mut u8) + Send + Sync>>,
) -> Option<Handle<T>> {
let size = std::mem::size_of::<T>();
let align = std::mem::align_of::<T>();
let layout = Layout::from_size_align(size, align).ok()?;
let ptr = unsafe { alloc(layout) };
if ptr.is_null() {
return None;
}
let (index, generation) = self.allocate_slot(ptr, size, relocatable, on_relocate);
self.total_allocated.fetch_add(size as u64, Ordering::Relaxed);
self.active_count.fetch_add(1, Ordering::Relaxed);
Some(Handle {
index,
generation,
_marker: PhantomData,
})
}
fn allocate_slot(
&self,
ptr: *mut u8,
size: usize,
relocatable: bool,
on_relocate: Option<Box<dyn Fn(*mut u8, *mut u8) + Send + Sync>>,
) -> (u32, Generation) {
let mut free_list = self.free_list.lock();
let mut slots = self.slots.lock();
if let Some(index) = free_list.pop() {
let slot = &mut slots[index as usize];
slot.ptr = ptr;
slot.size = size;
slot.generation = slot.generation.wrapping_add(1);
slot.in_use = true;
slot.relocatable = relocatable;
slot.on_relocate = on_relocate;
(index, slot.generation)
} else {
let index = slots.len() as u32;
slots.push(Slot {
ptr,
size,
generation: 1,
in_use: true,
relocatable,
on_relocate,
});
(index, 1)
}
}
pub fn free<T>(&self, handle: Handle<T>) {
if handle.is_dangling() {
return;
}
let mut slots = self.slots.lock();
let mut free_list = self.free_list.lock();
if let Some(slot) = slots.get_mut(handle.index as usize) {
if slot.in_use && slot.generation == handle.generation {
let layout = Layout::from_size_align(slot.size, 1).expect("Invalid layout");
unsafe {
dealloc(slot.ptr, layout);
}
self.total_allocated.fetch_sub(slot.size as u64, Ordering::Relaxed);
self.active_count.fetch_sub(1, Ordering::Relaxed);
slot.ptr = std::ptr::null_mut();
slot.in_use = false;
slot.on_relocate = None;
free_list.push(handle.index);
}
}
}
pub fn resolve<T>(&self, handle: Handle<T>) -> Option<*const T> {
if handle.is_dangling() {
return None;
}
let slots = self.slots.lock();
slots.get(handle.index as usize).and_then(|slot| {
if slot.in_use && slot.generation == handle.generation {
Some(slot.ptr as *const T)
} else {
None
}
})
}
pub fn resolve_mut<T>(&self, handle: Handle<T>) -> Option<*mut T> {
if handle.is_dangling() {
return None;
}
let slots = self.slots.lock();
slots.get(handle.index as usize).and_then(|slot| {
if slot.in_use && slot.generation == handle.generation {
Some(slot.ptr as *mut T)
} else {
None
}
})
}
pub fn is_valid<T>(&self, handle: Handle<T>) -> bool {
if handle.is_dangling() {
return false;
}
let slots = self.slots.lock();
slots.get(handle.index as usize).map_or(false, |slot| {
slot.in_use && slot.generation == handle.generation
})
}
pub fn pin<T>(&self, handle: Handle<T>) {
if handle.is_dangling() {
return;
}
let mut slots = self.slots.lock();
if let Some(slot) = slots.get_mut(handle.index as usize) {
if slot.in_use && slot.generation == handle.generation {
slot.relocatable = false;
}
}
}
pub fn unpin<T>(&self, handle: Handle<T>) {
if handle.is_dangling() {
return;
}
let mut slots = self.slots.lock();
if let Some(slot) = slots.get_mut(handle.index as usize) {
if slot.in_use && slot.generation == handle.generation {
slot.relocatable = true;
}
}
}
pub fn defragment(&self) -> usize {
let mut slots = self.slots.lock();
let mut relocations = 0;
let relocatable: Vec<usize> = slots
.iter()
.enumerate()
.filter(|(_, s)| s.in_use && s.relocatable)
.map(|(i, _)| i)
.collect();
for idx in relocatable {
let slot = &mut slots[idx];
let layout = match Layout::from_size_align(slot.size, 16) {
Ok(l) => l,
Err(_) => continue,
};
let new_ptr = unsafe { alloc(layout) };
if new_ptr.is_null() {
continue;
}
unsafe {
std::ptr::copy_nonoverlapping(slot.ptr, new_ptr, slot.size);
}
let old_ptr = slot.ptr;
if let Some(ref callback) = slot.on_relocate {
callback(old_ptr, new_ptr);
}
unsafe {
dealloc(old_ptr, layout);
}
slot.ptr = new_ptr;
relocations += 1;
}
self.relocation_count.fetch_add(relocations as u64, Ordering::Relaxed);
relocations
}
pub fn total_allocated(&self) -> u64 {
self.total_allocated.load(Ordering::Relaxed)
}
pub fn active_count(&self) -> u32 {
self.active_count.load(Ordering::Relaxed)
}
pub fn relocation_count(&self) -> u64 {
self.relocation_count.load(Ordering::Relaxed)
}
pub fn stats(&self) -> HandleAllocatorStats {
let slots = self.slots.lock();
let free_list = self.free_list.lock();
let relocatable_count = slots.iter().filter(|s| s.in_use && s.relocatable).count();
let pinned_count = slots.iter().filter(|s| s.in_use && !s.relocatable).count();
HandleAllocatorStats {
total_allocated: self.total_allocated.load(Ordering::Relaxed),
active_handles: self.active_count.load(Ordering::Relaxed),
total_slots: slots.len(),
free_slots: free_list.len(),
relocatable_count,
pinned_count,
relocation_count: self.relocation_count.load(Ordering::Relaxed),
}
}
}
impl Default for HandleAllocator {
fn default() -> Self {
Self::new()
}
}
unsafe impl Send for HandleAllocator {}
unsafe impl Sync for HandleAllocator {}
#[derive(Debug, Clone, Default)]
pub struct HandleAllocatorStats {
pub total_allocated: u64,
pub active_handles: u32,
pub total_slots: usize,
pub free_slots: usize,
pub relocatable_count: usize,
pub pinned_count: usize,
pub relocation_count: u64,
}
pub struct PinGuard<'a, T> {
allocator: &'a HandleAllocator,
handle: Handle<T>,
}
impl<'a, T> PinGuard<'a, T> {
pub fn new(allocator: &'a HandleAllocator, handle: Handle<T>) -> Self {
allocator.pin(handle);
Self { allocator, handle }
}
pub fn handle(&self) -> Handle<T> {
self.handle
}
pub fn get(&self) -> Option<*const T> {
self.allocator.resolve(self.handle)
}
pub fn get_mut(&self) -> Option<*mut T> {
self.allocator.resolve_mut(self.handle)
}
}
impl<'a, T> Drop for PinGuard<'a, T> {
fn drop(&mut self) {
self.allocator.unpin(self.handle);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_basic_alloc_free() {
let allocator = HandleAllocator::new();
let handle: Handle<u64> = allocator.alloc().unwrap();
assert!(allocator.is_valid(handle));
let ptr = allocator.resolve_mut(handle).unwrap();
unsafe { *ptr = 42; }
let value = unsafe { *allocator.resolve(handle).unwrap() };
assert_eq!(value, 42);
allocator.free(handle);
assert!(!allocator.is_valid(handle));
}
#[test]
fn test_generation_invalidation() {
let allocator = HandleAllocator::new();
let handle1: Handle<u64> = allocator.alloc().unwrap();
allocator.free(handle1);
let handle2: Handle<u64> = allocator.alloc().unwrap();
assert_eq!(handle1.index, handle2.index);
assert_ne!(handle1.generation, handle2.generation);
assert!(!allocator.is_valid(handle1));
assert!(allocator.is_valid(handle2));
}
#[test]
fn test_pin_unpin() {
let allocator = HandleAllocator::new();
let handle: Handle<u64> = allocator.alloc().unwrap();
allocator.pin(handle);
let stats = allocator.stats();
assert_eq!(stats.pinned_count, 1);
assert_eq!(stats.relocatable_count, 0);
allocator.unpin(handle);
let stats = allocator.stats();
assert_eq!(stats.pinned_count, 0);
assert_eq!(stats.relocatable_count, 1);
}
#[test]
fn test_dangling_handle() {
let allocator = HandleAllocator::new();
let handle: Handle<u64> = Handle::dangling();
assert!(handle.is_dangling());
assert!(!allocator.is_valid(handle));
assert!(allocator.resolve(handle).is_none());
}
}