use crate::types::Result;
use core::ptr::NonNull;
pub struct MemoryPool {
start_ptr: NonNull<u8>,
block_size: usize,
total_blocks: usize,
used_blocks: usize,
free_list: Option<NonNull<u8>>,
}
impl MemoryPool {
pub unsafe fn new(start_ptr: *mut u8, block_size: usize, total_blocks: usize) -> Self {
let start_ptr = NonNull::new_unchecked(start_ptr);
let aligned_block_size = (block_size + 7) & !7;
let mut free_list = None;
let mut current_ptr = start_ptr;
for _ in 0..total_blocks {
let block_ptr = current_ptr;
current_ptr = NonNull::new_unchecked(
(current_ptr.as_ptr() as usize + aligned_block_size) as *mut u8,
);
let next_ptr = free_list;
core::ptr::write(block_ptr.as_ptr() as *mut Option<NonNull<u8>>, next_ptr);
free_list = Some(block_ptr);
}
MemoryPool {
start_ptr,
block_size: aligned_block_size,
total_blocks,
used_blocks: 0,
free_list,
}
}
pub unsafe fn allocate(&mut self) -> Result<NonNull<u8>> {
if let Some(block_ptr) = self.free_list {
let next_ptr = core::ptr::read(block_ptr.as_ptr() as *const Option<NonNull<u8>>);
self.free_list = next_ptr;
self.used_blocks += 1;
Ok(block_ptr)
} else {
Err(crate::types::RemDbError::OutOfMemory)
}
}
pub unsafe fn free(&mut self, ptr: NonNull<u8>) {
let ptr_addr = ptr.as_ptr() as usize;
let start_addr = self.start_ptr.as_ptr() as usize;
let end_addr = start_addr + self.block_size * self.total_blocks;
assert!(
ptr_addr >= start_addr && ptr_addr < end_addr,
"Pointer not in memory pool"
);
core::ptr::write(ptr.as_ptr() as *mut Option<NonNull<u8>>, self.free_list);
self.free_list = Some(ptr);
self.used_blocks -= 1;
}
pub fn used_blocks(&self) -> usize {
self.used_blocks
}
pub fn total_blocks(&self) -> usize {
self.total_blocks
}
pub fn usage(&self) -> f32 {
self.used_blocks as f32 / self.total_blocks as f32
}
pub fn contains(&self, ptr: NonNull<u8>) -> bool {
let ptr_addr = ptr.as_ptr() as usize;
let start_addr = self.start_ptr.as_ptr() as usize;
let end_addr = start_addr + self.block_size * self.total_blocks;
ptr_addr >= start_addr && ptr_addr < end_addr
}
}
pub struct MultiPoolManager<'a> {
pools: &'a mut [MemoryPool],
pool_count: usize,
}
impl<'a> MultiPoolManager<'a> {
pub unsafe fn new(pools: &'a mut [MemoryPool]) -> Self {
let pool_count = pools.len();
MultiPoolManager { pools, pool_count }
}
pub unsafe fn allocate(&mut self, size: usize) -> Result<NonNull<u8>> {
for pool in &mut self.pools[..self.pool_count] {
if pool.block_size >= size {
return pool.allocate();
}
}
Err(crate::types::RemDbError::OutOfMemory)
}
pub unsafe fn free(&mut self, ptr: NonNull<u8>) -> Result<()> {
for pool in &mut self.pools[..self.pool_count] {
if pool.contains(ptr) {
pool.free(ptr);
return Ok(());
}
}
Err(crate::types::RemDbError::InvalidPointer)
}
pub fn total_usage(&self) -> f32 {
let mut total_used = 0;
let mut total_blocks = 0;
for pool in &self.pools[..self.pool_count] {
total_used += pool.used_blocks();
total_blocks += pool.total_blocks();
}
if total_blocks == 0 {
0.0
} else {
total_used as f32 / total_blocks as f32
}
}
}